/*
 * SCHTICK: Reaction Wheel Inverted Pendulum
 * Firmware Skeleton - ESP32 Motor & IMU Control
 * 
 * Hardware:
 * - ESP32 WROOM-32 @ 240MHz dual-core
 * - TB6612FNG H-bridge motor driver (module)
 * - N20 gearmotor 6V 500RPM with Hall encoder
 * - MPU-6050 IMU (GY-521 breakout) on I2C
 * 6V wall adapter powers motor; ESP32 via USB
 * 
 * Author: SCHTICK Project
 * Date: 2025
 */

#include <Wire.h>
#include <Adafruit_MPU6050.h>
#include <Adafruit_Sensor.h>

// ============================================================================
// PIN DEFINITIONS
// ============================================================================

// Motor control (TB6612FNG H-bridge)
#define PIN_PWMA      4   // PWM speed control (0-255)
#define PIN_AIN1      12  // Direction control bit 1
#define PIN_AIN2      13  // Direction control bit 2
#define PIN_STBY      5   // Standby/enable (active high)

// Hall encoder (motor speed feedback)
#define PIN_HALL      35  // GPIO 35 (input from Hall sensor with pull-up)

// I2C bus (MPU-6050)
#define PIN_SDA       21  // I2C data
#define PIN_SCL       22  // I2C clock
#define I2C_FREQ      100000  // 100 kHz (safe in noisy environment)

// Serial telemetry
#define SERIAL_BAUD   115200

// ============================================================================
// CONTROL PARAMETERS
// ============================================================================

#define UPDATE_RATE_HZ    100     // Control loop frequency
#define UPDATE_PERIOD_MS  (1000 / UPDATE_RATE_HZ)

#define PWM_MIN           0       // Motor off
#define PWM_MAX           255     // Motor full speed
#define PWM_RAMP_RATE     10      // PWM change per 10ms (soft start)
#define PWM_STARTUP_DELAY 100     // ms to wait before ramping motor

// Hall encoder pulse counting
#define HALL_DEBOUNCE_MS  2       // Ignore pulses within 2ms

// ============================================================================
// GLOBAL STATE
// ============================================================================

// IMU sensor object
Adafruit_MPU6050 mpu;

// Motor control state
volatile uint32_t motor_pwm_command = 0;
volatile uint32_t motor_pwm_actual = 0;
volatile int motor_direction = 0;  // -1 = reverse, 0 = stop, 1 = forward

// Hall encoder pulse tracking
volatile uint32_t hall_pulse_count = 0;
volatile uint32_t hall_last_pulse_time = 0;
volatile uint8_t hall_pulse_per_rev = 1;  // Adjust if your motor has multiple poles

// IMU readings
float angle_rad = 0.0;      // Inverted pendulum angle (rad, from accelerometer)
float angle_rate_rad_s = 0.0;  // Angle rate (rad/s, from gyroscope)
float gyro_x, gyro_y, gyro_z;
float accel_x, accel_y, accel_z;

// Complementary filter state
float angle_from_gyro = 0.0;
float angle_from_accel = 0.0;
float complementary_alpha = 0.95;  // Weight of gyro vs accel (0.95 = 95% gyro, 5% accel correction)

// Telemetry timing
uint32_t last_telemetry_time = 0;
uint32_t telemetry_interval_ms = 100;  // Send telemetry every 100ms

// Safety watchdog
uint32_t last_control_loop_time = 0;
uint32_t control_timeout_ms = 500;  // If control loop stalls > 500ms, kill motor

// ============================================================================
// INTERRUPT HANDLERS
// ============================================================================

// Hall encoder ISR: Count pulses, debounce to avoid double-counting
void IRAM_ATTR hall_interrupt() {
  uint32_t now = millis();
  if (now - hall_last_pulse_time > HALL_DEBOUNCE_MS) {
    hall_pulse_count++;
    hall_last_pulse_time = now;
  }
}

// ============================================================================
// SETUP
// ============================================================================

void setup() {
  // Initialize serial for telemetry
  Serial.begin(SERIAL_BAUD);
  delay(500);
  Serial.println("\n\n=== SCHTICK Firmware Boot ===\n");

  // Initialize I2C for MPU-6050
  Wire.begin(PIN_SDA, PIN_SCL, I2C_FREQ);
  
  // Initialize MPU-6050
  if (!mpu.begin(0x68, &Wire, 0)) {
    Serial.println("FATAL: MPU-6050 not found on I2C bus!");
    Serial.println("Check wiring: SDA=GPIO21, SCL=GPIO22, VCC=3.3V, GND");
    while (1) {
      delay(1000);
      Serial.println("Waiting for MPU-6050...");
    }
  }
  Serial.println("✓ MPU-6050 initialized");
  
  // Configure MPU-6050 ranges
  mpu.setAccelerometerRange(MPU6050_RANGE_8_G);
  mpu.setGyroRange(MPU6050_RANGE_500_DEG);
  mpu.setFilterBandwidth(MPU6050_BAND_21_HZ);  // Low-pass filter for noise
  
  // Initialize motor control pins
  pinMode(PIN_PWMA, OUTPUT);
  pinMode(PIN_AIN1, OUTPUT);
  pinMode(PIN_AIN2, OUTPUT);
  pinMode(PIN_STBY, OUTPUT);
  pinMode(PIN_HALL, INPUT_PULLUP);
  
  // Set motor to standby (disabled)
  digitalWrite(PIN_STBY, LOW);
  analogWrite(PIN_PWMA, 0);
  digitalWrite(PIN_AIN1, LOW);
  digitalWrite(PIN_AIN2, LOW);
  Serial.println("✓ Motor control pins initialized (standby mode)");
  
  // Attach Hall encoder interrupt
  attachInterrupt(digitalPinToInterrupt(PIN_HALL), hall_interrupt, RISING);
  Serial.println("✓ Hall encoder ISR attached (GPIO 35)");
  
  Serial.println("\n=== Entering main loop ===\n");
  delay(PWM_STARTUP_DELAY);
}

// ============================================================================
// MAIN CONTROL LOOP
// ============================================================================

void loop() {
  uint32_t loop_start = millis();
  
  // Read IMU (6-axis: accel + gyro)
  sensors_event_t a, g, temp;
  mpu.getEvent(&a, &g, &temp);
  
  accel_x = a.acceleration.x;
  accel_y = a.acceleration.y;
  accel_z = a.acceleration.z;
  
  gyro_x = g.gyro.x;
  gyro_y = g.gyro.y;
  gyro_z = g.gyro.z;
  
  // Estimate pendulum angle using complementary filter
  // Accelerometer gives direct angle (theta = atan2(ax, az))
  // Gyro integrates for continuous rate
  angle_from_accel = atan2(accel_x, accel_z);  // rad, [-pi/2, pi/2]
  angle_rate_rad_s = gyro_y;  // Rotation around Y-axis (pitch)
  
  // Integrate gyro (simple forward Euler for now)
  float dt = UPDATE_PERIOD_MS / 1000.0;
  angle_from_gyro += angle_rate_rad_s * dt;
  
  // Complementary filter: fuse accel + gyro
  angle_rad = complementary_alpha * angle_from_gyro + 
              (1.0 - complementary_alpha) * angle_from_accel;
  
  // ========================================================================
  // PLACEHOLDER: YOUR CONTROL LAW HERE
  // ========================================================================
  // This is where PID, LQR, or RL policy will compute motor command
  // For now: simple proportional control (P-controller)
  // 
  // Control goal: Keep pendulum upright (angle_rad ≈ 0)
  // Command: motor torque proportional to angle error
  
  float kp = 50.0;  // Proportional gain (tune this)
  float angle_error = angle_rad - 0.0;  // Error from vertical
  float motor_voltage_command = kp * angle_error;  // Volts
  
  // Convert voltage command to PWM (assume 6V max = 255 PWM)
  int pwm_target = (int)(motor_voltage_command * 255.0 / 6.0);
  pwm_target = constrain(pwm_target, -255, 255);
  
  // Determine direction from sign
  if (pwm_target > 0) {
    motor_direction = 1;
    motor_pwm_command = pwm_target;
  } else if (pwm_target < 0) {
    motor_direction = -1;
    motor_pwm_command = -pwm_target;
  } else {
    motor_direction = 0;
    motor_pwm_command = 0;
  }
  
  // ========================================================================
  // SOFT-START RAMP (prevent inrush current)
  // ========================================================================
  if (motor_pwm_actual < motor_pwm_command) {
    motor_pwm_actual = min(motor_pwm_actual + PWM_RAMP_RATE, motor_pwm_command);
  } else if (motor_pwm_actual > motor_pwm_command) {
    motor_pwm_actual = max((int)motor_pwm_actual - PWM_RAMP_RATE, (int)motor_pwm_command);
  }
  
  // Apply motor command to H-bridge
  set_motor_pwm(motor_direction, motor_pwm_actual);
  
  // ========================================================================
  // SAFETY: WATCHDOG TIMEOUT
  // ========================================================================
  // If control loop stalls (should never happen, but safety first)
  uint32_t loop_end = millis();
  if (loop_end - last_control_loop_time > control_timeout_ms) {
    Serial.println("WARN: Control loop timeout! Killing motor.");
    set_motor_pwm(0, 0);  // Force motor off
  }
  last_control_loop_time = loop_end;
  
  // ========================================================================
  // TELEMETRY (serial output for logging)
  // ========================================================================
  if (millis() - last_telemetry_time > telemetry_interval_ms) {
    // CSV format: timestamp_ms, angle_rad, angle_rate_rad_s, hall_count, pwm_command, direction
    Serial.print(millis());
    Serial.print(",");
    Serial.print(angle_rad, 4);
    Serial.print(",");
    Serial.print(angle_rate_rad_s, 4);
    Serial.print(",");
    Serial.print(hall_pulse_count);
    Serial.print(",");
    Serial.print(motor_pwm_command);
    Serial.print(",");
    Serial.println(motor_direction);
    
    last_telemetry_time = millis();
  }
  
  // ========================================================================
  // LOOP TIMING: Maintain constant 100 Hz update rate
  // ========================================================================
  uint32_t loop_elapsed = millis() - loop_start;
  if (loop_elapsed < UPDATE_PERIOD_MS) {
    delay(UPDATE_PERIOD_MS - loop_elapsed);
  }
}

// ============================================================================
// MOTOR CONTROL HELPER
// ============================================================================

void set_motor_pwm(int direction, uint32_t pwm_value) {
  /*
   * Direction: -1 = reverse, 0 = stop, 1 = forward
   * PWM value: 0-255
   * TB6612FNG control logic:
   *   STBY=1 (enabled)
   *   Forward:  AIN1=1, AIN2=0, PWMA=value
   *   Reverse:  AIN1=0, AIN2=1, PWMA=value
   *   Stop:     AIN1=0, AIN2=0, PWMA=0 (coasting)
   */
  
  pwm_value = constrain(pwm_value, 0, 255);
  
  // Enable motor driver
  digitalWrite(PIN_STBY, HIGH);
  
  if (direction == 1) {
    // Forward
    digitalWrite(PIN_AIN1, HIGH);
    digitalWrite(PIN_AIN2, LOW);
    analogWrite(PIN_PWMA, pwm_value);
  } else if (direction == -1) {
    // Reverse
    digitalWrite(PIN_AIN1, LOW);
    digitalWrite(PIN_AIN2, HIGH);
    analogWrite(PIN_PWMA, pwm_value);
  } else {
    // Stop (coast)
    digitalWrite(PIN_AIN1, LOW);
    digitalWrite(PIN_AIN2, LOW);
    analogWrite(PIN_PWMA, 0);
  }
}

// ============================================================================
// DEBUG COMMANDS (future: serial input for manual testing)
// ============================================================================
// Can extend this to accept commands like:
// "M<pwm>,<dir>" — manual motor control
// "C" — calibrate accelerometer zero-point
// "S" — status dump (angle, rates, hall count, etc.)
// 
// For now, telemetry-only. Firmware ready for extension.