#include <Arduino.h>


typedef enum {
  MODE_NORMAL,
  MODE_STEP_TEST,
  MODE_AWAIT_TUNE_CONFIRM
} Mode;


Mode currentMode = MODE_NORMAL;
bool simulationRunning = false;

float Kp = 7.0f;
float Ki = 0.5f;
float Kd = 0.5f;

float tau = 60.0f;
float Kplant = 2.5f;
float disturbanceAmp = 20.0f;


const unsigned long dt_ms = 100;
const float dt_s = dt_ms / 1000.0f;


float plantState = 20.0f;
const float ambient = 20.0f;

static unsigned long last = 0;
float integrator = 0.0f;
float prevMeasured = 0.0f;


static uint32_t rng_state;

#define MAX_SAMPLES 150
float timeData[MAX_SAMPLES];
float outputData[MAX_SAMPLES];
int sampleIndex = 0;

bool testRunning = false;
float stepInput = 0.0f;
float initialValue = 0.0f;
unsigned long testStartTime = 0;

float K_identified = 0.0f;
float tau_identified = 0.0f;
float theta_identified = 0.0f;

float Kp_new_pending = 0.0f;
float Ki_new_pending = 0.0f;
float Kd_new_pending = 0.0f;
static float prevDerivative = 0;

void seed_rng() {
  uint32_t s = analogRead(A0);
  if (s == 0) {
    s = 0x1234567;
  }
  rng_state = s;
}

static inline float randf_signed() {
  rng_state ^= rng_state << 13;
  rng_state ^= rng_state >> 17;
  rng_state ^= rng_state << 5;
  float v = (rng_state & 0x7FFFFFFF) / 2147483647.0f;
  return v * 2.0f - 1.0f;
}

int readRawADC() {
  return analogRead(A0);
}

float readSetpointFromADC(int raw) {
  float denom = (raw > 2000) ? 4095.0f : 1023.0f;
  return (raw / denom) * 50.0f;
}

float pidCompute(float setpoint, float measured) {
  float error = setpoint - measured;
  float P = Kp * error;

  float rawDerivative = (measured - prevMeasured) / dt_s;

  prevMeasured = measured;

  integrator += error * dt_s;
  float I = Ki * integrator;

  float n = 10;

  float tau;
  if (Kp != 0) {
    tau = Kd / (Kp * n);
  } else {
    tau = 0.05;
  }

  float alpha;
  if (tau > 0.00001f) {
    alpha = dt_s / (dt_s + tau);
  } else {
    alpha = 1.0f;
  }

  float derivative = alpha * rawDerivative + (1.0f - alpha) * prevDerivative;

  prevDerivative = derivative;

  float D = -Kd * derivative;
  float out_unsat = P + I + D;
  float out_sat = constrain(out_unsat, -100.0f, 100.0f);

  float excess = out_sat - out_unsat;
  integrator += excess / Kp * dt_s;
  
  return out_sat;
}

void updatePlant(float u) {
  float disturbance = randf_signed() * disturbanceAmp;
  float dy = -(plantState - ambient) / tau + Kplant * (u / 100.0f) + disturbance * 0.05f;
  plantState += dy * dt_s;
}

void pidReset() {
  integrator = 0.0f;
  prevMeasured = plantState;
}

void collectStepData() {
  if (!testRunning) return;

  if (sampleIndex >= MAX_SAMPLES) {
    testRunning = false;
    analyzeStepResponse();
    return;
  }

  float elapsedTime = (millis() - testStartTime) / 1000.0f;
  if (elapsedTime >= 0.2f) {
    timeData[sampleIndex] = elapsedTime;
    outputData[sampleIndex] = plantState;

    Serial.print(elapsedTime, 3);
    Serial.print(F(","));
    Serial.println(plantState, 3);

    sampleIndex++;

    if (sampleIndex > 50) {
      float sum = 0;
      for (int i = sampleIndex - 20; i < sampleIndex; i++) {
        sum += outputData[i];
      }
      float avg = sum / 20.0f;

      bool stable = true;
      for (int i = sampleIndex - 20; i < sampleIndex; i++) {
        if (abs(outputData[i] - avg) > abs(avg * 0.005f)) {
          stable = false;
          break;
        }
      }

      if (stable) {
        testRunning = false;
        analyzeStepResponse();
      }
    }
  }
}

void startStepTest(float stepValue) {
  currentMode = MODE_STEP_TEST;
  testRunning = true;
  stepInput = stepValue;
  initialValue = plantState;
  sampleIndex = 0;
  testStartTime = millis();
  
  pidReset();
  simulationRunning = true;
  Serial.println(F("Step test"));
  collectStepData();
}


void autoTunePID() {
  float ratio = (tau_identified > 0) ? theta_identified / tau_identified : 0;

  if (ratio < 0.1f) {
    Kp_new_pending = 1.2f * tau_identified / (K_identified * theta_identified);
    Ki_new_pending = Kp_new_pending / (2.0f * theta_identified);
    Kd_new_pending = Kp_new_pending * 0.5f * theta_identified;
  } else {
    Kp_new_pending = 0.9f * tau_identified / (K_identified * theta_identified);
    Ki_new_pending = Kp_new_pending / (3.0f * theta_identified);
    Kd_new_pending = Kp_new_pending * 0.4f * theta_identified;
  }
  
  Serial.println(F("\nMeasured param"));
  Serial.print(F("New Kp = ")); Serial.println(Kp_new_pending, 3);
  Serial.print(F("New Ki = ")); Serial.println(Ki_new_pending, 3);
  Serial.print(F("New Kd = ")); Serial.println(Kd_new_pending, 3);
  Serial.println(F("\nTo apply - send 'a'"));
  Serial.println(F("To keep - send 'k'"));
  
  currentMode = MODE_AWAIT_TUNE_CONFIRM;
}

void analyzeStepResponse() {
  Serial.println(F("Analyzing data..."));

  if (sampleIndex < 10) {
    Serial.println(F("Not enough data!"));
    currentMode = MODE_NORMAL;
    return;
  }

  float y0 = outputData[0];
  
  int steadyStart = max(0, sampleIndex - 20);
  float ySteady = 0;
  for (int i = steadyStart; i < sampleIndex; i++) {
    ySteady += outputData[i];
  }
  ySteady /= (sampleIndex - steadyStart);

  float deltaY = ySteady - y0;
  float deltaU = stepInput;

  K_identified = deltaY / deltaU;

  float threshold5 = y0 + deltaY * 0.05f;
  theta_identified = 0;

  float target63 = y0 + deltaY * 0.632f;
  float time63 = 0;

  for (int i = 1; i < sampleIndex; i++) {
    bool crossed = (deltaY > 0) ? (outputData[i] >= threshold5) : (outputData[i] <= threshold5);
    if (crossed && outputData[i-1] != outputData[i]) {
      float t1 = timeData[i-1], t2 = timeData[i];
      float y1 = outputData[i-1], y2 = outputData[i];
      theta_identified = t1 + (threshold5 - y1) * (t2 - t1) / (y2 - y1);
      break;
    }
  }

  for (int i = 1; i < sampleIndex; i++) {
    bool crossed = (deltaY > 0) ? (outputData[i] >= target63) : (outputData[i] <= target63);
    if (crossed && outputData[i-1] != outputData[i]) {
      float t1 = timeData[i-1], t2 = timeData[i];
      float y1 = outputData[i-1], y2 = outputData[i];
      time63 = t1 + (target63 - y1) * (t2 - t1) / (y2 - y1);
      break;
    }
  }

  tau_identified = time63 - theta_identified;
  
  if (K_identified <= 0 || tau_identified <= 0 || theta_identified <= 0) {
      Serial.println(F("Could not identify model parameters!"));
      currentMode = MODE_NORMAL;
      return;
  }
  Serial.println(F("\n--- Comparison with Real Model ---"));
  Serial.print(F("K:  Model=")); Serial.print(Kplant, 3);
  Serial.print(F(" | Measured=")); Serial.print(K_identified, 3);
  Serial.print(F(" | Error=")); Serial.print(abs(Kplant - K_identified) / Kplant * 100, 1); Serial.println(F("%"));

  Serial.print(F("τ:  Model=")); Serial.print(tau, 3);
  Serial.print(F(" | Measured=")); Serial.print(tau_identified, 3);
  Serial.print(F(" | Error=")); Serial.print(abs(tau - tau_identified) / tau * 100, 1); Serial.println(F("%"));
  Serial.print(F("θ: ")); Serial.println(theta_identified, 3);
  autoTunePID();
}

void printHelp() {
  Serial.println(F("\n--- Commands ---"));
  Serial.println(F("g - Start/Resume simulation"));
  Serial.println(F("r - Reset to NORMAL mode and PAUSE simulation"));
  Serial.println(F("t - Start Step Response Test (50%)"));
  Serial.println(F("T - Start Step Response Test (custom value)"));
  Serial.println(F("s - Show current status"));
  Serial.println(F("p - Show PID parameters"));
  Serial.println(F("m - Show identified model"));
  Serial.println(F("h - Show this help menu"));
  Serial.println(F("------------------"));
  Serial.println(F("Simulation is PAUSED. Send 'g' to start."));
}

void handleCommands() {
    if (!Serial.available()) return;

    char cmd = Serial.read();
    while(Serial.available()) Serial.read();

    switch (cmd) {
        case 'g':{
            simulationRunning = true;
            Serial.println(F("Running..."));
            break;
            }
        case 't':{
            startStepTest(50.0f);
            break;
            }
        case 'T':{
            Serial.println(F("Enter step value(0-100):"));
            while(!Serial.available()) {}
            float customStep = Serial.parseFloat();
            customStep = constrain(customStep, 0.0f, 100.0f);
            startStepTest(customStep);
            break;
            }
        case 'r':{
            currentMode = MODE_NORMAL;
            testRunning = false;
            simulationRunning = false;
            pidReset();
            Serial.println(F("Send 'g' to resume."));
            break;
            }
        case 's':{
            Serial.println(F("\nCurrent Status"));
            Serial.print(F("Mode: "));
            if (currentMode == MODE_NORMAL) Serial.println(F("NORMAL (PID)"));
            else if (currentMode == MODE_STEP_TEST) Serial.println(F("STEP TEST"));
            else if (currentMode == MODE_AWAIT_TUNE_CONFIRM) Serial.println(F("AWAITING TUNE CONFIRMATION"));
            Serial.print(F("Simulation: ")); Serial.println(simulationRunning ? "RUNNING" : "PAUSED");
            Serial.print(F("Plant state: ")); Serial.print(plantState, 2); Serial.println(F(" C"));
            Serial.print(F("Integrator: ")); Serial.println(integrator, 3);
            break;
            }
        case 'p':{
            Serial.println(F("\nPID Parameters"));
            Serial.print(F("Kp = ")); Serial.println(Kp, 3);
            Serial.print(F("Ki = ")); Serial.println(Ki, 3);
            Serial.print(F("Kd = ")); Serial.println(Kd, 3);
            break;
            }
        case 'm':{
            if (K_identified != 0.0f) {
                Serial.println(F("\nIdentified Model"));
                Serial.print(F(" K = ")); Serial.print(K_identified, 4);
                Serial.print(F(" τ = ")); Serial.print(tau_identified, 3);
                Serial.print(F(" θ = ")); Serial.print(theta_identified, 3);
            } else {
                Serial.println(F("No model identified yet."));
            }
            break;
            }
        case 'a':{
            if (currentMode == MODE_AWAIT_TUNE_CONFIRM) {
                Kp = Kp_new_pending;
                Ki = Ki_new_pending;
                Kd = Kd_new_pending;
                currentMode = MODE_NORMAL;
                pidReset();
                Serial.println(F("New PID parameters APPLIED."));
            }
            break;
        }
        case 'k': {
            if (currentMode == MODE_AWAIT_TUNE_CONFIRM) {
                currentMode = MODE_NORMAL;
                pidReset();
                Serial.println(F("Current PID parameters KEPT."));
            }
            break;
            }
        case 'h': {
            printHelp();
            break;
        }
    }
}

void setup() {
  Serial.begin(115200);
  seed_rng();
  printHelp();
  prevMeasured = plantState;
}

void loop() {
  handleCommands();
  
  if (!simulationRunning) {
    last = millis();
    return;
  }

  unsigned long now = millis();
  if (now - last < dt_ms) {
    return;
  }
  last = now;

  float u = 0;
  if (currentMode == MODE_STEP_TEST) {
    if (testRunning) {
      u = stepInput;
      collectStepData();
    }
  } else if (currentMode == MODE_NORMAL) {
    int raw = readRawADC();
    float sp = readSetpointFromADC(raw);
    float meas = plantState;
    u = pidCompute(sp, meas);

    Serial.print(F("sp:"));
    Serial.print(sp, 2);
    Serial.print(F(" meas:"));
    Serial.print(meas, 2);
    Serial.print(F(" u:"));
    Serial.print(u, 2);
    Serial.print(F(" err:"));
    Serial.println(meas - sp, 2);
    Serial.print(F(" time:"));
    Serial.println((float)(millis() / 1000.0f), 2);
  }
  
  if(currentMode == MODE_STEP_TEST || currentMode == MODE_NORMAL){
      updatePlant(u);
  }
}