#include "BodyPose.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 {
  float roll  = radians(s_roll_.current + correction_roll_);
  float pitch = radians(s_pitch_.current + correction_pitch_);
  float yaw   = radians(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_x_.current;
  out.y = y3 - s_y_.current;
  out.z = z3 - s_z_.current;
  return out;
}