#include <Arduino_LSM6DS3.h> // IMU sensor
#include <ArduinoBLE.h>      // Bluetooth library

// UUID Definitions for the Service and Characteristics
#define STEP_SERVICE_UUID "bcb0e7dd-d9a9-464a-a417-f739ca977fab"  // Step Count Service UUID
#define STEP_COUNT_UUID "bcb0e7dd-d9a9-464a-a417-f739ca977fab"   // Step Count Characteristic UUID
#define MISTAKES_UUID "e4762867-2fa5-4582-be12-8c82436b8894"       // Mistakes Characteristic UUID
#define PUNISHMENT_REASON_UUID "a8bbba71-af75-485b-9e29-172d35e97611" // Punishment Reason Characteristic UUID
#define FRONT_SENSOR_UUID "fa0e678b-edab-4cdf-a7b1-acd55f841cde"   // Front Sensor Characteristic UUID
#define BACK_SENSOR_UUID "7d69c881-8cc4-470e-8e43-b63277be751a"    // Back Sensor Characteristic UUID

// 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

// Create BLE service and characteristics using the defined UUIDs
BLEService stepService(STEP_SERVICE_UUID);   // Step service
BLEIntCharacteristic stepCountChar(STEP_COUNT_UUID, BLERead | BLENotify);   // Step count characteristic
BLEIntCharacteristic wrongStepChar(MISTAKES_UUID, BLERead | BLENotify);   // Mistakes characteristic
BLEStringCharacteristic punishmentReasonChar(PUNISHMENT_REASON_UUID, BLERead | BLENotify, 20);  // Punishment reason characteristic
BLEIntCharacteristic frontSensorChar(FRONT_SENSOR_UUID, BLERead | BLENotify);  // Front sensor characteristic
BLEIntCharacteristic backSensorChar(BACK_SENSOR_UUID, BLERead | BLENotify);   // Back sensor characteristic

// Declare 'lastPunishmentReason' to store the reason for punishment
String lastPunishmentReason = "";  // Stores last punishment reason

void setup() {
    Serial.begin(115200);
    pinMode(motorPin, OUTPUT);
    digitalWrite(motorPin, LOW);
    pinMode(ledPin, OUTPUT);
    digitalWrite(ledPin, LOW); // LED off initially

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

    if (BLE.begin()) {
        BLE.setLocalName("StepTracker");
        BLE.setAdvertisedService(stepService);
        stepService.addCharacteristic(stepCountChar);
        stepService.addCharacteristic(wrongStepChar);
        stepService.addCharacteristic(punishmentReasonChar);
        stepService.addCharacteristic(frontSensorChar);  // Added to BLE service
        stepService.addCharacteristic(backSensorChar);   // Added to BLE service
        BLE.addService(stepService);
        BLE.advertise();
        Serial.println("Bluetooth active and advertising.");
    } else {
        Serial.println("Bluetooth failed!");
        while (1);
    }
}

void loop() {
    BLE.poll();

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

    // Print sensor values and punishment reason in the desired format
    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");
    }

    // Detect incorrect step: A1 before A0 (only if A1 comes first)
    if (frontValue > 1022 && !waitingForFront && !lastStepCorrect) {
        Serial.println("❌ Wrong Step Detected!");
        wrongSteps++;
        lastStepCorrect = false;
    }

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

    // ** 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);
        }
    }

    // Bluetooth updates for app
    if (BLE.connected()) {
        stepCountChar.writeValue(stepCount);
        wrongStepChar.writeValue(wrongSteps);
        punishmentReasonChar.writeValue(lastPunishmentReason);
        frontSensorChar.writeValue(frontValue); // Send front sensor data
        backSensorChar.writeValue(backValue);   // Send back sensor data
    }

    delay(100);
}

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

    if (reason != lastPunishmentReason) {
        lastPunishmentReason = reason;
        punishmentReasonChar.writeValue(reason);
    }
}