// Voltomni rotary inverted pendulum kit: firmware for the Arduino Nano (CH340).
//
// Swings the pendulum up from rest with energy control and balances it with LQR, at 500 Hz.
// The stepper is driven with a commanded arm ACCELERATION (the model in python/furuta.py), so the
// arm angle comes from the step count. Pendulum 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.
//
// Serial console 115200 baud, one command per line:
//   s              calibrate (pendulum hanging still) and start
//   x              stop and release the motor
//   l              toggle CSV logging at 100 Hz: ms,theta,phi,dtheta,dphi,u,mode
//   g k1 k2 k3 k4  set LQR gains (u = -(k1*theta + k2*phi + k3*dtheta + k4*dphi)); print from furuta.py
//   e ke kc kd sm  set swing-up gains (energy, arm centring, arm damping, pump limit)
//   p              print settings and the current angles

#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 float STEPS_PER_REV = 200.0 * 8;                // 1.8° motor, 1/8 microstepping
const float STEPS_PER_RAD = STEPS_PER_REV / (2 * PI);
const uint8_t AS5600 = 0x36;

// ---- control --------------------------------------------------------------------------------
const float DT = 0.002;        // 500 Hz
const float U_MAX = 60;        // arm acceleration limit, rad/s²
const float V_MAX = 20;        // arm speed limit, rad/s
const float THETA_LIMIT = 2.5; // arm travel limit (cable), rad
const float CATCH = 0.35;      // start balancing within this angle of upright, rad
const float DROP = 0.6;        // give up balancing beyond this, rad
const float W0SQ = 85.84;      // m*lc*g/Jp of the kit pendulum (furuta.py prints it)

float K[4] = {-15.384, -334.960, -10.907, -36.112};  // from furuta.py with the kit's default rig
float kEnergy = 1.0, kCentre = 80.0, kDamp = 5.0, swingMax = 60.0;

// ---- state ----------------------------------------------------------------------------------
volatile long steps = 0;        // arm position in microsteps
volatile int8_t stepDir = 1;
float vel = 0;                  // commanded arm speed, rad/s
float downOffset = 0;           // encoder reading with the pendulum hanging
int8_t encSign = 1;             // set by calibration so that phi follows the model's sign
float phiPrev = 0, dphiF = 0;
bool 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 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;
}

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

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

// Zero both angles with the pendulum hanging still, then find the encoder direction: a short
// positive arm acceleration must swing the hanging pendulum to positive angles (model sign).
bool calibrate() {
  Serial.println(F("Calibrating: keep the pendulum 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);
  noInterrupts();
  steps = 0;
  interrupts();
  encSign = 1;
  // Jog: +30 rad/s² for 80 ms, then back to rest.
  float peak = 0;
  vel = 0;
  for (int i = 0; i < 80; i++) {
    unsigned long t0 = micros();
    vel += (i < 40 ? 30 : -30) * DT;
    setStepRate(vel * STEPS_PER_RAD);
    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 pendulum move: check the magnet gap and wiring."));
    stopMotor();
    return false;
  }
  encSign = peak > 0 ? 1 : -1;
  delay(1500);  // let it settle
  noInterrupts();
  steps = 0;  // the jog moved the arm slightly: this is the arm's zero
  interrupts();
  phiPrev = pendulumAngle();
  dphiF = 0;
  Serial.print(F("Encoder direction "));
  Serial.println(encSign);
  return true;
}

void control() {
  float theta = steps / STEPS_PER_RAD;
  float phi = pendulumAngle();
  float dphi = wrapPi(phi - phiPrev) / DT;
  phiPrev = phi;
  dphiF += 0.385 * (dphi - dphiF);  // ~50 Hz low-pass on the derivative
  float dtheta = vel;

  if (fabs(theta) > THETA_LIMIT) {
    Serial.println(F("Arm limit reached: stopped. Send s to try again."));
    stopMotor();
    return;
  }
  if (fabs(phi) < CATCH) balancing = true;
  else if (fabs(phi) > DROP) balancing = false;

  float u;
  if (balancing) {
    u = -(K[0] * theta + K[1] * phi + K[2] * dtheta + 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 * theta - kDamp * dtheta;
  }
  u = constrain(u, -U_MAX, U_MAX);
  lastU = u;
  vel = constrain(vel + u * DT, -V_MAX, V_MAX);
  setStepRate(vel * STEPS_PER_RAD);

  if (logging && ++logDiv >= 5) {  // 100 Hz
    logDiv = 0;
    Serial.print(millis()); Serial.print(',');
    Serial.print(theta, 4); Serial.print(',');
    Serial.print(phi, 4); Serial.print(',');
    Serial.print(dtheta, 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(F("arm ")); Serial.print(steps / STEPS_PER_RAD, 3);
  Serial.print(F(" rad, pendulum ")); Serial.print(pendulumAngle(), 3); Serial.println(F(" rad from upright"));
}

void handleCommand(char *line) {
  char *cmd = strtok(line, " ");
  if (!cmd) return;
  switch (cmd[0]) {
    case 's':
      if (calibrate()) {
        running = true;
        balancing = false;
        Serial.println(F("Running: swing-up, then balance."));
      }
      break;
    case 'x': stopMotor(); Serial.println(F("Stopped.")); break;
    case 'l': logging = !logging; if (logging) Serial.println(F("ms,theta,phi,dtheta,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: s x l g e p"));
  }
}

void setup() {
  pinMode(STEP_PIN, OUTPUT);
  pinMode(DIR_PIN, OUTPUT);
  pinMode(EN_PIN, OUTPUT);
  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 rotary pendulum. Let the pendulum 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();
  }
}