#include <WiFiNINA.h>       // Wi-Fi module for Arduino Nano 33 IoT
#include <PubSubClient.h>   // MQTT Client for communication
#include <Arduino_LSM6DS3.h> // IMU sensor

// Wi-Fi Credentials
const char* ssid = "ALFA STUDIO";       // Replace with your Wi-Fi SSID
const char* password = "alfa2017"; // Replace with your Wi-Fi Password

// MQTT Broker Details (HiveMQ Cloud)
const char* mqtt_server = "broker.emqx.io";
const int mqtt_port = 1883;  // Non-secure MQTT Port for testing (Use 1883 instead of 8883)
const char* mqtt_user = "";  
const char* mqtt_pass = "";  

// MQTT Topics
const char* step_topic = "/GAITSENSE/StepCount";
const char* mistakes_topic = "/GAITSENSE/Mistakes";
const char* sensor1_topic = "/GAITSENSE/Sensor1";
const char* sensor2_topic = "/GAITSENSE/Sensor2";
const char* punishment_topic = "/GAITSENSE/PunishmentReason";

// WiFi and MQTT Clients
WiFiClient wifiClient; // Use WiFiClient for MQTT communication
PubSubClient client(wifiClient); // Use regular client for MQTT communication

// Pin Definitions
const int sensorBack = A0;  // Back pressure sensor (heel)
const int sensorFront = A1; // Front pressure sensor (toe)
const int motorPin = 9;     // Motor control pin
const int ledPin = LED_BUILTIN; // Built-in LED for indication

// Step Tracking Variables
bool waitingForFront = false; // Tracks if we're expecting A1 after A0
bool lastStepCorrect = false; // Prevents double counting of steps
int stepCount = 0;
int wrongSteps = 0;

// IMU Variables for Gyroscope (Yaw Change Detection)
float gyroX, gyroY, gyroZ;
float yaw = 0.0, initialYaw = 0.0; // Current and initial yaw angles
float yawThreshold = 15.0;         // Threshold for yaw deviation in degrees
unsigned long lastUpdateTime = 0;  // Time for integrating gyroscope data
unsigned long lightOnTime = 0;     // Time when the LED was turned on
unsigned long recalibrationDelay = 3000; // 3 seconds to recalibrate
bool isYawChanged = false;        // Flag to track if yaw has changed
bool calibratingDirection = false;  // Indicates if we're recalibrating direction

// Stores last punishment reason
String lastPunishmentReason = "";  

void setup() {
    Serial.begin(115200);

    // WiFi Setup
    WiFi.begin(ssid, password);
    Serial.print("Connecting to WiFi...");
    while (WiFi.status() != WL_CONNECTED) {
        delay(500);
        Serial.print(".");
    }
    Serial.println("\nConnected to WiFi");

    // MQTT Setup
    client.setServer(mqtt_server, mqtt_port);

    connectMQTT(); // Connect to MQTT broker

    // Hardware Setup
    pinMode(motorPin, OUTPUT);
    digitalWrite(motorPin, LOW);
    pinMode(ledPin, OUTPUT);
    digitalWrite(ledPin, LOW); // LED off initially

    if (!IMU.begin()) {
        Serial.println("IMU failed!");
        while (1);
    }
}

void loop() {
    if (!client.connected()) {
        connectMQTT();
    }
    client.loop();

    int backValue = analogRead(sensorBack);
    int frontValue = analogRead(sensorFront);

    // Print sensor values in Serial Monitor
    Serial.print("Correct Steps: " + String(stepCount) + " | ");
    Serial.print("Mistakes: " + String(wrongSteps) + " | ");
    Serial.print("Front Sensor: " + String(frontValue) + " | ");
    Serial.print("Back Sensor: " + String(backValue) + " | ");
    Serial.println("Punishment Reason: " + lastPunishmentReason);

    // Detect correct step sequence (A0 → A1)
    if (backValue > 1022 && !waitingForFront && !lastStepCorrect) {
        waitingForFront = true;
        Serial.println("🟡 A0 Pressed First, Waiting for A1...");
    }

    if (frontValue > 1022 && waitingForFront) {
        stepCount++;
        waitingForFront = false;
        lastStepCorrect = true;
        Serial.println("✅ Correct Step Counted");
        client.publish(step_topic, String(stepCount).c_str()); // Publish step count
    }

    // Detect incorrect step: A1 before A0
    if (frontValue > 1022 && !waitingForFront && !lastStepCorrect) {
        Serial.println("❌ Wrong Step Detected!");
        wrongSteps++;
        activateMotor("Wrong Step");
        lastStepCorrect = false;
        client.publish(mistakes_topic, String(wrongSteps).c_str()); // Publish mistake count
    }

    // Reset step tracking when both sensors are released
    if (backValue < 1022 && frontValue < 1022) {
        lastStepCorrect = false;
        waitingForFront = false;
    }

    // Publish sensor values
    client.publish(sensor1_topic, String(frontValue).c_str());
    client.publish(sensor2_topic, String(backValue).c_str());

    // ** Gyroscope Integration for Direction Change Detection **
    if (IMU.gyroscopeAvailable()) {
        unsigned long currentTime = millis();
        float deltaTime = (currentTime - lastUpdateTime) / 1000.0;
        lastUpdateTime = currentTime;

        IMU.readGyroscope(gyroX, gyroY, gyroZ);
        yaw += gyroZ * deltaTime;

        float yawDeviation = abs(yaw - initialYaw);

        if (yawDeviation > yawThreshold && !isYawChanged) {
            lightOnTime = millis();
            isYawChanged = true;
            calibratingDirection = true;
            digitalWrite(ledPin, HIGH);
        }

        if (isYawChanged && millis() - lightOnTime >= recalibrationDelay) {
            initialYaw = yaw;
            isYawChanged = false;
            calibratingDirection = false;
            digitalWrite(ledPin, LOW);
            activateMotor("Direction Change");
        }
    }

    delay(100);
}

// Function to activate motor and publish punishment reason
void activateMotor(String reason) {
    digitalWrite(motorPin, HIGH);
    delay(500);
    digitalWrite(motorPin, LOW);

    if (reason != lastPunishmentReason) {
        lastPunishmentReason = reason;
        client.publish(punishment_topic, reason.c_str()); // Publish punishment reason
    }
}

// Function to connect to MQTT broker
void connectMQTT() {
    while (!client.connected()) {
        Serial.print("Connecting to MQTT...");
        if (client.connect("ArduinoNanoClient", mqtt_user, mqtt_pass)) {
            Serial.println("Connected to MQTT");
        } else {
            Serial.print("Failed (Error ");
            Serial.print(client.state());
            Serial.println("), retrying in 5 seconds...");
            delay(500);
        }
    }
}