#include "BodyPose.h"
#include "Limits.h"

void BodyPose::begin() {
  s_roll_.reset(0); s_pitch_.reset(0); s_yaw_.reset(0);
  s_x_.reset(0); s_y_.reset(0); s_z_.reset(0);
}

void BodyPose::update(float dt) {
  s_roll_.setTarget(command_.roll_deg);
  s_pitch_.setTarget(command_.pitch_deg);
  s_yaw_.setTarget(command_.yaw_deg);
  s_x_.setTarget(command_.x_mm);
  s_y_.setTarget(command_.y_mm);
  s_z_.setTarget(command_.z_mm);

  s_roll_.update(dt); s_pitch_.update(dt); s_yaw_.update(dt);
  s_x_.update(dt); s_y_.update(dt); s_z_.update(dt);
}

// Combines the commanded (smoothed) tilt with the robot's own stabilization
// correction, then applies both to a body-frame foot target so the feet stay
// planted while the body tilts/shifts on top of them.
//
// Sign convention here (whether positive roll tips the body left or right,
// whether offsetting +X moves the feet the way you expect) depends on your
// leg numbering and mount angles - bench-check it the same way as the IK
// elbow direction, and flip signs if it goes the wrong way.
FootTarget BodyPose::apply(const FootTarget& in) const { return applyScaled(in, reach_scale_); }

FootTarget BodyPose::applyScaled(const FootTarget& in, float s) const {
  // The feet stay put in the world, so to rotate the body by (roll, pitch,
  // yaw) the body-frame foot targets rotate by the inverse: negated angles,
  // applied yaw -> pitch -> roll. Result: +roll = right side down,
  // +pitch = nose down, +yaw = CCW (HexapodProtocol.h), checked in sim_smoke.
  // `s` scales the whole pose (see fitToReach).
  float roll  = -radians(s * (s_roll_.current + correction_roll_));
  float pitch = -radians(s * (s_pitch_.current + correction_pitch_));
  float yaw   = -radians(s * s_yaw_.current);

  float x = in.x, y = in.y, z = in.z;

  // yaw (about Z)
  float x1 = x * cos(yaw) - y * sin(yaw);
  float y1 = x * sin(yaw) + y * cos(yaw);
  // pitch (about Y)
  float x2 =  x1 * cos(pitch) + z * sin(pitch);
  float z2 = -x1 * sin(pitch) + z * cos(pitch);
  // roll (about X)
  float y3 = y1 * cos(roll) - z2 * sin(roll);
  float z3 = y1 * sin(roll) + z2 * cos(roll);

  FootTarget out;
  out.x = x2 - s * s_x_.current;
  out.y = y3 - s * s_y_.current;
  out.z = z3 - s * s_z_.current;
  return out;
}

bool BodyPose::reachable(const FootTarget sets[][6], int count, float s) const {
  for (int k = 0; k < count; k++) {
    FootTarget posed[6];
    for (int i = 0; i < 6; i++) posed[i] = applyScaled(sets[k][i], s);
    if (!Limits::feetFeasible(posed)) return false;
  }
  return true;
}

bool BodyPose::feasible(const FootTarget feet[6]) const {
  FootTarget posed[6];
  for (int i = 0; i < 6; i++) posed[i] = apply(feet[i]);
  return Limits::feetFeasible(posed);
}

// Clamping each foot separately would bend the six feet out of their rigid
// shape on the ground (feet dragging, legs fighting each other) as soon as
// a combined pose asks for more than a leg can reach, and nothing would stop
// neighbouring legs from swinging into each other. Instead the whole pose
// is scaled down by one factor until every foot set is physically possible,
// so the body does the most of what was asked that it really can.
//
// The scale itself moves at a bounded rate: a sudden drop would snap the
// whole body, faster than the servos can follow together. While walking,
// the gait slows down against the current pose in the meantime (see
// GaitEngine::limitTwist), so it stays physical during the handover.
void BodyPose::fitToReach(const FootTarget sets[][6], int count, float dt) {
  float fit = 1.0f;
  if (!reachable(sets, count, 1.0f)) {
    float lo = 0.0f, hi = 1.0f;   // lo stays feasible (or 0 if even no pose isn't - leave that to Limits)
    for (int n = 0; n < 8; n++) {
      float mid = 0.5f * (lo + hi);
      if (reachable(sets, count, mid)) lo = mid; else hi = mid;
    }
    fit = lo;
  }
  reach_scale_ = constrain(fit, reach_scale_ - REACH_DROP_PER_SEC * dt, reach_scale_ + REACH_RECOVER_PER_SEC * dt);
}