// Voltomni cart-pole kit: firmware for the Arduino Nano (CH340).
//
// Homes the cart on the left limit switch, centres it, then swings the pole up with energy control and
// balances it with LQR, at 500 Hz. The stepper drives the cart through a GT2 belt with a commanded
// ACCELERATION (the model in python/cartpole.py), so the cart position comes from the step count.
// Pole angle: AS5600 magnetic encoder over I2C.
//
// Wiring (see the guide): TMC2209 STEP=D2, DIR=D3, EN=D4 (MS1, MS2 to GND: 1/8 microstepping),
// AS5600 SDA=A4, SCL=A5, 5 V, GND; DIR pin of the AS5600 to GND.
// Limit switches (normally open, to GND): left D5, right D6.
//
// Serial console 115200 baud, one command per line:
//   h              home the cart and park it at the centre
//   s              home if needed, calibrate (pole hanging still) and start
//   x              stop and release the motor
//   l              toggle CSV logging at 100 Hz: ms,x,phi,dx,dphi,u,mode
//   g k1 k2 k3 k4  set LQR gains (u = -(k1*x + k2*phi + k3*dx + k4*dphi)); print from cartpole.py
//   e ke kc kd sm  set swing-up gains (energy, cart centring, cart damping, pump limit)
//   p              print settings and the current state

#include <Wire.h>
#include <math.h>

// ---- hardware -------------------------------------------------------------------------------
const uint8_t STEP_PIN = 2, DIR_PIN = 3, EN_PIN = 4;  // STEP is PD2 (pulsed in the timer ISR)
const uint8_t LEFT_PIN = 5, RIGHT_PIN = 6;
const float STEPS_PER_M = 200.0 * 8 / 0.040;          // 1.8° motor, 1/8 microstepping, 20T GT2 = 40 mm/rev
const uint8_t AS5600 = 0x36;
const float HOME_X = -0.167;   // cart position (m from centre) where the left switch closes
const float HOME_SPEED = 0.05; // m/s

// ---- control --------------------------------------------------------------------------------
const float DT = 0.002;        // 500 Hz
const float U_MAX = 15;        // cart acceleration limit, m/s²
const float V_MAX = 0.8;       // cart speed limit, m/s
const float X_LIMIT = 0.15;    // soft travel limit each side of centre, m (the switches are at ±0.167)
const float BRAKE = 25;        // deceleration when the soft limit is reached, m/s² (stops within 13 mm)
const float CATCH = 0.3;       // start balancing within this angle of upright, rad
const float DROP = 0.6;        // give up balancing beyond this, rad
const float W0SQ = 51.50;      // m*lc*g/Jp of the kit pole (cartpole.py prints it)

float K[4] = {-94.660, -144.786, -53.187, -20.157};  // from cartpole.py with the kit's default rig
float kEnergy = 0.2, kCentre = 80.0, kDamp = 2.0, swingMax = 6.0;

// ---- state ----------------------------------------------------------------------------------
volatile long steps = 0;        // cart position in microsteps (0 = centre once homed)
volatile int8_t stepDir = 1;
float vel = 0;                  // commanded cart speed, m/s
float downOffset = 0;           // encoder reading with the pole hanging
int8_t encSign = 1;             // set by calibration so that phi follows the model's sign
float phiPrev = 0, dphiF = 0;
bool homed = false, running = false, balancing = false, logging = false;
uint8_t logDiv = 0;
float lastU = 0;

// Timer1 makes one step pulse per compare match; OCR1A sets the step rate.
ISR(TIMER1_COMPA_vect) {
  PORTD |= _BV(PD2);
  steps += stepDir;
  delayMicroseconds(2);
  PORTD &= ~_BV(PD2);
}

void setStepRate(float stepsPerSec) {
  float r = fabs(stepsPerSec);
  if (r < 31) {  // below the timer's slowest rate: hold still
    TIMSK1 &= ~_BV(OCIE1A);
    return;
  }
  int8_t d = stepsPerSec > 0 ? 1 : -1;
  if (d != stepDir) {
    stepDir = d;
    digitalWrite(DIR_PIN, d > 0 ? HIGH : LOW);
    delayMicroseconds(2);
  }
  uint16_t ocr = (uint16_t)(2000000.0 / r) - 1;  // 16 MHz / 8
  noInterrupts();
  OCR1A = ocr;
  if (TCNT1 > ocr) TCNT1 = 0;
  interrupts();
  TIMSK1 |= _BV(OCIE1A);
}

float cartX() {
  noInterrupts();
  long s = steps;
  interrupts();
  return s / STEPS_PER_M;
}

void setCartX(float x) {
  noInterrupts();
  steps = lround(x * STEPS_PER_M);
  interrupts();
}

bool leftHit() { return digitalRead(LEFT_PIN) == LOW; }
bool rightHit() { return digitalRead(RIGHT_PIN) == LOW; }

float readEncoder() {  // radians, 0..2π
  Wire.beginTransmission(AS5600);
  Wire.write(0x0C);  // RAW ANGLE
  Wire.endTransmission(false);
  Wire.requestFrom(AS5600, (uint8_t)2);
  uint16_t raw = ((uint16_t)Wire.read() << 8) | Wire.read();
  return (raw & 0x0FFF) * (2 * PI / 4096.0);
}

float wrapPi(float a) {
  while (a > PI) a -= 2 * PI;
  while (a < -PI) a += 2 * PI;
  return a;
}

// Pole angle from upright (0 = up, ±π = hanging).
float poleAngle() { return wrapPi(encSign * (readEncoder() - downOffset) + PI); }

// Decelerate to rest at BRAKE, then release the motor (the cart would coast into the end stop otherwise).
void brakeAndStop() {
  while (fabs(vel) > 0.01) {
    unsigned long t0 = micros();
    vel -= constrain(vel, -BRAKE * DT, BRAKE * DT);
    setStepRate(vel * STEPS_PER_M);
    while (micros() - t0 < DT * 1e6) {}
  }
  stopMotor();
}

void stopMotor() {
  running = balancing = false;
  vel = 0;
  setStepRate(0);
  digitalWrite(EN_PIN, HIGH);  // driver off: cart free
}

// Move at a constant speed until `done` is true, a switch closes or the time runs out.
bool moveUntil(float speed, bool (*done)(), unsigned long timeoutMs) {
  unsigned long t0 = millis();
  setStepRate(speed * STEPS_PER_M);
  while (!done()) {
    if (millis() - t0 > timeoutMs) {
      setStepRate(0);
      return false;
    }
  }
  setStepRate(0);
  return true;
}

float target;
bool atTarget() { return fabs(cartX() - target) < 0.001 || leftHit() || rightHit(); }
bool pastTarget() { return cartX() >= target; }

bool moveTo(float x) {
  target = x;
  return moveUntil(x > cartX() ? HOME_SPEED * 2 : -HOME_SPEED * 2, atTarget, 8000) && !leftHit() && !rightHit();
}

// Drive left slowly until the left switch closes: that point is HOME_X. Then park at the centre.
bool home() {
  Serial.println(F("Homing..."));
  digitalWrite(EN_PIN, LOW);
  delay(200);
  if (leftHit()) {  // already on the switch: back off first
    setStepRate(HOME_SPEED * STEPS_PER_M);
    delay(400);
    setStepRate(0);
  }
  if (!moveUntil(-HOME_SPEED, leftHit, 12000)) {
    Serial.println(F("Left switch not found: check its wiring (D5 to GND when pressed)."));
    stopMotor();
    return false;
  }
  setCartX(HOME_X);
  target = HOME_X + 0.010;  // step off the switch (it's still closed, so not atTarget)
  moveUntil(HOME_SPEED, pastTarget, 2000);
  if (!moveTo(0)) {
    Serial.println(F("Couldn't reach the centre: a switch closed on the way. Check the switches."));
    stopMotor();
    return false;
  }
  homed = true;
  Serial.println(F("Homed: cart at the centre."));
  return true;
}

// Zero the pole angle with the pole hanging still, then find the encoder direction: a short
// positive cart acceleration must swing the hanging pole to positive angles (model sign).
bool calibrate() {
  Serial.println(F("Calibrating: keep the pole hanging still..."));
  digitalWrite(EN_PIN, LOW);
  delay(500);
  float s = 0, c = 0;
  for (int i = 0; i < 200; i++) {
    float a = readEncoder();
    s += sin(a);
    c += cos(a);
    delay(2);
  }
  downOffset = atan2(s, c);
  encSign = 1;
  // Jog: +3 m/s² for 80 ms, then back to rest (about 5 mm of travel).
  float peak = 0;
  vel = 0;
  for (int i = 0; i < 80; i++) {
    unsigned long t0 = micros();
    vel += (i < 40 ? 3 : -3) * DT;
    setStepRate(vel * STEPS_PER_M);
    float a = wrapPi(readEncoder() - downOffset);
    if (fabs(a) > fabs(peak)) peak = a;
    while (micros() - t0 < DT * 1e6) {}
  }
  vel = 0;
  setStepRate(0);
  if (fabs(peak) < 0.01) {
    Serial.println(F("Encoder didn't see the pole move: check the magnet gap and wiring."));
    stopMotor();
    return false;
  }
  encSign = peak > 0 ? 1 : -1;
  delay(1500);  // let it settle
  phiPrev = poleAngle();
  dphiF = 0;
  Serial.print(F("Encoder direction "));
  Serial.println(encSign);
  return true;
}

void control() {
  float x = cartX();
  float phi = poleAngle();
  float dphi = wrapPi(phi - phiPrev) / DT;
  phiPrev = phi;
  dphiF += 0.385 * (dphi - dphiF);  // ~50 Hz low-pass on the derivative
  float dx = vel;

  if (fabs(x) > X_LIMIT || leftHit() || rightHit()) {
    Serial.println(F("Travel limit reached: stopped. Send s to try again."));
    brakeAndStop();
    homed = false;  // the cart may have been pushed: home again
    return;
  }
  if (fabs(phi) < CATCH) balancing = true;
  else if (fabs(phi) > DROP) balancing = false;

  float u;
  if (balancing) {
    u = -(K[0] * x + K[1] * phi + K[2] * dx + K[3] * dphiF);
  } else {
    float E = 0.5 * dphiF * dphiF - W0SQ * (1 - cos(phi));
    float pump = kEnergy * E * (dphiF * cos(phi) >= 0 ? 1 : -1);
    pump = constrain(pump, -swingMax, swingMax);
    u = pump - kCentre * x - kDamp * dx;
  }
  u = constrain(u, -U_MAX, U_MAX);
  lastU = u;
  vel = constrain(vel + u * DT, -V_MAX, V_MAX);
  setStepRate(vel * STEPS_PER_M);

  if (logging && ++logDiv >= 5) {  // 100 Hz
    logDiv = 0;
    Serial.print(millis()); Serial.print(',');
    Serial.print(x, 4); Serial.print(',');
    Serial.print(phi, 4); Serial.print(',');
    Serial.print(dx, 3); Serial.print(',');
    Serial.print(dphiF, 3); Serial.print(',');
    Serial.print(u, 2); Serial.print(',');
    Serial.println(balancing ? 'B' : 'S');
  }
}

void printSettings() {
  Serial.print(F("LQR g "));
  for (int i = 0; i < 4; i++) { Serial.print(K[i], 3); Serial.print(' '); }
  Serial.print(F("| swing-up e "));
  Serial.print(kEnergy); Serial.print(' '); Serial.print(kCentre); Serial.print(' '); Serial.print(kDamp); Serial.print(' '); Serial.println(swingMax);
  Serial.print(homed ? F("cart ") : F("not homed; cart ")); Serial.print(cartX() * 1000, 1);
  Serial.print(F(" mm, pole ")); Serial.print(poleAngle(), 3); Serial.print(F(" rad from upright, switches L"));
  Serial.print(leftHit()); Serial.print(F(" R")); Serial.println(rightHit());
}

void handleCommand(char *line) {
  char *cmd = strtok(line, " ");
  if (!cmd) return;
  switch (cmd[0]) {
    case 'h': home(); break;
    case 's':
      if ((homed || home()) && calibrate()) {
        running = true;
        balancing = false;
        Serial.println(F("Running: swing-up, then balance."));
      }
      break;
    case 'x': brakeAndStop(); Serial.println(F("Stopped.")); break;
    case 'l': logging = !logging; if (logging) Serial.println(F("ms,x,phi,dx,dphi,u,mode")); break;
    case 'g': for (int i = 0; i < 4; i++) { char *t = strtok(NULL, " "); if (t) K[i] = atof(t); } printSettings(); break;
    case 'e': {
      float *v[4] = {&kEnergy, &kCentre, &kDamp, &swingMax};
      for (int i = 0; i < 4; i++) { char *t = strtok(NULL, " "); if (t) *v[i] = atof(t); }
      printSettings();
      break;
    }
    case 'p': printSettings(); break;
    default: Serial.println(F("Commands: h s x l g e p"));
  }
}

void setup() {
  pinMode(STEP_PIN, OUTPUT);
  pinMode(DIR_PIN, OUTPUT);
  pinMode(EN_PIN, OUTPUT);
  pinMode(LEFT_PIN, INPUT_PULLUP);
  pinMode(RIGHT_PIN, INPUT_PULLUP);
  digitalWrite(EN_PIN, HIGH);
  digitalWrite(DIR_PIN, HIGH);
  Serial.begin(115200);
  Wire.begin();
  Wire.setClock(400000);
  // Timer1: CTC mode, prescaler 8 (2 MHz)
  TCCR1A = 0;
  TCCR1B = _BV(WGM12) | _BV(CS11);
  OCR1A = 0xFFFF;
  Serial.println(F("Voltomni cart-pole. Let the pole hang still, then send s."));
}

void loop() {
  static char buf[64];
  static uint8_t n = 0;
  while (Serial.available()) {
    char c = Serial.read();
    if (c == '\n' || c == '\r') {
      if (n) { buf[n] = 0; handleCommand(buf); n = 0; }
    } else if (n < sizeof(buf) - 1) {
      buf[n++] = c;
    }
  }
  static unsigned long next = micros();
  if (running && (long)(micros() - next) >= 0) {
    next += (unsigned long)(DT * 1e6);
    control();
  } else if (!running) {
    next = micros();
  }
}