#include "SimBody.h"
#include <math.h>

namespace {

struct V { float x = 0, y = 0, z = 0; };
struct Q { float w = 1, x = 0, y = 0, z = 0; };

V operator+(V a, V b) { return { a.x + b.x, a.y + b.y, a.z + b.z }; }
V operator-(V a, V b) { return { a.x - b.x, a.y - b.y, a.z - b.z }; }
V operator*(V a, float s) { return { a.x * s, a.y * s, a.z * s }; }
float dot(V a, V b) { return a.x * b.x + a.y * b.y + a.z * b.z; }
V cross(V a, V b) { return { a.y * b.z - a.z * b.y, a.z * b.x - a.x * b.z, a.x * b.y - a.y * b.x }; }
float len(V a) { return sqrtf(dot(a, a)); }

Q mul(Q a, Q b) {
  return { a.w * b.w - a.x * b.x - a.y * b.y - a.z * b.z,
           a.w * b.x + a.x * b.w + a.y * b.z - a.z * b.y,
           a.w * b.y - a.x * b.z + a.y * b.w + a.z * b.x,
           a.w * b.z + a.x * b.y - a.y * b.x + a.z * b.w };
}
Q conj(Q q) { return { q.w, -q.x, -q.y, -q.z }; }
Q normalized(Q q) {
  float n = sqrtf(q.w * q.w + q.x * q.x + q.y * q.y + q.z * q.z);
  return { q.w / n, q.x / n, q.y / n, q.z / n };
}
V rotate(Q q, V v) {
  Q p = mul(mul(q, Q{ 0, v.x, v.y, v.z }), conj(q));
  return { p.x, p.y, p.z };
}
Q axisAngle(V axis, float rad) {
  float s = sinf(rad * 0.5f);
  return { cosf(rad * 0.5f), axis.x * s, axis.y * s, axis.z * s };
}
// Shortest rotation taking unit vector a onto unit vector b.
Q fromTo(V a, V b) {
  float d = dot(a, b);
  if (d < -0.9999f) {   // opposite: any perpendicular axis
    V axis = fabsf(a.x) < 0.9f ? cross(a, V{ 1, 0, 0 }) : cross(a, V{ 0, 1, 0 });
    return axisAngle(axis * (1.0f / len(axis)), (float)M_PI);
  }
  V c = cross(a, b);
  return normalized(Q{ 1 + d, c.x, c.y, c.z });
}

const V UP = { 0, 0, 1 };
constexpr float DEG = 180.0f / (float)M_PI;
constexpr float CONTACT_EPS_MM = 3.0f;
constexpr float MAX_GROUND_TILT_COS = 0.5f;   // body can't be tilted more than 60 deg against its ground

// ---- State ----
bool g_init = false;
V g_normal_prev = UP;              // ground normal in the body frame, last update
V g_ground_feet_prev[6];           // feet in the ground-aligned (un-yawed) frame, last update
int g_contact = 0, g_contact_prev = 0;
float g_px = 0, g_py = 0, g_yaw = 0;   // planar body position (mm) and heading (rad)
Q g_rot, g_rot_prev;
V g_pos;
V g_gyro_dps;
float g_tilt_roll = 0, g_tilt_pitch = 0;

Q groundRotation() {
  return mul(axisAngle(V{ 1, 0, 0 }, g_tilt_roll / DEG), axisAngle(V{ 0, 1, 0 }, g_tilt_pitch / DEG));
}

// The support plane: the plane the body rests on is one that every foot is
// on or above (a bottom face of the feet's convex hull). Of those, take the
// one most feet touch, then the one closest to last update's ground (the
// body doesn't flip onto a different face without reason).
bool supportPlane(const V feet[6], V& n_out, float& d_out) {
  int best_count = -1;
  float best_align = -2;
  for (int i = 0; i < 6; i++)
    for (int j = i + 1; j < 6; j++)
      for (int k = j + 1; k < 6; k++) {
        V c = cross(feet[j] - feet[i], feet[k] - feet[i]);
        float l = len(c);
        if (l < 1e-3f * 100 * 100) continue;   // feet (nearly) in a line: no plane
        V n = c * (1.0f / l);
        if (dot(n, g_normal_prev) < 0) n = n * -1.0f;   // ground normal points up toward the body
        float align = dot(n, g_normal_prev);
        if (align < MAX_GROUND_TILT_COS) continue;
        float d = dot(n, feet[i]);
        bool all_above = true;
        int count = 0;
        for (int m = 0; m < 6; m++) {
          float h = dot(n, feet[m]) - d;
          if (h < -CONTACT_EPS_MM) { all_above = false; break; }
          if (h < CONTACT_EPS_MM) count++;
        }
        if (!all_above) continue;
        if (count > best_count || (count == best_count && align > best_align)) {
          best_count = count;
          best_align = align;
          n_out = n;
          d_out = d;
        }
      }
  return best_count >= 3;
}

}  // namespace

namespace simbody {

void reset() {
  g_init = false;
  g_normal_prev = UP;
  g_contact = g_contact_prev = 0;
  g_px = g_py = g_yaw = 0;
  g_rot = g_rot_prev = Q{};
  g_pos = V{};
  g_gyro_dps = V{};
}

void setGroundTilt(float roll_deg, float pitch_deg) { g_tilt_roll = roll_deg; g_tilt_pitch = pitch_deg; }

void update(const float feet_mm[18], float dt) {
  V feet[6];
  for (int i = 0; i < 6; i++) feet[i] = { feet_mm[i * 3], feet_mm[i * 3 + 1], feet_mm[i * 3 + 2] };

  V n;
  float d;
  if (!supportPlane(feet, n, d)) {   // degenerate (e.g. limp legs in a line): keep last ground
    n = g_normal_prev;
    float lowest = 1e9f;
    for (auto& f : feet) lowest = fminf(lowest, dot(n, f));
    d = lowest;
  }
  g_contact = 0;
  for (int i = 0; i < 6; i++)
    if (dot(n, feet[i]) - d < CONTACT_EPS_MM) g_contact |= 1 << i;

  float height = -d;                       // body origin's distance above the ground
  Q tilt = fromTo(n, UP);                  // body frame -> ground-aligned frame

  // Planar motion: the rigid 2D move (Kabsch) that best carries the feet
  // that stayed grounded from their last positions to their new ones; the
  // body moved by the inverse of it.
  V ground_feet[6];
  for (int i = 0; i < 6; i++) ground_feet[i] = rotate(tilt, feet[i]);
  if (g_init) {
    int both = g_contact & g_contact_prev, m = 0;
    V c0, c1;
    for (int i = 0; i < 6; i++)
      if (both >> i & 1) { c0 = c0 + g_ground_feet_prev[i]; c1 = c1 + ground_feet[i]; m++; }
    if (m >= 2) {
      c0 = c0 * (1.0f / m);
      c1 = c1 * (1.0f / m);
      float s_cross = 0, s_dot = 0;
      for (int i = 0; i < 6; i++) {
        if (!(both >> i & 1)) continue;
        V a = g_ground_feet_prev[i] - c0, b = ground_feet[i] - c1;
        s_cross += a.x * b.y - a.y * b.x;
        s_dot += a.x * b.x + a.y * b.y;
      }
      float theta = atan2f(s_cross, s_dot);   // feet turned by theta under the body...
      float tx = c1.x - (c0.x * cosf(theta) - c0.y * sinf(theta));
      float ty = c1.y - (c0.x * sinf(theta) + c0.y * cosf(theta));
      g_yaw -= theta;                         // ...so the body turned by -theta
      g_px -= tx * cosf(g_yaw) - ty * sinf(g_yaw);
      g_py -= tx * sinf(g_yaw) + ty * cosf(g_yaw);
    }
  }
  for (int i = 0; i < 6; i++) g_ground_feet_prev[i] = ground_feet[i];
  g_contact_prev = g_contact;
  g_normal_prev = n;

  Q ground = groundRotation();
  Q rot = normalized(mul(mul(ground, axisAngle(UP, g_yaw)), tilt));
  g_pos = rotate(ground, V{ g_px, g_py, height });

  if (g_init && dt > 0) {
    Q dq = mul(conj(g_rot), rot);          // body-frame rotation over dt
    if (dq.w < 0) dq = Q{ -dq.w, -dq.x, -dq.y, -dq.z };
    float angle = 2.0f * acosf(fminf(1.0f, dq.w));
    float s = sqrtf(fmaxf(0.0f, 1.0f - dq.w * dq.w));
    g_gyro_dps = s > 1e-6f ? V{ dq.x / s, dq.y / s, dq.z / s } * (angle / dt * DEG) : V{};
  }
  g_rot_prev = g_rot;
  g_rot = rot;
  g_init = true;
}

void position(float out[3]) { out[0] = g_pos.x; out[1] = g_pos.y; out[2] = g_pos.z; }
void orientation(float out[4]) { out[0] = g_rot.w; out[1] = g_rot.x; out[2] = g_rot.y; out[3] = g_rot.z; }
float yawDeg() { return g_yaw * DEG; }
int contactMask() { return g_contact; }

void imu(float accel[3], float gyro[3]) {
  V a = rotate(conj(g_rot), UP);   // gravity reaction, seen from the body
  accel[0] = a.x; accel[1] = a.y; accel[2] = a.z;
  gyro[0] = g_gyro_dps.x; gyro[1] = g_gyro_dps.y; gyro[2] = g_gyro_dps.z;
}

}  // namespace simbody