/*
 * ============================================================
 *  EEE532 AIoT Project
 *  Intelligent Mobility and Fall Prevention Robotic System
 *  ESP32 Firmware — Wokwi Simulation Version
 * ============================================================
 *  Hardware: ESP32-DevKit-V1 + MPU6050 + Potentiometer(FSR proxy)
 *            + LEDs + Servo(motor proxy) + Brake LED
 *  Cloud:    ThingsBoard Cloud (MQTT)
 * ============================================================
 */

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

// ============ WiFi & MQTT Configuration ============
const char* ssid        = "Wokwi-GUEST";
const char* password    = "";
const char* mqtt_server = "thingsboard.cloud";
const int   mqtt_port   = 1883;
const char* token       = "X4B9DpqbVIfUJPE5nLWb";  // <-- Replace with your ThingsBoard device token

// ============ Pin Definitions ============
#define IMU_SDA        21
#define IMU_SCL        22
#define PRESSURE_PIN   34   // Potentiometer simulates FSR-402
#define MOTOR_PIN      18   // Servo simulates DC motor speed
#define BRAKE_PIN      26   // LED simulates electromagnetic brake
#define LED_GREEN      2    // Normal walking indicator
#define LED_YELLOW     4    // Motor assist active indicator
#define LED_RED        5    // Fall detected alarm indicator

// ============ MPU6050 Object ============
Adafruit_MPU6050 mpu;

// ============ Servo Object (Motor Proxy) ============
Servo motorServo;

// ============ WiFi & MQTT Clients ============
WiFiClient espClient;
PubSubClient client(espClient);

// ============ Sensor Data ============
float accel_x = 0, accel_y = 0, accel_z = 0;
float gyro_x = 0, gyro_y = 0, gyro_z = 0;
float accel_mag = 0;
float tilt_angle = 0;
int   pressure_val = 0;
int   motor_speed = 0;

// ============ Fall Detection Parameters ============
const float FALL_ACCEL_THRESHOLD = 2.5 * 9.81;  // 2.5g in m/s^2
const float FALL_TILT_THRESHOLD  = 45.0;         // degrees
const unsigned long FALL_WINDOW  = 500;           // ms
const unsigned long FALL_COOLDOWN = 5000;         // 5s cooldown after fall

bool  fall_detected  = false;
bool  brake_engaged  = false;
unsigned long fall_start_time   = 0;
unsigned long last_fall_time    = 0;
bool  in_fall_window = false;

// ============ PID Controller ============
float pid_kp = 2.0;
float pid_ki = 0.5;
float pid_kd = 0.1;
float pid_integral    = 0;
float pid_prev_error  = 0;
float pid_output      = 0;
float desired_speed   = 0;
float actual_speed    = 0;
unsigned long pid_last_time = 0;

// ============ MQTT Timing ============
unsigned long last_telemetry_time = 0;
const unsigned long TELEMETRY_INTERVAL = 2000;  // 2 seconds

// ============ Complementary Filter ============
float comp_angle = 0;
const float COMP_ALPHA = 0.98;

// ============ Function Prototypes ============
void setup_wifi();
void callback(char* topic, byte* payload, unsigned int length);
void reconnect_mqtt();
void read_imu();
void detect_fall();
void pid_control();
void update_leds();
void send_telemetry();
void send_attributes();
void handle_rpc(char* topic, byte* payload, unsigned int length);

// ============================================================
//  SETUP
// ============================================================
void setup() {
  Serial.begin(115200);
  delay(1000);
  Serial.println("========================================");
  Serial.println(" Smart Walker - ESP32 Firmware v1.0");
  Serial.println(" EEE532 AIoT Project");
  Serial.println("========================================");

  // Initialize I2C for MPU6050
  Wire.begin(IMU_SDA, IMU_SCL);

  // Initialize MPU6050
  if (!mpu.begin()) {
    Serial.println("[ERROR] MPU6050 not found! Check wiring.");
    while (1) { delay(10); }
  }
  Serial.println("[OK] MPU6050 initialized.");
  mpu.setAccelerometerRange(MPU6050_RANGE_8_G);
  mpu.setGyroRange(MPU6050_RANGE_500_DEG);
  mpu.setFilterBandwidth(MPU6050_BAND_21_HZ);

  // Initialize GPIO
  pinMode(PRESSURE_PIN, INPUT);
  pinMode(BRAKE_PIN, OUTPUT);
  pinMode(LED_GREEN, OUTPUT);
  pinMode(LED_YELLOW, OUTPUT);
  pinMode(LED_RED, OUTPUT);

  // Initialize Servo (motor proxy) with explicit 50Hz and pulse range
  motorServo.setPeriodHertz(50);
  int8_t servo_channel = motorServo.attach(MOTOR_PIN, 500, 2400);
  if (servo_channel < 0) {
    Serial.println("[ERROR] Servo attach failed! Check pin/wiring.");
  } else {
    Serial.printf("[OK] Servo attached on GPIO%d (channel=%d).\n", MOTOR_PIN, servo_channel);

    // Power-on self-test sweep (~2.2s): 0 -> 180 -> 0 (demo-style smooth sweep)
    Serial.println("[SELFTEST] Servo sweep start (0 -> 180 -> 0)");
    for (int pos = 0; pos <= 180; pos += 5) {
      motorServo.write(pos);
      delay(30);
    }
    for (int pos = 180; pos >= 0; pos -= 5) {
      motorServo.write(pos);
      delay(30);
    }
    Serial.println("[SELFTEST] Servo sweep done.");
  }
  motorServo.write(0);

  // Initialize outputs
  digitalWrite(BRAKE_PIN, LOW);
  digitalWrite(LED_GREEN, HIGH);   // Normal state
  digitalWrite(LED_YELLOW, LOW);
  digitalWrite(LED_RED, LOW);

  // Connect WiFi
  setup_wifi();

  // Configure MQTT
  client.setServer(mqtt_server, mqtt_port);
  client.setCallback(callback);
  reconnect_mqtt();

  pid_last_time = millis();
  Serial.println("[OK] System ready. Starting main loop...");
  Serial.println("----------------------------------------");
}

// ============================================================
//  MAIN LOOP
// ============================================================
void loop() {
  // Ensure MQTT connection
  if (!client.connected()) {
    reconnect_mqtt();
  }
  client.loop();

  // 1. Read sensors
  read_imu();
  pressure_val = analogRead(PRESSURE_PIN);

  // 2. Fall detection
  detect_fall();

  // 3. PID motor control (only if no fall)
  if (!fall_detected) {
    pid_control();
  } else {
    // Fall detected: stop motor and engage brake
    motor_speed = 0;
    motorServo.write(0);
    brake_engaged = true;
    digitalWrite(BRAKE_PIN, HIGH);
  }

  // 4. Update LED indicators
  update_leds();

  // 5. Send telemetry to ThingsBoard every 2 seconds
  if (millis() - last_telemetry_time >= TELEMETRY_INTERVAL) {
    send_telemetry();
    send_attributes();
    last_telemetry_time = millis();
  }

  delay(50);  // 20Hz main loop
}

// ============================================================
//  WiFi Setup
// ============================================================
void setup_wifi() {
  Serial.printf("[WiFi] Connecting to %s", ssid);
  WiFi.begin(ssid, password);
  int attempts = 0;
  while (WiFi.status() != WL_CONNECTED && attempts < 40) {
    delay(500);
    Serial.print(".");
    attempts++;
  }
  if (WiFi.status() == WL_CONNECTED) {
    Serial.println(" CONNECTED!");
    Serial.printf("[WiFi] IP: %s\n", WiFi.localIP().toString().c_str());
  } else {
    Serial.println(" FAILED! Will retry in loop.");
  }
}

// ============================================================
//  MQTT Callback (RPC Handler)
// ============================================================
void callback(char* topic, byte* payload, unsigned int length) {
  Serial.printf("[RPC] Received on topic: %s\n", topic);

  // Parse JSON payload manually (lightweight)
  String msg = "";
  for (unsigned int i = 0; i < length; i++) {
    msg += (char)payload[i];
  }
  Serial.printf("[RPC] Payload: %s\n", msg.c_str());

  // Handle setBrake RPC
  if (msg.indexOf("\"setBrake\"") >= 0 || msg.indexOf("setBrake") >= 0) {
    if (msg.indexOf("true") >= 0) {
      brake_engaged = true;
      digitalWrite(BRAKE_PIN, HIGH);
      motorServo.write(0);
      motor_speed = 0;
      Serial.println("[RPC] Brake ENGAGED remotely!");
    } else if (msg.indexOf("false") >= 0) {
      brake_engaged = false;
      digitalWrite(BRAKE_PIN, LOW);
      fall_detected = false;
      Serial.println("[RPC] Brake RELEASED remotely!");
    }
  }

  // Handle setPID RPC
  if (msg.indexOf("\"setPID\"") >= 0 || msg.indexOf("setPID") >= 0) {
    // Parse Kp, Ki, Kd from JSON
    int kp_idx = msg.indexOf("\"Kp\"");
    int ki_idx = msg.indexOf("\"Ki\"");
    int kd_idx = msg.indexOf("\"Kd\"");
    if (kp_idx >= 0) pid_kp = msg.substring(msg.indexOf(":", kp_idx) + 1).toFloat();
    if (ki_idx >= 0) pid_ki = msg.substring(msg.indexOf(":", ki_idx) + 1).toFloat();
    if (kd_idx >= 0) pid_kd = msg.substring(msg.indexOf(":", kd_idx) + 1).toFloat();
    Serial.printf("[RPC] PID updated: Kp=%.2f Ki=%.2f Kd=%.2f\n", pid_kp, pid_ki, pid_kd);
    pid_integral = 0;
    pid_prev_error = 0;
  }
}

// ============================================================
//  MQTT Reconnect
// ============================================================
void reconnect_mqtt() {
  while (!client.connected()) {
    Serial.printf("[MQTT] Connecting to %s:%d ...\n", mqtt_server, mqtt_port);
    if (client.connect("SmartWalker-ESP32", token, NULL)) {
      Serial.println("[MQTT] Connected!");
      // Subscribe to RPC topic
      client.subscribe("v1/devices/me/rpc/request/+");
      Serial.println("[MQTT] Subscribed to RPC topic.");
    } else {
      Serial.printf("[MQTT] Failed, rc=%d. Retry in 5s...\n", client.state());
      delay(5000);
    }
  }
}

// ============================================================
//  Read IMU Data
// ============================================================
void read_imu() {
  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;

  // Calculate acceleration magnitude
  accel_mag = sqrt(accel_x * accel_x + accel_y * accel_y + accel_z * accel_z);

  // Complementary filter for tilt angle
  float dt = 0.05;  // 50ms loop time
  float accel_angle = atan2(accel_x, accel_z) * 180.0 / PI;
  comp_angle = COMP_ALPHA * (comp_angle + gyro_y * dt) + (1 - COMP_ALPHA) * accel_angle;
  tilt_angle = comp_angle;

  // Simulated actual speed (simple model: decays toward 0, PID pushes it up)
  actual_speed = actual_speed * 0.95 + pid_output * 0.05;
}

// ============================================================
//  Fall Detection Algorithm
// ============================================================
void detect_fall() {
  unsigned long now = millis();

  // Check if in cooldown period
  if (fall_detected && (now - last_fall_time < FALL_COOLDOWN)) {
    return;  // Still in cooldown
  }

  // Check for fall conditions
  bool high_accel = (accel_mag > FALL_ACCEL_THRESHOLD);
  bool high_tilt  = (fabs(tilt_angle) > FALL_TILT_THRESHOLD);

  if (high_accel && high_tilt) {
    if (!in_fall_window) {
      in_fall_window = true;
      fall_start_time = now;
    }
    // Check if both conditions persist for the window duration
    if (now - fall_start_time >= FALL_WINDOW) {
      fall_detected = true;
      last_fall_time = now;
      in_fall_window = false;
      Serial.println("=====================================");
      Serial.println("  ** FALL DETECTED! **");
      Serial.printf("  Accel: %.2f m/s^2 (threshold: %.2f)\n", accel_mag, FALL_ACCEL_THRESHOLD);
      Serial.printf("  Tilt:  %.1f deg (threshold: %.1f)\n", tilt_angle, FALL_TILT_THRESHOLD);
      Serial.println("  Brake ENGAGED! Alert sending...");
      Serial.println("=====================================");
    }
  } else {
    in_fall_window = false;

    // Auto-recover after cooldown if conditions return to normal
    if (fall_detected && (now - last_fall_time >= FALL_COOLDOWN)) {
      if (!high_accel && !high_tilt) {
        fall_detected = false;
        brake_engaged = false;
        digitalWrite(BRAKE_PIN, LOW);
        Serial.println("[RECOVERY] Fall cleared, brake released.");
      }
    }
  }
}

// ============================================================
//  PID Motor Control
// ============================================================
void pid_control() {
  if (brake_engaged) return;

  unsigned long now = millis();
  float dt = (now - pid_last_time) / 1000.0;
  if (dt < 0.01) dt = 0.01;  // Minimum dt
  pid_last_time = now;

  // Desired speed from pressure sensor (0-4095 → 0-180)
  desired_speed = map(pressure_val, 0, 4095, 0, 180);

  // Only activate motor assist if pressure exceeds threshold
  if (pressure_val < 500) {
    desired_speed = 0;
    pid_integral = 0;
  }

  // PID calculation
  float error = desired_speed - actual_speed;
  pid_integral += error * dt;
  // Anti-windup: clamp integral
  if (pid_integral > 200) pid_integral = 200;
  if (pid_integral < -200) pid_integral = -200;

  float derivative = (error - pid_prev_error) / dt;
  pid_output = pid_kp * error + pid_ki * pid_integral + pid_kd * derivative;
  pid_prev_error = error;

  // Clamp output
  if (pid_output > 180) pid_output = 180;
  if (pid_output < 0) pid_output = 0;

  // Apply to servo (motor proxy)
  motor_speed = (int)pid_output;
  motorServo.write(motor_speed);
}

// ============================================================
//  Update LED Indicators
// ============================================================
void update_leds() {
  if (fall_detected) {
    digitalWrite(LED_GREEN, LOW);
    digitalWrite(LED_YELLOW, LOW);
    digitalWrite(LED_RED, HIGH);
  } else if (motor_speed > 10) {
    digitalWrite(LED_GREEN, LOW);
    digitalWrite(LED_YELLOW, HIGH);
    digitalWrite(LED_RED, LOW);
  } else {
    digitalWrite(LED_GREEN, HIGH);
    digitalWrite(LED_YELLOW, LOW);
    digitalWrite(LED_RED, LOW);
  }
}

// ============================================================
//  Send Telemetry to ThingsBoard
// ============================================================
void send_telemetry() {
  if (!client.connected()) return;

  // Build JSON payload
  char payload[512];
  snprintf(payload, sizeof(payload),
    "{\"accel_x\":%.2f,\"accel_y\":%.2f,\"accel_z\":%.2f,"
    "\"accel_mag\":%.2f,\"tilt_angle\":%.1f,"
    "\"pressure\":%d,\"motor_speed\":%d,"
    "\"fall_detected\":%s,\"brake_engaged\":%s}",
    accel_x, accel_y, accel_z,
    accel_mag, tilt_angle,
    pressure_val, motor_speed,
    fall_detected ? "true" : "false",
    brake_engaged ? "true" : "false"
  );

  // Publish to ThingsBoard
  bool success = client.publish("v1/devices/me/telemetry", payload);

  // Also print to Serial for monitoring
  Serial.printf("[TEL] ax=%.2f ay=%.2f az=%.2f | mag=%.2f | tilt=%.1f | pres=%d | motor=%d | fall=%s | brake=%s %s\n",
    accel_x, accel_y, accel_z,
    accel_mag, tilt_angle,
    pressure_val, motor_speed,
    fall_detected ? "YES" : "no",
    brake_engaged ? "YES" : "no",
    success ? "(sent)" : "(FAIL)"
  );
}

// ============================================================
//  Send Attributes to ThingsBoard
// ============================================================
void send_attributes() {
  if (!client.connected()) return;

  // Mirror key status as shared telemetry + attributes for widget compatibility.
  char payload[128];
  snprintf(payload, sizeof(payload),
    "{\"fall_detected\":%s,\"brake_engaged\":%s}",
    fall_detected ? "true" : "false",
    brake_engaged ? "true" : "false"
  );

  bool success = client.publish("v1/devices/me/attributes", payload);
  Serial.printf("[ATTR] fall=%s | brake=%s %s\n",
    fall_detected ? "YES" : "no",
    brake_engaged ? "YES" : "no",
    success ? "(sent)" : "(FAIL)"
  );
}