#include "HeaveController.hpp"
#include <Arduino.h>
#include <avr/pgmspace.h>

static const float MAX_INTEGRAL = 1e4f;

// Gap [mm] → equilibrium PWM (+ = repel, - = attract).
// 64 entries over 3–20 mm at 0.269841 mm steps. Copied from Controller.cpp.
static const int16_t HEAVE_FF_LUT[HEAVE_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
};

HeaveController::HeaveController(
    IndSensorL& f, IndSensorL& b,
    HeavePIDGains g, float avgRef, bool useFeedforward)
  : oor(false), outputOn(false),
    Front(f), Back(b),
    gains(g), state({0, 0, 0}),
    AvgRef(avgRef), avg(0), PWM(0), ffPWM(0),
    fullAttract(false), ffEnabled(useFeedforward)
{}

void HeaveController::update() {
  Front.readMM();
  Back.readMM();

  oor = Front.oor || Back.oor;

  avg = (Front.mmVal + Back.mmVal) * 0.5f;

  float e = AvgRef - avg;
  state.eDiff = e - state.e;
  if (!oor && !fullAttract) {
    state.eInt += e;
    state.eInt = constrain(state.eInt, -MAX_INTEGRAL, MAX_INTEGRAL);
  }
  state.e = e;

  ffPWM = ffEnabled ? feedforward(avg) : 0;

  if (fullAttract) {
    PWM = -HEAVE_CAP;  // manual override — ignore OOR, drive coils unconditionally
  } else if (oor) {
    // Sensor out of LUT range → trust feedforward alone (PID input is clamped).
    // avg is clamped to [mmMin, mmMax] so FF saturates toward repel when too close,
    // attract when too far — which is exactly the bring-up behavior.
    PWM = ffPWM;
  } else {
    // Gain-schedule PID output by (z/z_ref)² to linearize 1/z² force law.
    // Floor z to 2mm so sensor dropouts don't collapse control authority.
    // float z = max(avg, 2.0f);
    // float scale = (z * z) / (AvgRef * AvgRef);
    // PWM = constrain(ffPWM + (int16_t)(scale * pidCompute()), -HEAVE_CAP, HEAVE_CAP);    // In range: feedforward provides gravity/equilibrium bias, PID corrects residual.
    PWM = constrain(ffPWM + pidCompute(), -HEAVE_CAP, HEAVE_CAP);
  }
}

// Linearly interpolate the PROGMEM feedforward LUT.
int16_t HeaveController::feedforward(float gapMM) {
  if (gapMM <= HEAVE_FF_GAP_MIN) return (int16_t)pgm_read_word(&HEAVE_FF_LUT[0]);
  if (gapMM >= HEAVE_FF_GAP_MAX) return (int16_t)pgm_read_word(&HEAVE_FF_LUT[HEAVE_FF_LUT_SIZE - 1]);

  float idx_f = (gapMM - HEAVE_FF_GAP_MIN) / HEAVE_FF_GAP_STEP;
  uint8_t idx = (uint8_t)idx_f;
  if (idx >= HEAVE_FF_LUT_SIZE - 1) idx = HEAVE_FF_LUT_SIZE - 2;

  int16_t v0 = (int16_t)pgm_read_word(&HEAVE_FF_LUT[idx]);
  int16_t v1 = (int16_t)pgm_read_word(&HEAVE_FF_LUT[idx + 1]);
  float frac = idx_f - (float)idx;
  return (int16_t)(v0 + frac * (v1 - v0));
}

int16_t HeaveController::pidCompute() {
  if (oor) return 0;
  float out = gains.kp * state.e
            + gains.ki * state.eInt
            + gains.kd * state.eDiff;
  return (int16_t)constrain(out, -HEAVE_CAP, HEAVE_CAP);
}

void HeaveController::zeroPWMs() {
  PWM = 0;
}

// Drive all four motor channels identically. Direction bit mirrors the sign
// convention from Controller.cpp: LOW = repelling (PWM > 0), HIGH = attracting.
void HeaveController::sendOutputs() {
  if (!outputOn) zeroPWMs();

  bool attract = (PWM < 0);
  uint8_t mag = (uint8_t)abs(PWM);

  digitalWrite(dirFL, attract);
  digitalWrite(dirBL, attract);
  digitalWrite(dirFR, attract);
  digitalWrite(dirBR, attract);

  OCR2A = mag;  // Pin 11 (FL)
  OCR1A = mag;  // Pin 9  (FR)
  OCR2B = mag;  // Pin 3  (BL)
  OCR1B = mag;  // Pin 10 (BR)
}

void HeaveController::report() {
  // CSV: Front,Back,Avg,PWM,ControlOn
  Serial.print(Front.mmVal); Serial.print(',');
  Serial.print(Back.mmVal);  Serial.print(',');
  Serial.print(avg);         Serial.print(',');
  Serial.print(PWM);         Serial.print(',');
  Serial.println(outputOn);
}

void HeaveController::updatePID(HeavePIDGains g) { gains = g; }
void HeaveController::updateReference(float avgReference) { AvgRef = avgReference; }
void HeaveController::setFullAttract(bool enabled) {
  fullAttract = enabled;
  if (enabled) state.eInt = 0;  // drop stale integral so PID resume is clean
}