// Voltomni self-balancing robot kit: firmware for the ESP32-C3 SuperMini (Arduino core 3.x).
//
// Balances with LQR at 500 Hz on the tilt from an MPU-6050 (complementary filter). The steppers are
// driven with a commanded wheel ACCELERATION (the model in python/balancer.py), so the distance
// travelled comes from the step counts. The robot is its own Wi-Fi access point and serves a page
// with a joystick, live tilt and battery, and the controller gains.
//
// Wiring (see the guide): left TMC2209 STEP=GPIO0 DIR=GPIO1, right TMC2209 STEP=GPIO3 DIR=GPIO10,
// both EN=GPIO21 (MS1, MS2 to GND: 1/8 microstepping); MPU-6050 SDA=GPIO6 SCL=GPIO7, 3.3 V, GND,
// X arrow pointing forward; battery divider (100 k / 22 k) to GPIO4.
//
// Wi-Fi: network "Voltomni-Balancer-XXXX", password "voltomni"; then open http://192.168.4.1
// Serial (USB) at 115200 prints the network name and status.

#include <WiFi.h>
#include <WebServer.h>
#include <Wire.h>
#include <Preferences.h>
#include <soc/gpio_struct.h>
#include <math.h>
#include "page.h"

// ---- hardware -------------------------------------------------------------------------------
const int STEP_L = 0, DIR_L = 1, STEP_R = 3, DIR_R = 10, EN_PIN = 21;
const int SDA_PIN = 6, SCL_PIN = 7, VBAT_PIN = 4;
const float VBAT_DIVIDER = (100.0 + 22.0) / 22.0;
const float VBAT_LOW = 10.5;            // 3S: 3.5 V per cell
const uint8_t MPU = 0x68;
const float WHEEL_R = 0.045;            // printed wheel with its rubber tyre, m
const float STEPS_PER_M = 200.0 * 8 / (2 * PI * WHEEL_R);
const uint32_t TICK_HZ = 40000;         // step timer; up to 20 000 steps/s per motor
const char *AP_PASSWORD = "voltomni";

// ---- control --------------------------------------------------------------------------------
const float DT = 0.002;                 // 500 Hz
const float A_MAX = 6.0;                // wheel acceleration limit, m/s²
const float V_MAX = 1.0;                // wheel speed limit, m/s
const float FALL = 0.6;                 // tilt at which it gives up, rad
const float DRIVE_SPEED = 0.4;          // full joystick forward, m/s
const float TURN_SPEED = 0.25;          // full joystick sideways: wheel speed difference, m/s

float K[4] = {-5.974, -68.247, -11.940, -6.968};  // from balancer.py with the kit's default robot
float trim = 0;                         // tilt reading when balanced (set by Calibrate), rad
float gyroBias = 0;                     // rad/s

// ---- state ----------------------------------------------------------------------------------
volatile int32_t posL = 0, posR = 0;    // steps, forward positive
volatile uint32_t incL = 0, incR = 0;   // phase increments per tick (2^32 = one step)
volatile int8_t dirL = 1, dirR = 1;
uint32_t accL = 0, accR = 0;
bool pulseL = false, pulseR = false;

volatile float tilt = 0, tiltRate = 0, vbat = 0, vel = 0;
volatile float vRef = 0, turnRef = 0;
volatile bool balancing = false, lowBattery = false, motorTest = false;
volatile bool calRequest = false, calDone = false;  // calibration runs in the control task (it owns the I2C bus)
volatile uint32_t lastDrive = 0;
float xRef = 0;

WebServer server(80);
Preferences prefs;
hw_timer_t *timer;

// ---- steppers -------------------------------------------------------------------------------
void IRAM_ATTR onTick() {
  uint32_t set = 0, clr = 0;
  if (pulseL) { clr |= 1 << STEP_L; pulseL = false; }
  else if (incL) {
    uint32_t a = accL + incL;
    if (a < accL) { set |= 1 << STEP_L; pulseL = true; posL += dirL; }
    accL = a;
  }
  if (pulseR) { clr |= 1 << STEP_R; pulseR = false; }
  else if (incR) {
    uint32_t a = accR + incR;
    if (a < accR) { set |= 1 << STEP_R; pulseR = true; posR += dirR; }
    accR = a;
  }
  if (set) GPIO.out_w1ts.val = set;
  if (clr) GPIO.out_w1tc.val = clr;
}

// Wheel speeds in m/s, forward positive. The right motor faces the other way, so its DIR is inverted.
void setWheels(float vl, float vr) {
  int8_t dl = vl >= 0 ? 1 : -1, dr = vr >= 0 ? 1 : -1;
  if (dl != dirL) { digitalWrite(DIR_L, dl > 0 ? HIGH : LOW); dirL = dl; }
  if (dr != dirR) { digitalWrite(DIR_R, dr > 0 ? LOW : HIGH); dirR = dr; }
  incL = (uint32_t)(fminf(fabsf(vl) * STEPS_PER_M, TICK_HZ / 2) * (4294967296.0f / TICK_HZ));
  incR = (uint32_t)(fminf(fabsf(vr) * STEPS_PER_M, TICK_HZ / 2) * (4294967296.0f / TICK_HZ));
}

void motors(bool on) { digitalWrite(EN_PIN, on ? LOW : HIGH); }

// ---- IMU ------------------------------------------------------------------------------------
void mpuWrite(uint8_t reg, uint8_t v) {
  Wire.beginTransmission(MPU);
  Wire.write(reg);
  Wire.write(v);
  Wire.endTransmission();
}

bool mpuBegin() {
  mpuWrite(0x6B, 0x00);  // wake
  mpuWrite(0x1A, 0x03);  // DLPF 44 Hz
  mpuWrite(0x1B, 0x08);  // gyro ±500 °/s
  mpuWrite(0x1C, 0x08);  // accel ±4 g
  Wire.beginTransmission(MPU);
  return Wire.endTransmission() == 0;
}

// Tilt from the accelerometer and pitch rate from the gyro (X forward, Z up: leaning forward is +).
bool mpuRead(float &accTilt, float &rate) {
  Wire.beginTransmission(MPU);
  Wire.write(0x3B);
  if (Wire.endTransmission(false) != 0) return false;
  if (Wire.requestFrom(MPU, (uint8_t)14) != 14) return false;
  int16_t d[7];
  for (int i = 0; i < 7; i++) d[i] = (Wire.read() << 8) | Wire.read();
  accTilt = atan2f(-(float)d[0], (float)d[2]);
  rate = d[5] / 65.5f * (PI / 180) - gyroBias;  // gyro Y
  return true;
}

// ---- control task (500 Hz) ------------------------------------------------------------------
void control(void *) {
  TickType_t wake = xTaskGetTickCount();
  float est = 0, calTilt = 0, calRate = 0;
  int calLeft = 0;
  uint32_t uprightSince = 0;
  for (;;) {
    vTaskDelayUntil(&wake, pdMS_TO_TICKS(2));
    float accTilt, rate;
    if (!mpuRead(accTilt, rate)) continue;
    est = 0.98f * (est + rate * DT) + 0.02f * accTilt;
    tilt = est - trim;
    tiltRate = rate;

    // Calibrate (robot held upright and still): average the tilt reading and the gyro over 0.8 s.
    if (calRequest) {
      calRequest = false;
      balancing = false;
      setWheels(0, 0);
      motors(false);
      calLeft = 400;
      calTilt = calRate = 0;
    }
    if (calLeft > 0) {
      calTilt += accTilt;
      calRate += rate + gyroBias;
      if (--calLeft == 0) {
        trim = calTilt / 400;
        gyroBias = calRate / 400;
        calDone = true;
      }
      uprightSince = 0;
      continue;
    }

    if (motorTest) continue;  // the web page is running the motors
    if (millis() - lastDrive > 500) { vRef = 0; turnRef = 0; }  // joystick released or page closed

    if (!balancing) {
      // Stand it up and hold it still for half a second to start.
      if (fabsf(tilt) < 0.05f && !lowBattery) {
        if (!uprightSince) uprightSince = millis();
        if (millis() - uprightSince > 500) {
          posL = posR = 0;
          xRef = 0;
          vel = 0;
          motors(true);
          balancing = true;
        }
      } else uprightSince = 0;
      continue;
    }
    if (fabsf(tilt) > FALL || lowBattery) {
      balancing = false;
      uprightSince = 0;
      setWheels(0, 0);
      motors(false);
      continue;
    }
    float x = (posL + posR) * 0.5f / STEPS_PER_M;
    xRef += vRef * DT;
    float a = -(K[0] * (x - xRef) + K[1] * tilt + K[2] * (vel - vRef) + K[3] * tiltRate);
    a = constrain(a, -A_MAX, A_MAX);
    vel = constrain(vel + a * DT, -V_MAX, V_MAX);
    if (fabsf(x - xRef) > 0.5f) xRef = x - copysignf(0.5f, x - xRef);  // don't wind up after a push
    setWheels(vel + turnRef, vel - turnRef);
  }
}

// ---- web page -------------------------------------------------------------------------------
void sendState() {
  char buf[220];
  snprintf(buf, sizeof buf,
           "{\"tilt\":%.3f,\"rate\":%.3f,\"vbat\":%.2f,\"balancing\":%s,\"low\":%s,\"k\":[%.3f,%.3f,%.3f,%.3f],\"trim\":%.4f}",
           tilt, tiltRate, vbat, balancing ? "true" : "false", lowBattery ? "true" : "false", K[0], K[1], K[2], K[3], trim);
  server.send(200, "application/json", buf);
}

void setupWeb() {
  server.on("/", [] { server.send_P(200, "text/html", PAGE); });
  server.on("/state", sendState);
  server.on("/drive", HTTP_POST, [] {  // joystick: x, y in -1..1, sent every 100 ms while held
    float jx = constrain(server.arg("x").toFloat(), -1, 1), jy = constrain(server.arg("y").toFloat(), -1, 1);
    vRef = jy * DRIVE_SPEED;
    turnRef = jx * TURN_SPEED;
    lastDrive = millis();
    server.send(204);
  });
  server.on("/gains", HTTP_POST, [] {
    const char *names[4] = {"k1", "k2", "k3", "k4"};
    for (int i = 0; i < 4; i++)
      if (server.hasArg(names[i])) { K[i] = server.arg(names[i]).toFloat(); prefs.putFloat(names[i], K[i]); }
    if (server.hasArg("trim")) { trim = server.arg("trim").toFloat(); prefs.putFloat("trim", trim); }
    sendState();
  });
  server.on("/calibrate", HTTP_POST, [] {  // hold the robot upright and still while it runs (about 1 s)
    calDone = false;
    calRequest = true;
    for (int i = 0; i < 150 && !calDone; i++) delay(10);
    if (!calDone) return server.send(500, "text/plain", "IMU not responding");
    prefs.putFloat("trim", trim);
    prefs.putFloat("gbias", gyroBias);
    Serial.printf("Calibrated: trim %.4f rad, gyro bias %.4f rad/s\n", trim, gyroBias);
    sendState();
  });
  server.on("/test", HTTP_POST, [] {  // both wheels forward slowly for 2 s (robot lying on its back)
    if (balancing) return server.send(409, "text/plain", "Lay it down first");
    motorTest = true;
    motors(true);
    setWheels(0.1, 0.1);
    delay(2000);
    setWheels(0, 0);
    motors(false);
    motorTest = false;
    server.send(204);
  });
  server.begin();
}

void setup() {
  Serial.begin(115200);
  pinMode(STEP_L, OUTPUT);
  pinMode(DIR_L, OUTPUT);
  pinMode(STEP_R, OUTPUT);
  pinMode(DIR_R, OUTPUT);
  pinMode(EN_PIN, OUTPUT);
  motors(false);
  digitalWrite(DIR_L, HIGH);
  digitalWrite(DIR_R, LOW);
  analogReadResolution(12);

  prefs.begin("balancer");
  const char *names[4] = {"k1", "k2", "k3", "k4"};
  for (int i = 0; i < 4; i++) K[i] = prefs.getFloat(names[i], K[i]);
  trim = prefs.getFloat("trim", 0);
  gyroBias = prefs.getFloat("gbias", 0);

  Wire.begin(SDA_PIN, SCL_PIN, 400000);
  if (!mpuBegin()) Serial.println("MPU-6050 not found: check SDA (GPIO6), SCL (GPIO7) and 3.3 V.");

  uint8_t mac[6];
  WiFi.macAddress(mac);
  char ssid[32];
  snprintf(ssid, sizeof ssid, "Voltomni-Balancer-%02X%02X", mac[4], mac[5]);
  WiFi.softAP(ssid, AP_PASSWORD);
  setupWeb();
  Serial.printf("Wi-Fi \"%s\", password \"%s\": open http://%s\n", ssid, AP_PASSWORD, WiFi.softAPIP().toString().c_str());

  timer = timerBegin(1000000);
  timerAttachInterrupt(timer, &onTick);
  timerAlarm(timer, 1000000 / TICK_HZ, true, 0);
  xTaskCreate(control, "control", 4096, nullptr, configMAX_PRIORITIES - 1, nullptr);
  Serial.println("Stand the robot upright and hold it still to start balancing.");
}

void loop() {
  server.handleClient();
  static uint32_t lastBat = 0;
  if (millis() - lastBat > 500) {
    lastBat = millis();
    vbat = analogReadMilliVolts(VBAT_PIN) / 1000.0f * VBAT_DIVIDER;
    if (vbat > 5 && vbat < VBAT_LOW) {  // above 5 V: a battery is connected (not just USB)
      if (!lowBattery) Serial.println("Battery low: stopped. Charge it.");
      lowBattery = true;
    }
  }
  delay(2);
}