#include "Kinematics.h"
#include "Config.h"

// Standard analytical 3-DOF leg IK: coxa handles yaw in the horizontal plane,
// femur/tibia form a 2-link planar problem solved with the law of cosines.
//
// The elbow direction (which way tibia_rad's sign comes out) and the overall
// roll/pitch sign convention both depend on which way your leg is physically
// mounted - bench-check this on one leg with known foot positions before
// trusting all six.
JointAngles Kinematics::solve(int i, const FootTarget& target) {
  JointAngles out{};

  // Body-frame target -> leg-local frame: translate to the coxa pivot, then
  // rotate by -mount angle so the leg's local +X points along its resting
  // outward direction.
  float dx = target.x - LEG_MOUNT_X_MM[i];
  float dy = target.y - LEG_MOUNT_Y_MM[i];
  float a  = radians(-LEG_MOUNT_ANGLE_DEG[i]);
  float local_x = dx * cos(a) - dy * sin(a);
  float local_y = dx * sin(a) + dy * cos(a);
  float local_z = target.z;

  float coxa_rad = atan2(local_y, local_x);

  // Horizontal distance from the femur pivot (after the coxa segment) to the foot.
  float horiz = sqrt(local_x * local_x + local_y * local_y) - COXA_LEN_MM;
  float dist  = sqrt(horiz * horiz + local_z * local_z);

  out.valid = true;
  float max_reach = FEMUR_LEN_MM + TIBIA_LEN_MM;
  float min_reach = fabs(FEMUR_LEN_MM - TIBIA_LEN_MM);
  if (dist > max_reach) { dist = max_reach * 0.999f; out.valid = false; }
  if (dist < min_reach) { dist = min_reach * 1.001f; out.valid = false; }

  // Law of cosines: knee (tibia) angle from the femur/tibia/foot-distance triangle.
  float cos_tibia = (FEMUR_LEN_MM * FEMUR_LEN_MM + TIBIA_LEN_MM * TIBIA_LEN_MM - dist * dist)
                     / (2.0f * FEMUR_LEN_MM * TIBIA_LEN_MM);
  cos_tibia = constrain(cos_tibia, -1.0f, 1.0f);
  float tibia_rad = -(PI - acos(cos_tibia));   // knee-down convention; flip sign if yours mounts the other way

  // Femur angle: elevation to the foot, plus the triangle's femur-side angle.
  float elevation = atan2(-local_z, horiz);
  float cos_femur_tri = (FEMUR_LEN_MM * FEMUR_LEN_MM + dist * dist - TIBIA_LEN_MM * TIBIA_LEN_MM)
                         / (2.0f * FEMUR_LEN_MM * dist);
  cos_femur_tri = constrain(cos_femur_tri, -1.0f, 1.0f);
  float femur_rad = elevation + acos(cos_femur_tri);

  out.coxa_deg  = degrees(coxa_rad);
  out.femur_deg = degrees(femur_rad);
  out.tibia_deg = degrees(tibia_rad);

  if (out.coxa_deg  < COXA_MIN_DEG  || out.coxa_deg  > COXA_MAX_DEG)  out.valid = false;
  if (out.femur_deg < FEMUR_MIN_DEG || out.femur_deg > FEMUR_MAX_DEG) out.valid = false;
  if (out.tibia_deg < TIBIA_MIN_DEG || out.tibia_deg > TIBIA_MAX_DEG) out.valid = false;

  return out;
}

FootTarget Kinematics::neutralStance(int i) {
  float a = radians(LEG_MOUNT_ANGLE_DEG[i]);
  FootTarget t;
  t.x = LEG_MOUNT_X_MM[i] + cos(a) * (COXA_LEN_MM + STANCE_RADIUS_MM);
  t.y = LEG_MOUNT_Y_MM[i] + sin(a) * (COXA_LEN_MM + STANCE_RADIUS_MM);
  t.z = -STANCE_HEIGHT_MM;
  return t;
}