#include "Controller.hpp"
#include <Arduino.h>
#include <math.h>

// ── Integral windup limit ────────────────────────────────────
static const float MAX_INTEGRAL = 1e4f;

// ── Feedforward LUT in PROGMEM ───────────────────────────────
// Generated by gen_ff_lut.py  (pod 9.4 kg, R=1.1Ω, V=12V, 3–20 mm)
// Positive = repelling, Negative = attracting
static const int16_t FF_PWM_LUT[FF_LUT_SIZE] PROGMEM = {
   238,  238,  238,  238,  238,  238,  238,  238,
   238,  238,  238,  238,  238,  238,  238,  238,
   238,  238,  234,  219,  204,  188,  172,  157,
   141,  125,  109,   93,   77,   61,   45,   29,
    13,   -3,  -19,  -35,  -51,  -67,  -84, -100,
  -116, -133, -150, -166, -183, -200, -217, -234,
  -250, -250, -250, -250, -250, -250, -250, -250,
  -250, -250, -250, -250, -250, -250, -250, -250
};

// ── Constructor ──────────────────────────────────────────────
FullController::FullController(
    IndSensor& l, IndSensor& r, IndSensor& f, IndSensor& b,
    PIDGains hGains, PIDGains rGains, PIDGains pGains,
    float avgRef, bool useFeedforward)
  : Left(l), Right(r), Front(f), Back(b),
    heightGains(hGains), rollGains(rGains), pitchGains(pGains),
    heightErr({0,0,0}), rollErr({0,0,0}), pitchErr({0,0,0}),
    AvgRef(avgRef), avg(0), ffEnabled(useFeedforward),
    oor(false), outputOn(false),
    FLPWM(0), BLPWM(0), FRPWM(0), BRPWM(0)
{}

// ── Main update (called each control tick) ───────────────────
void FullController::update() {
  // 1. Read all sensors (updates mmVal, oor)
  Left.readMM();
  Right.readMM();
  Front.readMM();
  Back.readMM();

  oor = Left.oor || Right.oor || Front.oor || Back.oor;

  // 2. Compute average gap (mm)
  avg = (Left.mmVal + Right.mmVal + Front.mmVal + Back.mmVal) * 0.25f;

  // 3. Feedforward: base PWM from equilibrium LUT
  int16_t ffPWM = ffEnabled ? feedforward(avg) : 0;

  // 4. Height PID: error = reference - average gap
  float heightE = AvgRef - avg;
  heightErr.eDiff = heightE - heightErr.e;
  if (!oor) {
    heightErr.eInt += heightE;
    heightErr.eInt = constrain(heightErr.eInt, -MAX_INTEGRAL, MAX_INTEGRAL);
  }
  heightErr.e = heightE;
  int16_t heightPWM = pidCompute(heightGains, heightErr, CAP);

  // 5. Roll PID: angle-based error (matches simulation)
  //    roll_angle = asin((left - right) / y_distance)
  //    error = -roll_angle  (drive roll toward zero)
  float rollRatio = (Left.mmVal - Right.mmVal) / Y_DISTANCE_MM;
  rollRatio = constrain(rollRatio, -1.0f, 1.0f);
  float rollAngle = asinf(rollRatio);
  float rollE = -rollAngle;

  rollErr.eDiff = rollE - rollErr.e;
  if (!oor) {
    rollErr.eInt += rollE;
    rollErr.eInt = constrain(rollErr.eInt, -MAX_INTEGRAL, MAX_INTEGRAL);
  }
  rollErr.e = rollE;
  int16_t rollPWM = pidCompute(rollGains, rollErr, CAP / 2);

  // 6. Pitch PID: angle-based error (matches simulation)
  //    pitch_angle = asin((back - front) / x_distance)
  //    error = -pitch_angle  (drive pitch toward zero)
  float pitchRatio = (Back.mmVal - Front.mmVal) / X_DISTANCE_MM;
  pitchRatio = constrain(pitchRatio, -1.0f, 1.0f);
  float pitchAngle = asinf(pitchRatio);
  float pitchE = -pitchAngle;

  pitchErr.eDiff = pitchE - pitchErr.e;
  if (!oor) {
    pitchErr.eInt += pitchE;
    pitchErr.eInt = constrain(pitchErr.eInt, -MAX_INTEGRAL, MAX_INTEGRAL);
  }
  pitchErr.e = pitchE;
  int16_t pitchPWM = pidCompute(pitchGains, pitchErr, CAP / 2);

  // 7. Mix outputs (same sign convention as simulation)
  //    FL: ff + height - roll - pitch
  //    FR: ff + height + roll - pitch
  //    BL: ff + height - roll + pitch
  //    BR: ff + height + roll + pitch
  FLPWM = constrain(ffPWM + heightPWM - rollPWM - pitchPWM, -CAP, CAP);
  FRPWM = constrain(ffPWM + heightPWM + rollPWM - pitchPWM, -CAP, CAP);
  BLPWM = constrain(ffPWM + heightPWM - rollPWM + pitchPWM, -CAP, CAP);
  BRPWM = constrain(ffPWM + heightPWM + rollPWM + pitchPWM, -CAP, CAP);
}

// ── PID compute (single symmetric set of gains) ─────────────
int16_t FullController::pidCompute(PIDGains& gains, PIDState& state, float maxOutput) {
  if (oor) return 0;
  float out = gains.kp * state.e
            + gains.ki * state.eInt
            + gains.kd * state.eDiff;
  return (int16_t)constrain(out, -maxOutput, maxOutput);
}

// ── Feedforward: linearly interpolate PROGMEM LUT ────────────
int16_t FullController::feedforward(float gapMM) {
  if (gapMM <= FF_GAP_MIN) return (int16_t)pgm_read_word(&FF_PWM_LUT[0]);
  if (gapMM >= FF_GAP_MAX) return (int16_t)pgm_read_word(&FF_PWM_LUT[FF_LUT_SIZE - 1]);

  float idx_f = (gapMM - 1.0 - FF_GAP_MIN) / FF_GAP_STEP;
  uint8_t idx = (uint8_t)idx_f;
  if (idx >= FF_LUT_SIZE - 1) idx = FF_LUT_SIZE - 2;

  int16_t v0 = (int16_t)pgm_read_word(&FF_PWM_LUT[idx]);
  int16_t v1 = (int16_t)pgm_read_word(&FF_PWM_LUT[idx + 1]);
  float frac = idx_f - (float)idx;

  return (int16_t)(v0 + frac * (v1 - v0));
}

// ── Zero all PWMs ────────────────────────────────────────────
void FullController::zeroPWMs() {
  FLPWM = 0;
  BLPWM = 0;
  FRPWM = 0;
  BRPWM = 0;
}

// ── Send PWM values to hardware ──────────────────────────────
void FullController::sendOutputs() {
  if (!outputOn) {
    zeroPWMs();
  }

  // Direction: LOW = repelling (positive PWM), HIGH = attracting (negative PWM)
  // Using direct register writes for fast PWM mode set by setupFastPWM()
  digitalWrite(dirFL, FLPWM < 0);
  OCR2A = abs(FLPWM); // Pin 11 → Timer 2A
  digitalWrite(dirBL, BLPWM < 0);
  OCR1A = abs(BLPWM); // Pin 9  → Timer 1A
  digitalWrite(dirFR, FRPWM < 0);
  OCR2B = abs(FRPWM); // Pin 3  → Timer 2B
  digitalWrite(dirBR, BRPWM < 0);
  OCR1B = abs(BRPWM); // Pin 10 → Timer 1B
}

// ── Serial report ────────────────────────────────────────────
void FullController::report() {
  // CSV: Left,Right,Front,Back,Avg,FLPWM,BLPWM,FRPWM,BRPWM,ControlOn
  Serial.print(Left.mmVal);  Serial.print(',');
  Serial.print(Right.mmVal); Serial.print(',');
  Serial.print(Front.mmVal); Serial.print(',');
  Serial.print(Back.mmVal);  Serial.print(',');
  Serial.print(avg);         Serial.print(',');
  Serial.print(FLPWM);      Serial.print(',');
  Serial.print(BLPWM);      Serial.print(',');
  Serial.print(FRPWM);      Serial.print(',');
  Serial.print(BRPWM);      Serial.print(',');
  Serial.println(outputOn);
}

// ── Runtime tuning methods ───────────────────────────────────
void FullController::updateHeightPID(PIDGains gains) { heightGains = gains; }
void FullController::updateRollPID(PIDGains gains)   { rollGains   = gains; }
void FullController::updatePitchPID(PIDGains gains)  { pitchGains  = gains; }

void FullController::updateReference(float avgReference) {
  AvgRef = avgReference;
}

void FullController::setFeedforward(bool enabled) {
  ffEnabled = enabled;
}