// Arduino Electronic Boost Controller
// https://forum.arduino.cc/t/arduino-electronic-boost-controller/1417566

#include <Wire.h>
#include <LiquidCrystal_I2C.h>
#include <PID_v1.h>
#include "sTune.h"
#include <math.h>

// ================== USER SETTINGS ==================
static const double setpoint_kPa_init = 180.0;  // target boost
static double setpoint_kPa = setpoint_kPa_init;

// Air source (battery inflator) soft limit (it cuts ~250 kPa)
static const double SUPPLY_MAX_KPA = 245.0;   // safety headroom for guard
static const double OVERSHOOT_MARGIN = 10.0;  // start backing off above SP + margin

// PID gains for TRIM (15 ms sample time, feed-forward enabled)
static double Kp = 2.2, Ki = 0.34, Kd = 0.32;

// Feed-forward base and trim ranges
static const double BASE_DUTY_FRAC = 0.58;   // 60% base duty
static const double MAX_BASE_BACKOFF = 0.1;  // up to -10% base when guarding
static const double TRIM_FRAC = 0.14;        // PID trims ±14% around base

// Keep the valve in its linear band (final duty clamp)
static const uint8_t DUTY_CLAMP_MIN = 102;  // 40% of 255
static const uint8_t DUTY_CLAMP_MAX = 204;  // 80% of 255

// Pressure filtering
static const uint8_t FILTER_SAMPLES = 5;
static const unsigned FILTER_US_DELAY = 300;  // µs between ADC reads

// Integral zone (don’t integrate when far from target)
static const double I_ZONE_KPA = 15.0;

// Logging / LCD periods
static const unsigned LOG_MS = 80;   // ~12.5 Hz serial log
static const unsigned LCD_MS = 500;  // 2 Hz LCD
// =====================================================

LiquidCrystal_I2C lcd(0x3F, 16, 2);

const int pressurePin = A0;  // MPX sensor
const int solenoidPin = 11;  // OC2A (Timer2 PWM @ ~31 Hz)
const int buttonPin = 7;     // Autotune trigger button (to GND)

double PV_kPa, OUT_trim;
PID boostPID(&PV_kPa, &OUT_trim, &setpoint_kPa, Kp, Ki, Kd, DIRECT);

unsigned long lastLCD = 0, lastLog = 0;
uint8_t lastDuty = 0;
bool buttonPrevPressed = false;

// ---------- sTune integration ----------
enum Mode : uint8_t { MODE_NORMAL = 0,
                      MODE_AUTOTUNE = 1,
                      MODE_APPLY = 2
                    };
static Mode runMode = MODE_NORMAL;

// sTune input/output variables
static float tuneInput = 0.0f;   // mirrors PV
static float tuneOutput = 0.0f;  // raw duty during tuning

// construct tuner ONCE globally (no assignment!)
sTune tuner(&tuneInput, &tuneOutput,
            sTune::DampedOsc_PID,
            sTune::directIP,
            sTune::printSUMMARY);

// plant setup constants
static const float STUNE_INPUT_SPAN = 500.0f;    // kPa full span (MPX5500)
static const float STUNE_OUTPUT_SPAN = 255.0f;   // PWM 8-bit span
static const float STUNE_OUTPUT_START = 130.0f;  // ~51%
static const float STUNE_OUTPUT_STEP = 165.0f;   // ~65%
static const uint32_t STUNE_TEST_SEC = 3;
static const uint32_t STUNE_SETTLE_SEC = 1;
static const uint16_t STUNE_SAMPLES = 250;

// ---------- simple serial command parser ----------
static char _lineBuf[64];
static uint8_t _lineLen = 0;

void printHelp() {
  Serial.println(F("Cmds:"));
  Serial.println(F("  P=#  I=#  D=#          (set individual gains)"));
  Serial.println(F("  PID kp ki kd           (set all three)"));
  Serial.println(F("  GET                    (print current PID gains)"));
  Serial.println(F("  ATSTART                (start sTune autotune)"));
  Serial.println(F("  ATSTOP                 (abort autotune and return to normal)"));
  Serial.println(F("  ATAPPLY                (re-apply last found sTune gains)"));
  Serial.println(F("  SET TGT=###            (set setpoint kPa)"));
}

void handleSerial() {
  while (Serial.available()) {
    char c = Serial.read();
    if (c == '\r') continue;
    if (c == '\n') {
      _lineBuf[_lineLen] = '\0';
      if (_lineLen > 0) {
        char* s = _lineBuf;
        while (*s == ' ') s++;

        if ((s[0] == 'P' || s[0] == 'p') && s[1] == '=') {
          double v = atof(s + 2);
          Kp = v;
          boostPID.SetTunings(Kp, Ki, Kd);
          Serial.print(F("Kp="));
          Serial.println(Kp, 6);
        } else if ((s[0] == 'I' || s[0] == 'i') && s[1] == '=') {
          double v = atof(s + 2);
          Ki = v;
          boostPID.SetTunings(Kp, Ki, Kd);
          Serial.print(F("Ki="));
          Serial.println(Ki, 6);
        } else if ((s[0] == 'D' || s[0] == 'd') && s[1] == '=') {
          double v = atof(s + 2);
          Kd = v;
          boostPID.SetTunings(Kp, Ki, Kd);
          Serial.print(F("Kd="));
          Serial.println(Kd, 6);
        } else if ((s[0] == 'P' || s[0] == 'p') && (s[1] == 'I' || s[1] == 'i') && (s[2] == 'D' || s[2] == 'd')) {
          double nkp = Kp, nki = Ki, nkd = Kd;
          char* p = s + 3;
          while (*p == ' ') p++;
          if (*p) {
            nkp = atof(p);
          }
          while (*p && *p != ' ') p++;
          while (*p == ' ') p++;
          if (*p) {
            nki = atof(p);
          }
          while (*p && *p != ' ') p++;
          while (*p == ' ') p++;
          if (*p) {
            nkd = atof(p);
          }
          Kp = nkp;
          Ki = nki;
          Kd = nkd;
          boostPID.SetTunings(Kp, Ki, Kd);
          Serial.print(F("PID set -> Kp="));
          Serial.print(Kp, 6);
          Serial.print(F(", Ki="));
          Serial.print(Ki, 6);
          Serial.print(F(", Kd="));
          Serial.println(Kd, 6);
        } else if (strncasecmp(s, "GET", 3) == 0) {
          Serial.print(F("Kp="));
          Serial.print(Kp, 6);
          Serial.print(F(" Ki="));
          Serial.print(Ki, 6);
          Serial.print(F(" Kd="));
          Serial.println(Kd, 6);
        } else if (strncasecmp(s, "SET TGT=", 8) == 0) {
          double v = atof(s + 8);
          setpoint_kPa = v;
          Serial.print(F("Setpoint_kPa="));
          Serial.println(setpoint_kPa, 1);
        } else if (strncasecmp(s, "ATSTART", 7) == 0) {
          startAutotune();
        } else if (strncasecmp(s, "ATSTOP", 6) == 0) {
          abortAutotune();
        } else if (strncasecmp(s, "ATAPPLY", 7) == 0) {
          applyLastTunings(true);
        } else {
          printHelp();
        }
      }
      _lineLen = 0;
    } else {
      if (_lineLen < sizeof(_lineBuf) - 1) _lineBuf[_lineLen++] = c;
    }
  }
}

// ---------- Pressure read ----------
float getFilteredPressure_kPa() {
  long sum = 0;
  // uint32_t analogAquisitionTime = micros();
  for (uint8_t i = 0; i < FILTER_SAMPLES; i++) {
    sum += analogRead(pressurePin);
    delayMicroseconds(FILTER_US_DELAY);
  }
  // analogAquisitionTime = micros() - analogAquisitionTime;
  // Serial.print("Aqu: ");
  // Serial.println(analogAquisitionTime);

  const float avg = sum / (float)FILTER_SAMPLES;

  const float voltage = avg * (5.0f / 1023.0f);
  const float V0 = 0.20f;
  const float spanV = 4.70f;
  float kPa = (voltage - V0) * (500.0f / spanV);

  if (kPa < -50) kPa = -50;
  if (kPa > 550) kPa = 550;
  return kPa;
}

// ---------- Timer2 ~31 Hz PWM ----------
void setupTimer2_31Hz() {
  pinMode(solenoidPin, OUTPUT);
  TCCR2A = 0;
  TCCR2B = 0;
  TCCR2A |= _BV(WGM20);
  TCCR2A |= _BV(COM2A1);
  TCCR2B |= _BV(CS22) | _BV(CS21) | _BV(CS20);
  OCR2A = 0;
}
inline void pwmWrite31Hz(uint8_t duty) {
  OCR2A = duty;
}

// ---------- sTune helpers ----------
void startAutotune() {
  if (runMode == MODE_AUTOTUNE) {
    Serial.println(F("Autotune already running."));
    return;
  }

  tuneInput = (float)PV_kPa;
  tuneOutput = STUNE_OUTPUT_START;

  tuner.Configure(STUNE_INPUT_SPAN, STUNE_OUTPUT_SPAN,
                  STUNE_OUTPUT_START, STUNE_OUTPUT_STEP,
                  STUNE_TEST_SEC, STUNE_SETTLE_SEC, STUNE_SAMPLES);

  tuner.SetEmergencyStop((float)SUPPLY_MAX_KPA);

  runMode = MODE_AUTOTUNE;
  boostPID.SetMode(MANUAL);
  Serial.println(F("\n=== sTune: starting autotune (open-loop step test) ==="));
}

void abortAutotune() {
  if (runMode != MODE_AUTOTUNE) {
    Serial.println(F("Autotune not running."));
    return;
  }
  runMode = MODE_NORMAL;
  boostPID.SetMode(AUTOMATIC);
  Serial.println(F("sTune: aborted. Back to normal PID."));
}

void applyLastTunings(bool announce) {
  float kp_s, kii, kdd;
  tuner.GetAutoTunings(&kp_s, &kii, &kdd);

  double newKp = kp_s;
  double newKi = kp_s * kii;
  double newKd = (kdd > 0.0f) ? (kp_s / kdd) : 0.0;

  if (isfinite(newKp) && isfinite(newKi) && isfinite(newKd) && newKp > 0) {
    Kp = newKp;
    Ki = newKi;
    Kd = newKd;
    boostPID.SetTunings(Kp, Ki, Kd);
    if (announce) {
      Serial.print(F("sTune applied -> Kp="));
      Serial.print(Kp, 6);
      Serial.print(F(", Ki="));
      Serial.print(Ki, 6);
      Serial.print(F(", Kd="));
      Serial.println(Kd, 6);
    }
  } else {
    Serial.println(F("sTune: invalid tunings; keeping old gains."));
  }
}

// ---------- setup ----------
void setup() {
  setupTimer2_31Hz();
  lcd.init();
  lcd.backlight();
  Serial.begin(115200);

  pinMode(buttonPin, INPUT_PULLUP);  // button to GND, active LOW

  const int TRIM_MAX = (int)(TRIM_FRAC * 255.0 + 0.5);
  boostPID.SetOutputLimits(-TRIM_MAX, +TRIM_MAX);
  boostPID.SetSampleTime(15);
  boostPID.SetMode(AUTOMATIC);

  printHelp();
  Serial.println(F("PID serial ready."));
}

// ---------- loop ----------
void loop() {
  handleSerial();

  PV_kPa = getFilteredPressure_kPa();
  tuneInput = (float)PV_kPa;

  // Read button (active LOW)
  bool buttonNowPressed = (digitalRead(buttonPin) == LOW);

  // On rising edge of button (user presses), start autotune
  if (buttonNowPressed && !buttonPrevPressed && runMode == MODE_NORMAL) {
    startAutotune();
    buttonPrevPressed = buttonNowPressed;
  }

  if (PV_kPa >= SUPPLY_MAX_KPA) {
    pwmWrite31Hz(DUTY_CLAMP_MIN);
  }

  if (runMode == MODE_AUTOTUNE) {
    // Only run autotune while the button is held down
    if (buttonNowPressed) {
      uint8_t duty = (uint8_t)constrain((int)roundf(tuneOutput), DUTY_CLAMP_MIN, DUTY_CLAMP_MAX);
      pwmWrite31Hz(duty);
      lastDuty = duty;

      uint8_t st = tuner.Run();

      if (st == sTune::tunings) {
        applyLastTunings(true);
        runMode = MODE_APPLY;
      }
    } else {
      // Button released during autotune -> abort and return to normal
      abortAutotune();
    }
  } else {
    if (runMode == MODE_APPLY) {
      boostPID.SetMode(AUTOMATIC);
      runMode = MODE_NORMAL;
      Serial.println(F("sTune: switching to NORMAL with new gains."));
    }

    double baseFrac = BASE_DUTY_FRAC;
    double backoff = 0.0;
    if (PV_kPa > setpoint_kPa + OVERSHOOT_MARGIN) {
      backoff += (PV_kPa - (setpoint_kPa + OVERSHOOT_MARGIN)) / 40.0;
    }
    if (PV_kPa > SUPPLY_MAX_KPA - 8.0) {
      backoff += (PV_kPa - (SUPPLY_MAX_KPA - 8.0)) / 16.0;
    }
    if (backoff > MAX_BASE_BACKOFF) backoff = MAX_BASE_BACKOFF;
    baseFrac -= backoff;
    if (baseFrac < 0.40) baseFrac = 0.40;

    double error = setpoint_kPa - PV_kPa;
    double Ki_eff = Ki;
    if ((lastDuty >= DUTY_CLAMP_MAX && error > 0) || (lastDuty <= DUTY_CLAMP_MIN && error < 0)) {
      Ki_eff = 0;
    }
    if (fabs(error) > I_ZONE_KPA) {
      Ki_eff = 0;
    }
    boostPID.SetTunings(Kp, Ki_eff, Kd);
    boostPID.Compute();
    boostPID.SetTunings(Kp, Ki, Kd);

    const int baseDuty = (int)(baseFrac * 255.0 + 0.5);
    int duty = baseDuty + (int)OUT_trim;
    if (duty < DUTY_CLAMP_MIN) duty = DUTY_CLAMP_MIN;
    if (duty > DUTY_CLAMP_MAX) duty = DUTY_CLAMP_MAX;
    pwmWrite31Hz((uint8_t)duty);
    lastDuty = (uint8_t)duty;
  }

  const unsigned long now = millis();
  if (now - lastLog >= LOG_MS) {
    lastLog = now;
    const float dutyPct = lastDuty * 100.0f / 255.0f;
    const float trimPct = OUT_trim * 100.0f / 255.0f;
    Serial.print("P:");
    Serial.print(PV_kPa, 1);
    Serial.print("\tT:");
    Serial.print(setpoint_kPa, 0);
    Serial.print("\tD:");
    Serial.print(dutyPct, 1);
    Serial.print("\tTRIM%:");
    Serial.print(trimPct, 1);
    Serial.print("\tMODE:");
    Serial.print(runMode == MODE_NORMAL ? "N" : (runMode == MODE_AUTOTUNE ? "AT" : "AP"));
    Serial.println();
  }

  if (now - lastLCD >= LCD_MS) {
    lastLCD = now;
    const int dutyPctRounded = (int)(lastDuty * 100.0f / 255.0f + 0.5f);

    lcd.setCursor(0, 0);
    lcd.print("P:        kPa ");
    lcd.setCursor(3, 0);
    lcd.print(PV_kPa, 1);

    lcd.setCursor(0, 1);
    if (runMode == MODE_AUTOTUNE) {
      lcd.print("AT RUN  D:");
      lcd.setCursor(10, 1);
      lcd.print(dutyPctRounded);
    } else {
      lcd.print("T:      D:   ");
      lcd.setCursor(3, 1);
      lcd.print(setpoint_kPa, 0);
      lcd.setCursor(10, 1);
      lcd.print(dutyPctRounded);
    }
  }
}