#include "GaitEngine.h"
#include "Limits.h"

const GaitPattern GAIT_TRIPOD = {
  0.5f,
  { 0.5f, 0.0f, 0.5f, 0.0f, 0.5f, 0.0f }
};

const GaitPattern GAIT_WAVE = {
  5.0f / 6.0f,
  { 0.0f, 1.0f / 6, 2.0f / 6, 3.0f / 6, 4.0f / 6, 5.0f / 6 }
};

void GaitEngine::begin() {
  resetFeet();
}

// Every foot starts at neutral, which is the *middle* of a stride. So the
// cycle restarts with the pattern shifted to put leg 0 at mid-stance (for
// the tripod that puts every grounded leg at mid-stance); a grounded leg
// that started at the beginning of its stance would instead run a full
// stride past neutral into the workspace clamp and stop pushing. Legs that
// land in their swing phase lift off from neutral over the remaining swing.
void GaitEngine::resetFeet() {
  for (int i = 0; i < 6; i++) {
    legs_pattern_[i] = requested_pattern_;
    const GaitPattern* p = legs_pattern_[i];
    float phase = p->leg_phase_offset[i] + p->duty_factor * 0.5f;
    leg_phase_[i] = phase - floorf(phase);
    foot_[i] = Kinematics::neutralStance(i);
    liftoff_[i] = foot_[i];
    swing_start_[i] = leg_phase_[i] < p->duty_factor ? p->duty_factor : leg_phase_[i];
  }
}

static float smoothstep(float t) { return t * t * (3.0f - 2.0f * t); }

void GaitEngine::strideEnvelope(const Twist& twist, FootTarget sets[ENVELOPE_SETS][6]) const {
  float half = cycle_sec_ * requested_pattern_->duty_factor * 0.5f;
  for (int i = 0; i < 6; i++) {
    FootTarget n = Kinematics::neutralStance(i);
    float dx = (twist.vx - twist.omega * n.y) * half;   // half-stride: body velocity at the foot
    float dy = (twist.vy + twist.omega * n.x) * half;
    float sign = (i % 2 == 0) ? 1.0f : -1.0f;
    sets[0][i] = { n.x + sign * dx, n.y + sign * dy, n.z };
    sets[1][i] = { n.x - sign * dx, n.y - sign * dy, n.z };
    sets[2][i] = { n.x, n.y, n.z + SWING_LIFT_MM };
  }
}

const BodyPose GaitEngine::NO_POSE{};

bool GaitEngine::envelopeFeasible(const Twist& twist, const BodyPose& pose) const {
  FootTarget sets[ENVELOPE_SETS][6];
  strideEnvelope(twist, sets);
  for (int k = 0; k < ENVELOPE_SETS; k++)
    if (!pose.feasible(sets[k])) return false;
  return true;
}

// Neighbouring legs step in opposite phases, so each one's stride eats into
// the gap between them from both sides; turning on the spot at full rate
// asks for ~60 mm half-strides at the outer feet, enough to make the legs
// meet. Rather than letting single feet clamp (or collide), the whole
// command is slowed until the stride fits.
Twist GaitEngine::limitTwist(const Twist& twist, const BodyPose& pose, bool chained) {
  float s = 1.0f;
  if (!envelopeFeasible(twist, pose)) {
    float lo = 0.0f, hi = 1.0f;
    for (int n = 0; n < 8; n++) {
      float mid = 0.5f * (lo + hi);
      Twist t = { twist.vx * mid, twist.vy * mid, twist.omega * mid };
      if (envelopeFeasible(t, pose)) lo = mid; else hi = mid;
    }
    s = lo;
  }
  twist_scale_ = chained ? twist_scale_ * s : s;
  return { twist.vx * s, twist.vy * s, twist.omega * s };
}

void GaitEngine::update(float dt, const Twist& twist, const BodyPose& pose, FootTarget out[6]) {
  // Pass 1: advance phases, place swing feet, and work out how far each
  // planted foot would move this tick.
  bool stance[6];
  FootTarget step[6] = {};
  for (int i = 0; i < 6; i++) {
    bool was_stance = leg_phase_[i] < legs_pattern_[i]->duty_factor;
    leg_phase_[i] += dt / cycle_sec_;

    // Touchdown: this leg's cycle wraps back to 0 - the one moment it's safe
    // to hand it a newly requested pattern, since it's grounded with zero
    // relative velocity right here regardless of which two patterns are involved.
    if (leg_phase_[i] >= 1.0f) {
      leg_phase_[i] -= 1.0f;
      legs_pattern_[i] = requested_pattern_;
    }

    const GaitPattern* p = legs_pattern_[i];
    bool in_stance = leg_phase_[i] < p->duty_factor;
    stance[i] = in_stance;

    FootTarget neutral = Kinematics::neutralStance(i);
    float stride_time = cycle_sec_ * p->duty_factor;
    FootTarget& foot = foot_[i];

    if (in_stance) {
      // Planted foot: stays put in the world, so in the body frame it moves
      // opposite to the body's velocity at that point (v + omega x r).
      // Integrated every tick, so the body moves exactly as commanded even
      // when the command changes mid-stride.
      step[i].x = -(twist.vx - twist.omega * foot.y) * dt;
      step[i].y = -(twist.vy + twist.omega * foot.x) * dt;
      foot.z = neutral.z;
    } else {
      if (was_stance) {                      // just lifted off
        liftoff_[i] = foot;
        swing_start_[i] = p->duty_factor;
      }

      // Swing toward the leading end of the next stride: the point that a
      // full stance at the current command carries back through neutral.
      float half = stride_time * 0.5f;
      FootTarget touchdown = {
        neutral.x + (twist.vx - twist.omega * neutral.y) * half,
        neutral.y + (twist.vy + twist.omega * neutral.x) * half,
        neutral.z
      };
      float t = (leg_phase_[i] - swing_start_[i]) / (1.0f - swing_start_[i]);
      float st = smoothstep(t);
      foot.x = liftoff_[i].x + (touchdown.x - liftoff_[i].x) * st;
      foot.y = liftoff_[i].y + (touchdown.y - liftoff_[i].y) * st;
      foot.z = neutral.z + SWING_LIFT_MM * sin(PI * t);
    }
  }

  // Pass 2: the planted feet move together - they're what carries the body.
  // limitTwist() sizes the command for a steady stride, but when the command
  // changes mid-step (a key pressed while a foot has just landed) a planted
  // foot can be asked to travel further than its leg reaches, or into its
  // neighbour. Then all planted feet move less, by one common factor: the
  // body slows down rather than one foot being clamped out of formation.
  float s = 1.0f;
  FootTarget proposed[6];
  auto propose = [&](float k) {
    for (int i = 0; i < 6; i++)
      proposed[i] = stance[i] ? FootTarget{ foot_[i].x + step[i].x * k, foot_[i].y + step[i].y * k, foot_[i].z }
                              : foot_[i];
    return pose.feasible(proposed);
  };
  if (!propose(1.0f)) {
    float lo = 0.0f, hi = 1.0f;
    for (int n = 0; n < 6; n++) {
      float mid = 0.5f * (lo + hi);
      if (propose(mid)) lo = mid; else hi = mid;
    }
    s = lo;
    propose(s);
  }
  applied_twist_ = { twist.vx * s, twist.vy * s, twist.omega * s };
  for (int i = 0; i < 6; i++) {
    foot_[i] = proposed[i];
    out[i] = foot_[i];
  }
}