/*
 * ============================================================
 * EEE532 AIoT Project - Smart Walker Firmware
 * Intelligent Mobility and Fall Prevention Robotic System
 * ============================================================
 * Platform : ESP32-DevKitC (Wokwi Simulation)
 * Sensors  : MPU6050 (IMU), Potentiometer (Pressure proxy)
 * Actuators: Servo (Motor), LED (Brake), RGB LEDs (Status)
 * Cloud    : ThingsBoard via MQTT
 * ============================================================
 */

#include <Wire.h>
#include <WiFi.h>
#include <PubSubClient.h>
#include <ESP32Servo.h>

// ===================== PIN DEFINITIONS =====================
#define IMU_SDA          21
#define IMU_SCL          22
#define PRESSURE_PIN     34   // Potentiometer (FSR proxy)
#define MOTOR_SERVO_PIN  25   // Servo simulates DC motor speed
#define BRAKE_LED_PIN    26   // LED simulates electromagnetic brake
#define LED_GREEN_PIN     2   // Normal walking indicator
#define LED_YELLOW_PIN    4   // Motor assist active indicator
#define LED_RED_PIN       5   // Fall detected indicator

// ===================== MPU6050 REGISTERS =====================
#define MPU6050_ADDR      0x68
#define MPU6050_PWR_MGMT  0x6B
#define MPU6050_ACCEL_X   0x3B
#define MPU6050_GYRO_X    0x43
#define MPU6050_CONFIG    0x1A
#define MPU6050_GYRO_CFG  0x1B
#define MPU6050_ACCEL_CFG 0x1C

// ===================== WIFI & MQTT CONFIG =====================
// Wokwi simulated WiFi
const char* WIFI_SSID     = "Wokwi-GUEST";
const char* WIFI_PASSWORD = "";

// ThingsBoard Cloud MQTT
const char* MQTT_SERVER   = "thingsboard.cloud";
const int   MQTT_PORT     = 1883;
// >>> REPLACE WITH YOUR DEVICE ACCESS TOKEN <<<
const char* MQTT_TOKEN    = "X4B9DpqbVIfUJPE5nLWb";

// ===================== FALL DETECTION THRESHOLDS =====================
#define ACCEL_FALL_THRESHOLD   2.5f   // g-force threshold for fall
#define TILT_FALL_THRESHOLD    45.0f  // degrees threshold for fall
#define FALL_WINDOW_MS         500    // ms time window for fall detection
#define FALL_COOLDOWN_MS       5000   // 5s cooldown after fall event

// ===================== PID PARAMETERS =====================
float Kp = 2.0f;
float Ki = 0.5f;
float Kd = 0.1f;

// ===================== SYSTEM PARAMETERS =====================
#define PRESSURE_ASSIST_THRESHOLD  2000  // ADC threshold for motor assist
#define MQTT_UPLOAD_INTERVAL_MS    2000  // Upload telemetry every 2s
#define IMU_SAMPLE_INTERVAL_MS     10    // 100Hz sampling rate
#define SERIAL_BAUD                115200

// ===================== GLOBAL VARIABLES =====================
// IMU raw data
float accel_x = 0, accel_y = 0, accel_z = 0;
float gyro_x = 0, gyro_y = 0, gyro_z = 0;

// Processed data
float accel_magnitude = 0;
float tilt_angle = 0;
float complementary_angle = 0;
float alpha_complementary = 0.98f; // Complementary filter coefficient

// Pressure
int pressure_value = 0;

// Fall detection
bool fall_detected = false;
unsigned long fall_start_time = 0;
unsigned long last_fall_time = 0;
bool in_fall_cooldown = false;

// PID control
float desired_speed = 0;
float actual_speed = 0;
float pid_output = 0;
float pid_integral = 0;
float pid_prev_error = 0;
unsigned long pid_prev_time = 0;

// Motor
int motor_pwm = 0;
bool brake_engaged = false;

// MQTT timing
unsigned long last_mqtt_upload = 0;

// Servo
Servo motorServo;

// WiFi & MQTT clients
WiFiClient espClient;
PubSubClient mqtt(espClient);

// Moving average buffer for pressure
#define MA_WINDOW 5
int pressure_buffer[MA_WINDOW] = {0};
int pressure_index = 0;

// ===================== MPU6050 FUNCTIONS =====================

void mpu6050_init() {
  Wire.begin(IMU_SDA, IMU_SCL);
  delay(100);

  // Wake up MPU6050
  Wire.beginTransmission(MPU6050_ADDR);
  Wire.write(MPU6050_PWR_MGMT);
  Wire.write(0x00);  // Clear sleep mode
  Wire.endTransmission(true);
  delay(50);

  // Configure accelerometer: ±4g (AFS_SEL=1)
  Wire.beginTransmission(MPU6050_ADDR);
  Wire.write(MPU6050_ACCEL_CFG);
  Wire.write(0x08);  // AFS_SEL=1 => ±4g
  Wire.endTransmission(true);

  // Configure gyroscope: ±500°/s (FS_SEL=1)
  Wire.beginTransmission(MPU6050_ADDR);
  Wire.write(MPU6050_GYRO_CFG);
  Wire.write(0x08);  // FS_SEL=1 => ±500°/s
  Wire.endTransmission(true);

  // Configure DLPF (Digital Low Pass Filter)
  Wire.beginTransmission(MPU6050_ADDR);
  Wire.write(MPU6050_CONFIG);
  Wire.write(0x03);  // ACC: 44Hz, GYRO: 42Hz
  Wire.endTransmission(true);

  Serial.println("[MPU6050] Initialized successfully");
}

int16_t readRegister16(uint8_t reg) {
  Wire.beginTransmission(MPU6050_ADDR);
  Wire.write(reg);
  Wire.endTransmission(false);
  Wire.requestFrom(MPU6050_ADDR, (uint8_t)2, true);
  int16_t value = (Wire.read() << 8) | Wire.read();
  return value;
}

void read_mpu6050() {
  // Read accelerometer (6 registers starting from 0x3B)
  int16_t raw_ax = readRegister16(MPU6050_ACCEL_X);
  int16_t raw_ay = readRegister16(MPU6050_ACCEL_X + 2);
  int16_t raw_az = readRegister16(MPU6050_ACCEL_X + 4);

  // Read gyroscope (6 registers starting from 0x43)
  int16_t raw_gx = readRegister16(MPU6050_GYRO_X);
  int16_t raw_gy = readRegister16(MPU6050_GYRO_X + 2);
  int16_t raw_gz = readRegister16(MPU6050_GYRO_X + 4);

  // Convert to physical units
  // ACC: ±4g range => sensitivity = 8192 LSB/g
  // GYRO: ±500°/s => sensitivity = 65.5 LSB/(°/s)
  accel_x = (float)raw_ax / 8192.0f;
  accel_y = (float)raw_ay / 8192.0f;
  accel_z = (float)raw_az / 8192.0f;
  gyro_x  = (float)raw_gx / 65.5f;
  gyro_y  = (float)raw_gy / 65.5f;
  gyro_z  = (float)raw_gz / 65.5f;

  // Calculate total acceleration magnitude
  accel_magnitude = sqrt(accel_x * accel_x + accel_y * accel_y + accel_z * accel_z);
}

// ===================== DATA PROCESSING =====================

void process_sensor_data() {
  // 1. Complementary filter for tilt angle estimation
  //    angle = alpha * (angle + gyro * dt) + (1 - alpha) * accel_angle
  float dt = IMU_SAMPLE_INTERVAL_MS / 1000.0f;

  // Calculate tilt from accelerometer (angle from vertical)
  float accel_angle = atan2(accel_x, sqrt(accel_y * accel_y + accel_z * accel_z)) * 180.0f / PI;

  // Complementary filter
  complementary_angle = alpha_complementary * (complementary_angle + gyro_x * dt)
                      + (1.0f - alpha_complementary) * accel_angle;

  tilt_angle = complementary_angle;

  // 2. Moving average filter for pressure sensor
  pressure_buffer[pressure_index] = analogRead(PRESSURE_PIN);
  pressure_index = (pressure_index + 1) % MA_WINDOW;
  long sum = 0;
  for (int i = 0; i < MA_WINDOW; i++) {
    sum += pressure_buffer[i];
  }
  pressure_value = sum / MA_WINDOW;
}

// ===================== FALL DETECTION =====================

void detect_fall() {
  unsigned long current_time = millis();

  // Check cooldown period
  if (in_fall_cooldown) {
    if (current_time - last_fall_time > FALL_COOLDOWN_MS) {
      in_fall_cooldown = false;
    } else {
      return;  // Still in cooldown
    }
  }

  // Fall detection: acceleration exceeds threshold AND tilt exceeds threshold
  bool accel_trigger = (accel_magnitude > ACCEL_FALL_THRESHOLD);
  bool tilt_trigger = (fabs(tilt_angle) > TILT_FALL_THRESHOLD);

  if (accel_trigger && tilt_trigger) {
    if (fall_start_time == 0) {
      fall_start_time = current_time;
    }
    // Check if both conditions persist within the time window
    if (current_time - fall_start_time < FALL_WINDOW_MS) {
      // FALL DETECTED!
      fall_detected = true;
      last_fall_time = current_time;
      in_fall_cooldown = true;
      fall_start_time = 0;

      // Activate emergency response
      engage_brake();
      digitalWrite(LED_RED_PIN, HIGH);
      digitalWrite(LED_GREEN_PIN, LOW);
      digitalWrite(LED_YELLOW_PIN, LOW);

      Serial.println("╔══════════════════════════════════╗");
      Serial.println("║    ⚠️  FALL DETECTED!  ⚠️         ║");
      Serial.println("╠══════════════════════════════════╣");
      Serial.print("║ Accel: "); Serial.print(accel_magnitude, 2); Serial.println(" g");
      Serial.print("║ Tilt:  "); Serial.print(tilt_angle, 1); Serial.println(" °");
      Serial.println("║ Brake: ENGAGED");
      Serial.println("║ Alert: SENDING via MQTT...");
      Serial.println("╚══════════════════════════════════╝");

      // Immediately send fall alert via MQTT
      send_fall_alert();
    }
  } else {
    // Reset fall timer if conditions no longer met
    if (current_time - fall_start_time > FALL_WINDOW_MS) {
      fall_start_time = 0;
    }

    // If not in cooldown and no fall, clear detection
    if (!in_fall_cooldown) {
      fall_detected = false;
      digitalWrite(LED_RED_PIN, LOW);
      if (!brake_engaged) {
        // Restore appropriate LED state
        if (motor_pwm > 0) {
          digitalWrite(LED_YELLOW_PIN, HIGH);
          digitalWrite(LED_GREEN_PIN, LOW);
        } else {
          digitalWrite(LED_GREEN_PIN, HIGH);
          digitalWrite(LED_YELLOW_PIN, LOW);
        }
      }
    }
  }
}

// ===================== BRAKE CONTROL =====================

void engage_brake() {
  brake_engaged = true;
  digitalWrite(BRAKE_LED_PIN, HIGH);
  motor_pwm = 0;
  motorServo.write(0);  // Stop motor
}

void release_brake() {
  brake_engaged = false;
  digitalWrite(BRAKE_LED_PIN, LOW);
  digitalWrite(LED_RED_PIN, LOW);
  digitalWrite(LED_GREEN_PIN, HIGH);
}

// ===================== PID MOTOR CONTROL =====================

void pid_motor_control() {
  if (fall_detected || brake_engaged) {
    // Safety: do not run motor during fall/brake
    motor_pwm = 0;
    motorServo.write(0);
    pid_integral = 0;
    pid_prev_error = 0;
    return;
  }

  unsigned long current_time = millis();
  float dt = (current_time - pid_prev_time) / 1000.0f;
  if (dt <= 0 || dt > 1.0f) {
    pid_prev_time = current_time;
    return;
  }

  // Determine desired speed from pressure sensor
  // Map pressure (0-4095) to desired speed (0-180 for servo)
  if (pressure_value > PRESSURE_ASSIST_THRESHOLD) {
    desired_speed = map(pressure_value, PRESSURE_ASSIST_THRESHOLD, 4095, 30, 180);
    digitalWrite(LED_YELLOW_PIN, HIGH);
    digitalWrite(LED_GREEN_PIN, LOW);
  } else {
    desired_speed = 0;
    digitalWrite(LED_YELLOW_PIN, LOW);
    digitalWrite(LED_GREEN_PIN, HIGH);
  }

  // Simulate actual speed (in real system, this comes from wheel encoder)
  // For Wokwi simulation, we use a simple first-order model
  float motor_time_constant = 0.3f;  // seconds
  actual_speed += (motor_pwm - actual_speed) * (dt / motor_time_constant);

  // PID calculation
  float error = desired_speed - actual_speed;

  // Proportional
  float p_term = Kp * error;

  // Integral (with anti-windup)
  pid_integral += error * dt;
  pid_integral = constrain(pid_integral, -100, 100);  // Anti-windup
  float i_term = Ki * pid_integral;

  // Derivative
  float d_term = Kd * (error - pid_prev_error) / dt;
  pid_prev_error = error;

  // Total PID output
  pid_output = p_term + i_term + d_term;

  // Apply to motor PWM
  motor_pwm = constrain((int)(actual_speed + pid_output), 0, 180);

  // Set servo position (simulating motor speed)
  motorServo.write(motor_pwm);

  pid_prev_time = current_time;
}

// ===================== WIFI & MQTT =====================

void setup_wifi() {
  Serial.print("[WiFi] Connecting to ");
  Serial.println(WIFI_SSID);
  WiFi.begin(WIFI_SSID, WIFI_PASSWORD);

  int attempts = 0;
  while (WiFi.status() != WL_CONNECTED && attempts < 40) {
    delay(500);
    Serial.print(".");
    attempts++;
  }

  if (WiFi.status() == WL_CONNECTED) {
    Serial.println();
    Serial.print("[WiFi] Connected! IP: ");
    Serial.println(WiFi.localIP());
  } else {
    Serial.println();
    Serial.println("[WiFi] Connection failed - will retry in loop");
  }
}

void mqtt_callback(char* topic, byte* payload, unsigned int length) {
  // Handle RPC commands from ThingsBoard
  String message;
  for (unsigned int i = 0; i < length; i++) {
    message += (char)payload[i];
  }
  Serial.print("[MQTT] RPC received: ");
  Serial.println(message);

  // Parse simple JSON commands
  if (message.indexOf("\"setBrake\"") >= 0 || message.indexOf("\"brake\"") >= 0) {
    if (message.indexOf("true") >= 0) {
      engage_brake();
      Serial.println("[RPC] Brake ENGAGED remotely");
    } else if (message.indexOf("false") >= 0) {
      release_brake();
      Serial.println("[RPC] Brake RELEASED remotely");
    }
  }

  if (message.indexOf("\"setPID\"") >= 0 || message.indexOf("\"kp\"") >= 0) {
    // Simple PID parameter update parsing
    int kp_idx = message.indexOf("\"kp\"");
    if (kp_idx >= 0) {
      float new_kp = message.substring(kp_idx + 5).toFloat();
      if (new_kp > 0 && new_kp < 100) Kp = new_kp;
    }
    int ki_idx = message.indexOf("\"ki\"");
    if (ki_idx >= 0) {
      float new_ki = message.substring(ki_idx + 5).toFloat();
      if (new_ki >= 0 && new_ki < 100) Ki = new_ki;
    }
    int kd_idx = message.indexOf("\"kd\"");
    if (kd_idx >= 0) {
      float new_kd = message.substring(kd_idx + 5).toFloat();
      if (new_kd >= 0 && new_kd < 100) Kd = new_kd;
    }
    Serial.printf("[RPC] PID updated: Kp=%.2f, Ki=%.2f, Kd=%.2f\n", Kp, Ki, Kd);
  }
}

void mqtt_reconnect() {
  while (!mqtt.connected() && WiFi.status() == WL_CONNECTED) {
    Serial.print("[MQTT] Connecting to ThingsBoard...");
    if (mqtt.connect("SmartWalker-ESP32", MQTT_TOKEN, NULL)) {
      Serial.println(" Connected!");
      // Subscribe to RPC topic
      mqtt.subscribe("v1/devices/me/rpc/request/+");
      Serial.println("[MQTT] Subscribed to RPC");
    } else {
      Serial.print(" Failed, rc=");
      Serial.print(mqtt.state());
      Serial.println(" - retry in 5s");
      delay(5000);
    }
  }
}

void send_telemetry() {
  if (!mqtt.connected()) return;

  // Build JSON telemetry payload
  String payload = "{";
  payload += "\"accel_x\":" + String(accel_x, 3) + ",";
  payload += "\"accel_y\":" + String(accel_y, 3) + ",";
  payload += "\"accel_z\":" + String(accel_z, 3) + ",";
  payload += "\"accel_mag\":" + String(accel_magnitude, 3) + ",";
  payload += "\"gyro_x\":" + String(gyro_x, 2) + ",";
  payload += "\"gyro_y\":" + String(gyro_y, 2) + ",";
  payload += "\"gyro_z\":" + String(gyro_z, 2) + ",";
  payload += "\"tilt_angle\":" + String(tilt_angle, 2) + ",";
  payload += "\"pressure\":" + String(pressure_value) + ",";
  payload += "\"motor_speed\":" + String(motor_pwm) + ",";
  payload += "\"fall_detected\":" + String(fall_detected ? "true" : "false") + ",";
  payload += "\"brake_engaged\":" + String(brake_engaged ? "true" : "false") + ",";
  payload += "\"desired_speed\":" + String(desired_speed, 1) + ",";
  payload += "\"actual_speed\":" + String(actual_speed, 1);
  payload += "}";

  mqtt.publish("v1/devices/me/telemetry", payload.c_str());
  Serial.print("[MQTT] Telemetry sent: ");
  Serial.println(payload.substring(0, 80) + "...");
}

void send_fall_alert() {
  if (!mqtt.connected()) return;

  // Priority fall alert (also sent as attribute for alarm trigger)
  String alert = "{";
  alert += "\"fall_detected\":true,";
  alert += "\"fall_accel\":" + String(accel_magnitude, 3) + ",";
  alert += "\"fall_tilt\":" + String(tilt_angle, 2) + ",";
  alert += "\"fall_timestamp\":" + String(millis());
  alert += "}";

  // Send as telemetry (triggers ThingsBoard alarm rule)
  mqtt.publish("v1/devices/me/telemetry", alert.c_str());

  // Also send as attribute
  String attr = "{\"lastFallAccel\":" + String(accel_magnitude, 3) + ",\"lastFallTilt\":" + String(tilt_angle, 2) + "}";
  mqtt.publish("v1/devices/me/attributes", attr.c_str());

  Serial.println("[MQTT] Fall alert sent to ThingsBoard!");
}

// ===================== SERIAL DEBUG OUTPUT =====================

void print_debug_info() {
  Serial.println("─────────────────────────────────────────");
  Serial.printf("  Accel  (g):  X=%+.3f  Y=%+.3f  Z=%+.3f  |Mag|=%.3f\n",
                accel_x, accel_y, accel_z, accel_magnitude);
  Serial.printf("  Gyro (°/s):  X=%+.1f  Y=%+.1f  Z=%+.1f\n",
                gyro_x, gyro_y, gyro_z);
  Serial.printf("  Tilt Angle:  %.1f°\n", tilt_angle);
  Serial.printf("  Pressure:    %d (threshold: %d)\n", pressure_value, PRESSURE_ASSIST_THRESHOLD);
  Serial.printf("  Motor PWM:   %d / 180  | Desired: %.1f | Actual: %.1f\n",
                motor_pwm, desired_speed, actual_speed);
  Serial.printf("  Fall: %s | Brake: %s | PID Out: %.2f\n",
                fall_detected ? "⚠️ DETECTED" : "OK",
                brake_engaged ? "ENGAGED" : "OFF",
                pid_output);
  Serial.printf("  WiFi: %s | MQTT: %s\n",
                WiFi.status() == WL_CONNECTED ? "Connected" : "Disconnected",
                mqtt.connected() ? "Connected" : "Disconnected");
  Serial.println("─────────────────────────────────────────");
}

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

void setup() {
  // Initialize Serial
  Serial.begin(SERIAL_BAUD);
  Serial.println("\n\n╔══════════════════════════════════════╗");
  Serial.println("║  Smart Walker - AIoT System v1.0     ║");
  Serial.println("║  EEE532 Artificial Intelligence of    ║");
  Serial.println("║  Things Technologies                  ║");
  Serial.println("╚══════════════════════════════════════╝\n");

  // Initialize GPIO
  pinMode(LED_GREEN_PIN, OUTPUT);
  pinMode(LED_YELLOW_PIN, OUTPUT);
  pinMode(LED_RED_PIN, OUTPUT);
  pinMode(BRAKE_LED_PIN, OUTPUT);
  pinMode(PRESSURE_PIN, INPUT);

  // Initial LED state
  digitalWrite(LED_GREEN_PIN, HIGH);
  digitalWrite(LED_YELLOW_PIN, LOW);
  digitalWrite(LED_RED_PIN, LOW);
  digitalWrite(BRAKE_LED_PIN, LOW);

  // Initialize MPU6050
  mpu6050_init();

  // Initialize Servo (motor)
  motorServo.attach(MOTOR_SERVO_PIN);
  motorServo.write(0);  // Motor off

  // Initialize pressure buffer
  for (int i = 0; i < MA_WINDOW; i++) {
    pressure_buffer[i] = analogRead(PRESSURE_PIN);
  }

  // Initialize WiFi
  setup_wifi();

  // Initialize MQTT
  mqtt.setServer(MQTT_SERVER, MQTT_PORT);
  mqtt.setCallback(mqtt_callback);

  // Initialize PID timer
  pid_prev_time = millis();

  Serial.println("[System] Initialization complete!");
  Serial.println("[System] Starting main loop...\n");
}

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

unsigned long last_imu_read = 0;
unsigned long last_debug_print = 0;

void loop() {
  unsigned long current_time = millis();

  // 1. Read IMU at 100Hz (every 10ms)
  if (current_time - last_imu_read >= IMU_SAMPLE_INTERVAL_MS) {
    last_imu_read = current_time;
    read_mpu6050();
    process_sensor_data();
    detect_fall();
  }

  // 2. PID motor control at 50Hz (every 20ms)
  if (current_time - pid_prev_time >= 20) {
    pid_motor_control();
  }

  // 3. MQTT telemetry upload every 2 seconds
  if (current_time - last_mqtt_upload >= MQTT_UPLOAD_INTERVAL_MS) {
    last_mqtt_upload = current_time;

    // Maintain MQTT connection
    if (WiFi.status() == WL_CONNECTED) {
      if (!mqtt.connected()) {
        mqtt_reconnect();
      }
      mqtt.loop();
      send_telemetry();
    }
  }

  // 4. Debug print every 1 second
  if (current_time - last_debug_print >= 1000) {
    last_debug_print = current_time;
    print_debug_info();
  }

  // 5. Auto-release brake after cooldown (for demo purposes)
  //    In real system, brake is released by caregiver via RPC
  if (brake_engaged && !fall_detected && current_time - last_fall_time > FALL_COOLDOWN_MS) {
    // Auto-release for demo convenience
    release_brake();
    Serial.println("[System] Brake auto-released after cooldown");
  }

  // 6. MQTT loop (keep alive)
  if (mqtt.connected()) {
    mqtt.loop();
  }

  delay(1);  // Small delay to prevent watchdog
}