#include <Arduino.h>
#include <WiFi.h>
#include <WiFiUdp.h>

#define TELEPLOT_URL "teleplot.fr"
#define TELEPLOT_PORT 36948
WiFiUDP udp;

#define WIFI_SSID "Wokwi-GUEST"
#define WIFI_PASSWORD ""
// Defining the WiFi channel speeds up the connection:
#define WIFI_CHANNEL 6

volatile int motorPos = 0;

#define ppr 100 // pulses per revolution

#define PWM_FREQ 5000
#define PIN_MOTOR_P 25
#define PIN_MOTOR_N 26

#define PWM_CHANNEL_0 0
#define PWM_CHANNEL_1 1

#define PIN_ENCODER_A 34
#define PIN_ENCODER_B 35

#define MIN(a,b) (((a)<(b)) ? (a) : (b))
#define MAX(a,b) (((a)>(b)) ? (a) : (b))

void IRAM_ATTR pinChangeInterrupt();
void setSpeed(int speed);


float mse_best = 0.0;
float kp_test = 1.5, ti_test = 9.0, td_test = 0.1;
float kp_best = 1.5, ti_best = 9.0, td_best = 0.1;

float mse = 0.0;
unsigned int mse_cnt = 0;


typedef struct {
    float kp;
    float ti;
    float td;
    float integral;
    float prev_err;
    unsigned long prev_ms;
} controller_t;

void pid_init(controller_t* pid, float kp, float ti, float td) {
    pid->kp = kp;
    pid->ti = ti;
    pid->td = td;
    pid->integral = 0.0;
    pid->prev_err = 0.0;
    pid->prev_ms = 0;
}

float pid_update(controller_t* pid, float setpos, float motorPos) {
    unsigned long now_ms = millis();
    unsigned long dt_ms = now_ms - pid->prev_ms;
    float dt = dt_ms / 1000.0;

    if (dt == 0) return 0;  // Prevent division by zero

    float err = setpos - motorPos;

    pid->integral += err * dt;
    pid->integral = fmax(-1000, fmin(1000, pid->integral));  // Anti-windup

    float derivative = (err - pid->prev_err) / dt;

    float cmd = pid->kp * (err + pid->integral / pid->ti + pid->td * derivative);
    cmd = fmax(-255, fmin(255, cmd));

    pid->prev_err = err;
    pid->prev_ms = now_ms;

    return cmd;
}


// PID instance
controller_t servo_pid;

void setup() {
  Serial.begin(115200);
  Serial.printf("\n\n\n\n\n");

  WiFi.begin(WIFI_SSID, WIFI_PASSWORD, WIFI_CHANNEL);
  while (WiFi.status() != WL_CONNECTED) {
    delay(100); Serial.print(".");
  }
  Serial.printf("\nConnected as %s\n", WiFi.localIP().toString().c_str());

  /* motor outputs */
  pinMode(PIN_MOTOR_P, OUTPUT);
  pinMode(PIN_MOTOR_N, OUTPUT);

  /* 5 kHz PWM, 8-bit resolution */
/*
  ledcSetup(PWM_CHANNEL_0, PWM_FREQ, 8);
  ledcSetup(PWM_CHANNEL_1, PWM_FREQ, 8);

  ledcAttachPin(PIN_MOTOR_P, PWM_CHANNEL_0);
  ledcAttachPin(PIN_MOTOR_N, PWM_CHANNEL_1);
*/
  ledcAttach(PIN_MOTOR_P, 10000, 8);
  ledcAttach(PIN_MOTOR_N, 10000, 8);


  /* quadrature encoder inputs */
  pinMode(PIN_ENCODER_A, INPUT);
  pinMode(PIN_ENCODER_B, INPUT);

  motorPos = -100;
  attachInterrupt(PIN_ENCODER_A, pinChangeInterrupt, CHANGE);

  Serial.printf("setPos, motorPos, err, der, int, cmd\n");
  pid_init(&servo_pid, kp_best, ti_best, td_best);
}


/**
 * @brief Set motor speed
 * @param speed (-255: full speed CCW ... 0: stop ... +255: full speed CW)
 */
void setSpeed(int speed) {
  speed = MAX(-255, MIN(255, speed));

  if (speed > 0) {
    ledcWrite(PIN_MOTOR_P, 255);
    ledcWrite(PIN_MOTOR_N, 255 - speed);
  } else {
    ledcWrite(PIN_MOTOR_P, 255 + speed);
    ledcWrite(PIN_MOTOR_N, 255);
  }
}


void loop() {
  /* Read user input */
  unsigned long now_ms = millis();

  /* Generate square signal for setpos */
  int setpos = ((now_ms % 4000) < 2000) ? 20 : 0;
  
  /* PID regulation */
  float cmd = pid_update(&servo_pid, setpos, motorPos);  

  /* Actuate motor */
  setSpeed((int)cmd);
  
  /* Print out data for display (use the Serial plotter to graph out data) */
  Serial.printf("%9d, %9d, %9d\n", setpos, motorPos, (int)cmd);

  /* PID auto-tuning */
  float err = setpos - motorPos;
  mse += err * err;
  mse_cnt ++;

  static unsigned long last_tune_ms = 0;
  const unsigned long tune_interval_ms = 12000;
  if (now_ms - last_tune_ms > tune_interval_ms) {
    mse = mse / mse_cnt;

    setSpeed(0);

    Serial.printf("Params %3.1f, %3.1f, %3.1f:", 
        kp_test, ti_test, td_test);
    Serial.printf("MSE current %9.3f, best %9.3f\n", mse, mse_best);

    sendTelemetry("mse", mse);
    sendTelemetry("mse_best", mse_best);
    sendTelemetry("kp_test", kp_test);
    sendTelemetry("ti_test", ti_test);
    sendTelemetry("td_test", td_test);

    if (now_ms > 2 * tune_interval_ms) {
      if ((mse < mse_best) || (mse_best == 0.0)) {
        // use new parameter set
        mse_best = mse;
        kp_best = kp_test;
        ti_best = ti_test;
        td_best = td_test;
        Serial.printf("Updated set %3.1f, %3.1f, %3.1, %3.1f\n", kp_best, ti_best, td_best, mse_best);
      } else {
        // roll back
        kp_test = kp_best;
        ti_test = ti_best;
        td_test = td_best;
      }

      // randomly update one parameter
      int param = rand() % 3;
      float factor = (rand() % 2 == 0) ? 1.1 : 0.9;

      switch(param) {
        case 0: kp_test *= factor; break;
        case 1: ti_test *= factor; break;
        default: td_test *= factor; break;
      }
    }

    // reset PID regulation parameters
    pid_init(&servo_pid, kp_test, ti_test, td_test);
    mse = 0.0;
    mse_cnt = 0;
    last_tune_ms = now_ms;
  }

  delay(100);
}


/**
 * @brief Read quadrature encoder signals, and update shaft position
 */
void IRAM_ATTR pinChangeInterrupt() {
  int a = digitalRead(PIN_ENCODER_A);
  int b = digitalRead(PIN_ENCODER_B);

  motorPos += (a ^ b) ? 1 : -1;
}


void sendTelemetry(const char *signal, float value)
{
  static char buffer[256];
  uint32_t now_ms = millis();

  snprintf(buffer, sizeof(buffer), "%s:%3.3f",
      signal, value);
  udp.beginPacket(TELEPLOT_URL, TELEPLOT_PORT);
  udp.write((const uint8_t*)buffer, strlen(buffer));
  udp.endPacket();
}

void sendTelemetry(const char *signal, int value)
{
  static char buffer[256];
  uint32_t now_ms = millis();

  snprintf(buffer, sizeof(buffer), "%s:%d",
      signal, value);
  udp.beginPacket(TELEPLOT_URL, TELEPLOT_PORT);
  udp.write((const uint8_t*)buffer, strlen(buffer));
  udp.endPacket();
}