#include <LSM6DS3.h>
#include <Wire.h>
#include <math.h>
#include <TensorFlowLite.h>
#include <tensorflow/lite/micro/all_ops_resolver.h>
#include <tensorflow/lite/micro/micro_error_reporter.h>
#include <tensorflow/lite/micro/micro_interpreter.h>
#include <tensorflow/lite/schema/schema_generated.h>
#include "model.h"
#include <bluefruit.h> // Include the Bluefruit library

#define BUTTON_PIN D2       // Define button pin (D2)
#define LED_PIN LED_BUILTIN // Define LED pin (built-in LED)

const bool DEBUG_MODE = false; // Set this to false to disable serial output
const float accelerationThreshold = 5.5; // threshold of significant in G's
const int numSamples = 119;

int samplesRead = 0;
bool inSwing = false;
volatile bool sleepMode = false; // Flag to manage sleep mode
volatile bool buttonPressed = false; // Flag to manage button press

LSM6DS3 myIMU(I2C_MODE, 0x6A);

// Global variables used for TensorFlow Lite (Micro)
tflite::MicroErrorReporter tflErrorReporter;
tflite::AllOpsResolver tflOpsResolver;

const tflite::Model* tflModel = nullptr;
tflite::MicroInterpreter* tflInterpreter = nullptr;
TfLiteTensor* tflInputTensor = nullptr;
TfLiteTensor* tflOutputTensor = nullptr;

constexpr int tensorArenaSize = 8 * 1024;
byte tensorArena[tensorArenaSize] __attribute__((aligned(16)));

// Array to map gesture index to a name
const char* GESTURES[] = {
  "1",
  "2",
  "3",
  "4"
};

#define NUM_GESTURES (sizeof(GESTURES) / sizeof(GESTURES[0]))

// Parameters for moving average filter
const int filterSize = 10;
float aX_buffer[filterSize] = {0};
float aY_buffer[filterSize] = {0};
float aZ_buffer[filterSize] = {0};
int filterIndex = 0;

float velocity[3] = {0.0, 0.0, 0.0};  // Initial velocity
float accel_scale = 16384.0;  // Accelerometer scale factor
float g = 9.80665;  // Acceleration due to gravity in m/s^2
unsigned long lastTime = 0;
float aX_offset = 0.0, aY_offset = 0.0, aZ_offset = 0.0;
const int calibrationSamples = 100;

// Complementary filter parameters
const float alpha = 0.98; // Adjust this value based on sensor noise

// Variables for angle calculation
float pitch = 0.0;
float roll = 0.0;

// Low-pass filter parameters
const float cutoff_freq = 5.0; // Cutoff frequency (Hz)
float rc = 1.0 / (2 * PI * cutoff_freq);
float dt_filter = 0.0;
float aX_filtered = 0.0, aY_filtered = 0.0, aZ_filtered = 0.0;
float gX_filtered = 0.0, gY_filtered = 0.0, gZ_filtered = 0.0;

// List to store swing data
float swingData[1000][6];
int swingDataIndex = 0;

// BLE UART service
BLEUart bleuart;

// Define the DebugLog function
extern "C" void DebugLog(const char* s) {
  if (DEBUG_MODE) {
    Serial.println(s);
  }
}

// Define the test_over_serial namespace functions
namespace test_over_serial {
  void SerialWrite(const char* s) {
    if (DEBUG_MODE) {
      Serial.print(s);
    }
  }

  int SerialReadLine(int timeout_ms) {
    // Add implementation for reading a line from Serial
    return 0;
  }
}

// Charging state pin and LED
const int chargePin = 23; // P0.17 pin
const int chargeLED = 3;  // Pin for additional LED

void calibrateSensor() {
  float aX, aY, aZ;
  for (int i = 0; i < calibrationSamples; i++) {
    aX = myIMU.readFloatAccelX();
    aY = myIMU.readFloatAccelY();
    aZ = myIMU.readFloatAccelZ();
    aX_offset += aX;
    aY_offset += aY;
    aZ_offset += aZ;
    delay(10);
  }
  aX_offset /= calibrationSamples;
  aY_offset /= calibrationSamples;
  aZ_offset /= calibrationSamples;
}

void buttonISR() {
  buttonPressed = true; // Set flag to indicate button press
}

void setup() {
  // Initialize serial communication
  if (DEBUG_MODE) {
    Serial.begin(9600);
    while (!Serial);
  }

  // Initialize LED and button pins
  pinMode(LED_PIN, OUTPUT);
  pinMode(BUTTON_PIN, INPUT_PULLUP);
  attachInterrupt(digitalPinToInterrupt(BUTTON_PIN), buttonISR, FALLING); // Attach interrupt to button

  // Blink LED to indicate setup
  digitalWrite(LED_PIN, HIGH);
  delay(1000);
  digitalWrite(LED_PIN, LOW);
  delay(1000);
  digitalWrite(LED_PIN, HIGH);
  delay(1000);
  digitalWrite(LED_PIN, LOW);

  if (myIMU.begin() != 0) {
    if (DEBUG_MODE) {
      Serial.println("Device error");
    }
  } else {
    if (DEBUG_MODE) {
      Serial.println("Device OK!");
    }
  }

  calibrateSensor();

  // Get the TFL representation of the model byte array
  tflModel = tflite::GetModel(model);
  if (tflModel->version() != TFLITE_SCHEMA_VERSION) {
    if (DEBUG_MODE) {
      Serial.println("Model schema mismatch!");
    }
    while (1);
  }

  // Create an interpreter to run the model
  tflInterpreter = new tflite::MicroInterpreter(tflModel, tflOpsResolver, tensorArena, tensorArenaSize, &tflErrorReporter);

  // Allocate memory for the model's input and output tensors
  tflInterpreter->AllocateTensors();

  // Get pointers for the model's input and output tensors
  tflInputTensor = tflInterpreter->input(0);
  tflOutputTensor = tflInterpreter->output(0);

  lastTime = millis();  // Set the initial time in milliseconds

  // Initialize BLE
  Bluefruit.begin();
  Bluefruit.setTxPower(4); // Set maximum transmission power
  Bluefruit.setName("IMU_5");

  // Start BLE UART service
  bleuart.begin();

  // Advertising setup and start
  Bluefruit.Advertising.addFlags(BLE_GAP_ADV_FLAGS_LE_ONLY_GENERAL_DISC_MODE);
  Bluefruit.Advertising.addTxPower();
  Bluefruit.Advertising.addService(bleuart); // Advertising UART service
  Bluefruit.Advertising.addName();
  Bluefruit.Advertising.restartOnDisconnect(true);
  Bluefruit.Advertising.start(0); // 0 = Continue advertising without timeout

  // Setup charge detection and LED
  pinMode(chargePin, INPUT_PULLUP);
  pinMode(chargeLED, OUTPUT);
  digitalWrite(chargeLED, LOW);  // Turn off the LED initially

  if (DEBUG_MODE) {
    Serial.println("System initialized and running.");
  }
}

float movingAverage(float* buffer, float newValue) {
  buffer[filterIndex] = newValue;
  float sum = 0.0;
  for (int i = 0; i < filterSize; i++) {
    sum += buffer[i];
  }
  filterIndex = (filterIndex + 1) % filterSize;
  return sum / filterSize;
}

void calculateVelocity() {
  // Calculate velocity from saved swing data
  float dt = 0.0084; // Sampling rate (119 Hz)
  for (int i = 1; i < swingDataIndex; i++) {
    float aX_world = swingData[i][0] * cos(pitch * PI / 180.0) + swingData[i][1] * sin(pitch * PI / 180.0) * sin(roll * PI / 180.0) + swingData[i][2] * cos(roll * PI / 180.0) * sin(pitch * PI / 180.0);
    float aY_world = swingData[i][1] * cos(roll * PI / 180.0) - swingData[i][2] * sin(roll * PI / 180.0);
    float aZ_world = -swingData[i][0] * sin(pitch * PI / 180.0) + swingData[i][1] * cos(pitch * PI / 180.0) * sin(roll * PI / 180.0) + swingData[i][2] * cos(pitch * PI / 180.0) * cos(roll * PI / 180.0) - g;

    // Update velocity using trapezoidal integration
    velocity[0] += aX_world * dt;
    velocity[1] += aY_world * dt;
    velocity[2] += aZ_world * dt;
  }

  // Apply a scaling factor to adjust the velocity
  float scalingFactor = 2.0; // This factor can be adjusted based on observed results
  velocity[0] *= scalingFactor;
  velocity[1] *= scalingFactor;
  velocity[2] *= scalingFactor;

  // Calculate speed in km/h
  float speed_ms = sqrt(velocity[0] * velocity[0] + velocity[1] * velocity[1] + velocity[2] * velocity[2]);
  float speed_kmh = speed_ms * 3.6;
  if (DEBUG_MODE) {
    Serial.print("Estimated speed: ");
    Serial.print(speed_kmh, 2);
    Serial.println(" km/h");
  }

  // Send speed via BLE
  char buffer[32];
  snprintf(buffer, sizeof(buffer), "V:%.2f", speed_kmh);
  bleuart.write(buffer);

  // Reset velocity and swing data index
  velocity[0] = 0.0;
  velocity[1] = 0.0;
  velocity[2] = 0.0;
  swingDataIndex = 0;
}

// Function to enter sleep mode
void enterSleepMode() {
  // Indicate entering sleep mode by blinking LED
  for (int i = 0; i < 3; i++) {
    digitalWrite(LED_PIN, HIGH);
    delay(200);
    digitalWrite(LED_PIN, LOW);
    delay(200);
  }
  // Turn off LED
  digitalWrite(LED_PIN, LOW);

  // Indicate entering sleep mode on serial monitor
  if (DEBUG_MODE) {
    Serial.println("Entering sleep mode...");
  }

  // Stop serial communication and BLE
  if (DEBUG_MODE) {
    Serial.end();
  }
  Bluefruit.Advertising.stop();
  Bluefruit.disconnect(0); // Disconnect all connections

  // Configure wakeup from sleep mode on button press
  nrf_gpio_cfg_sense_input(digitalPinToInterrupt(BUTTON_PIN), NRF_GPIO_PIN_PULLUP, NRF_GPIO_PIN_SENSE_LOW);

  // Enter System ON sleep mode
  __WFI(); // Wait for wakeup
}

// Function to initialize after wakeup
void wakeup() {
  // Reinitialize serial communication
  if (DEBUG_MODE) {
    Serial.begin(9600);
    while (!Serial);
    Serial.println("Waking up...");
  }

  // Indicate wakeup by blinking LED
  digitalWrite(LED_PIN, HIGH);
  delay(1000);
  digitalWrite(LED_PIN, LOW);

  // Reinitialize BLE
  Bluefruit.begin();
  Bluefruit.setTxPower(4); // Set maximum transmission power
  Bluefruit.setName("BLE_Sample"); // Set device name

  // Restart BLE UART service
  bleuart.begin();

  // Setup and restart BLE advertising
  Bluefruit.Advertising.addFlags(BLE_GAP_ADV_FLAGS_LE_ONLY_GENERAL_DISC_MODE);
  Bluefruit.Advertising.addTxPower();
  Bluefruit.Advertising.addService(bleuart); // Add BLE UART service to advertising
  Bluefruit.Advertising.addName();
  Bluefruit.Advertising.restartOnDisconnect(true);
  Bluefruit.Advertising.start(0); // 0 = Continue advertising without timeout
}

void loop() {
  // Check if button was pressed
  if (buttonPressed) {
    buttonPressed = false; // Reset flag
    sleepMode = !sleepMode; // Toggle sleep mode state

    if (sleepMode) {
      if (DEBUG_MODE) {
        Serial.println("Switching to sleep mode...");
      }
      enterSleepMode(); // Enter sleep mode
    } else {
      if (DEBUG_MODE) {
        Serial.println("Switching to wakeup mode...");
      }
      wakeup(); // Wake up
    }
  }

  if (!sleepMode) {
    // Check charging state
    int chargeState = digitalRead(chargePin);
    if (chargeState == LOW) { // Detect charging state (LOW means charging)
      digitalWrite(chargeLED, HIGH); // Turn on LED when charging
    } else {
      digitalWrite(chargeLED, LOW); // Turn off LED when not charging
    }

    float aX, aY, aZ, gX, gY, gZ;
    unsigned long currentTime = millis();
    float dt = (currentTime - lastTime) / 1000.0;  // Convert to seconds
    lastTime = currentTime;

    dt_filter = dt; // Set dt for low-pass filter

    // Read the acceleration and gyroscope data
    aX = myIMU.readFloatAccelX() - aX_offset;
    aY = myIMU.readFloatAccelY() - aY_offset;
    aZ = myIMU.readFloatAccelZ() - aZ_offset;

    gX = myIMU.readFloatGyroX();
    gY = myIMU.readFloatGyroY();
    gZ = myIMU.readFloatGyroZ();

    // Apply moving average filter
    aX = movingAverage(aX_buffer, aX);
    aY = movingAverage(aY_buffer, aY);
    aZ = movingAverage(aZ_buffer, aZ);

    // Apply low-pass filter
    aX_filtered = (dt_filter / (rc + dt_filter)) * aX_filtered + (rc / (rc + dt_filter)) * aX;
    aY_filtered = (dt_filter / (rc + dt_filter)) * aY_filtered + (rc / (rc + dt_filter)) * aY;
    aZ_filtered = (dt_filter / (rc + dt_filter)) * aZ_filtered + (rc / (rc + dt_filter)) * aZ;
    gX_filtered = (dt_filter / (rc + dt_filter)) * gX_filtered + (rc / (rc + dt_filter)) * gX;
    gY_filtered = (dt_filter / (rc + dt_filter)) * gY_filtered + (rc / (rc + dt_filter)) * gY;
    gZ_filtered = (dt_filter / (rc + dt_filter)) * gZ_filtered + (rc / (rc + dt_filter)) * gZ;

    // Calculate pitch and roll from accelerometer data
    float pitchAcc = atan2(aY_filtered, sqrt(aX_filtered * aX_filtered + aZ_filtered * aZ_filtered)) * 180.0 / PI;
    float rollAcc = atan2(-aX_filtered, aZ_filtered) * 180.0 / PI;

    // Integrate gyro data to get angle
    pitch = pitch + gY_filtered * dt;
    roll = roll + gX_filtered * dt;

    // Apply complementary filter
    pitch = alpha * pitch + (1 - alpha) * pitchAcc;
    roll = alpha * roll + (1 - alpha) * rollAcc;

    // Detect swing start
    if (!inSwing && fabs(aX) + fabs(aY) + fabs(aZ) >= accelerationThreshold) {
      inSwing = true;
      samplesRead = 0;
      velocity[0] = 0.0;
      velocity[1] = 0.0;
      velocity[2] = 0.0;
      swingDataIndex = 0;
    }

    if (inSwing && samplesRead < numSamples) {
      // Save swing data
      if (swingDataIndex < 1000) {
        swingData[swingDataIndex][0] = aX_filtered;
        swingData[swingDataIndex][1] = aY_filtered;
        swingData[swingDataIndex][2] = aZ_filtered;
        swingData[swingDataIndex][3] = gX_filtered;
        swingData[swingDataIndex][4] = gY_filtered;
        swingData[swingDataIndex][5] = gZ_filtered;
        swingDataIndex++;
      }

      // Normalize the IMU data between 0 to 1 and store in the model's input tensor
      tflInputTensor->data.f[samplesRead * 6 + 0] = (aX + 4.0) / 8.0;
      tflInputTensor->data.f[samplesRead * 6 + 1] = (aY + 4.0) / 8.0;
      tflInputTensor->data.f[samplesRead * 6 + 2] = (aZ + 4.0) / 8.0;
      tflInputTensor->data.f[samplesRead * 6 + 3] = (gX + 2000.0) / 4000.0;
      tflInputTensor->data.f[samplesRead * 6 + 4] = (gY + 2000.0) / 4000.0;
      tflInputTensor->data.f[samplesRead * 6 + 5] = (gZ + 2000.0) / 4000.0;

      samplesRead++;
    }

    // Detect swing end and run inference
    if (samplesRead == numSamples) {
      inSwing = false;

      // Run inferencing
      TfLiteStatus invokeStatus = tflInterpreter->Invoke();
      if (invokeStatus != kTfLiteOk) {
        if (DEBUG_MODE) {
          Serial.println("Invoke failed!");
        }
        while (1);
      }

      // Find the gesture with the highest confidence
      int maxIndex = 0;
      float maxConfidence = tflOutputTensor->data.f[0];
      for (int i = 1; i < NUM_GESTURES; i++) {
        if (tflOutputTensor->data.f[i] > maxConfidence) {
          maxConfidence = tflOutputTensor->data.f[i];
          maxIndex = i;
        }
      }

      // Print the detected gesture
      if (DEBUG_MODE) {
        Serial.print("Detected gesture: ");
        Serial.println(GESTURES[maxIndex]);
      }

      // Send gesture via BLE
//      bleuart.write("S:");
//      bleuart.write(GESTURES[maxIndex]);
//      bleuart.write("\n");

        char buffer[32];
        snprintf(buffer, sizeof(buffer), "S:%s", GESTURES[maxIndex]);
        bleuart.write(buffer);

      // Calculate and print the velocity
      calculateVelocity();

      // Reset samplesRead for the next swing
      samplesRead = 0;
    }

    // Indicate system running on serial monitor
    //if (DEBUG_MODE) {
    //  Serial.println("System running...");
    //}
  }
}