// PixelBallCalibration.ino
// Interactive calibration system for mapping LEDs on a spherical ball
// using a 6DoF accelerometer to detect orientation

#include <FastLED.h>
#include <Wire.h>
#include <Adafruit_Sensor.h>
#include <Adafruit_LSM6DSOX.h>
#include "PixelBallLEDMap.h"

// Pin Definitions
#define LED_PIN     D6
#define PIN_BUTTON  D10

// LED Settings
#define NUM_LEDS        SPHERE_NUM_LEDS
#define BRIGHTNESS      64
#define LED_TYPE        WS2812B
#define COLOR_ORDER     RGB

// Button state
#define BUTTON_PRESSED  LOW
#define BUILTIN_LED_ON  LOW
#define BUILTIN_LED_OFF HIGH

// Calibration Settings
#define CALIB_BLINK_TIME_MS    3000  // Time to blink LED before capturing
#define CALIB_CONFIRM_TIME_MS  1000  // Time to show confirmation color
#define CALIB_STEP             4    // Calibrate every Nth LED (default)

// Calibration States
enum CalibrationState {
  IDLE,              // Not calibrating
  WAIT_FOR_POSITION, // Blinking current LED, waiting for positioning
  CAPTURING_DATA,    // Currently capturing accelerometer data
  CONFIRMING_SAVED,  // Showing confirmation LED
  COMPLETED          // Calibration completed
};

// LED array
CRGB leds[NUM_LEDS];

// Accelerometer sensor
Adafruit_LSM6DSOX lsm6ds;

// Current state variables
CalibrationState calibState = IDLE;
int currentLed = 0;
int ledStep = CALIB_STEP;
unsigned long stateStartTime = 0;
int lastCalibratedLed = -1; // Tracks the last LED that was calibrated

// Button debouncing
unsigned long lastButtonPressTime = 0;
const unsigned long debounceTime = 300;

// Serial Communication
#define SERIAL_BAUD_RATE  115200

// Message Headers
const char* MSG_START_CALIBRATION = "START_CALIBRATION";
const char* MSG_SET_INDEX = "SET_INDEX";
const char* MSG_LED_INFO = "LED_INFO";
const char* MSG_ACCEL_DATA = "ACCEL_DATA";
const char* MSG_SAVE_POSITION = "SAVE_POSITION";
const char* MSG_CALIBRATION_COMPLETE = "CALIBRATION_COMPLETE";
const char* MSG_GET_MAPPING = "GET_MAPPING";
const char* MSG_MAPPING_DATA = "MAPPING_DATA";
const char* MSG_RESET_CALIBRATION = "RESET_CALIBRATION";

// Function for blinking all LEDs to indicate status
void flashAllLEDs(CRGB color, int times, int delayMs) {
  for (int i = 0; i < times; i++) {
    fill_solid(leds, NUM_LEDS, color);
    FastLED.show();
    delay(delayMs);
    fill_solid(leds, NUM_LEDS, CRGB::Black);
    FastLED.show();
    delay(delayMs);
  }
}

// Read current accelerometer data
sensors_event_t accel;
sensors_event_t gyro;
sensors_event_t temp;

void readAccelerometer() {
  lsm6ds.getEvent(&accel, &gyro, &temp);

  // Send current accelerometer data periodically
  static unsigned long lastAccelSendTime = 0;
  if (millis() - lastAccelSendTime > 100) { // Send every 100ms
    String accelMsg = String(MSG_ACCEL_DATA) + "," +
                     String(accel.acceleration.x, 4) + "," +
                     String(accel.acceleration.y, 4) + "," +
                     String(accel.acceleration.z, 4);
    Serial.println(accelMsg);
    lastAccelSendTime = millis();
  }
}

// Start calibration process
void startCalibration(int startLed = 0, int step = CALIB_STEP) {
  if (calibState != IDLE && calibState != COMPLETED) {
    Serial.println("ERROR: Calibration already in progress");
    return;
  }

  // Set initial LED and step size
  currentLed = constrain(startLed, 0, NUM_LEDS - 1);
  ledStep = constrain(step, 1, NUM_LEDS);

  // Reset all LEDs
  fill_solid(leds, NUM_LEDS, CRGB::Black);

  // Start calibration
  calibState = WAIT_FOR_POSITION;
  stateStartTime = millis();

  // Send confirmation
  Serial.println("Calibration started. First LED: " + String(currentLed) + ", Step: " + String(ledStep));
  Serial.println(getLEDSummary());
  sendLEDInfo();
}

// Send current LED information
void sendLEDInfo() {
  String ledInfo = String(MSG_LED_INFO) + "," +
                 String(currentLed) + "," +
                 String((calibState == WAIT_FOR_POSITION) ? "POSITIONING" :
                        (calibState == CAPTURING_DATA) ? "CAPTURING" :
                        (calibState == CONFIRMING_SAVED) ? "CONFIRMING" :
                        (calibState == COMPLETED) ? "COMPLETED" : "IDLE");
  Serial.println(ledInfo);
}

// Save current position for the active LED
void saveCurrentPosition() {
  // Only save if in appropriate state
  if (calibState != WAIT_FOR_POSITION) {
    Serial.println("ERROR: Not ready to save position");
    return;
  }

  // Read current accelerometer data
  readAccelerometer();

  // Save this position for the current LED
  saveLEDPosition(currentLed, accel.acceleration.x, accel.acceleration.y, accel.acceleration.z);

  // Send confirmation of save
  String saveMsg = String(MSG_SAVE_POSITION) + "," +
                  String(currentLed) + "," +
                  String(accel.acceleration.x, 4) + "," +
                  String(accel.acceleration.y, 4) + "," +
                  String(accel.acceleration.z, 4);
  Serial.println(saveMsg);
  
  // Look for previous calibrated LED to interpolate from
  if (lastCalibratedLed >= 0 && lastCalibratedLed < currentLed) {
    // Interpolate positions for LEDs between the last calibrated one and this one
    bool interpolated = interpolateBetweenLEDs(lastCalibratedLed, currentLed);
    if (interpolated) {
      Serial.println("Interpolated positions for LEDs " + 
                    String(lastCalibratedLed + 1) + " to " + 
                    String(currentLed - 1));
    }
  }
  
  // Remember this as the last calibrated LED
  lastCalibratedLed = currentLed;

  // Print updated LED status
  Serial.println(getLEDSummary());

  // Change state to confirmation
  calibState = CONFIRMING_SAVED;
  stateStartTime = millis();
}

// Move to the next LED in calibration sequence
void moveToNextLED() {
  // Store the current LED index before we move to the next one
  int prevCalibratedLed = currentLed;
  
  // Increment the LED index
  currentLed += ledStep;

  // Check if calibration is complete
  if (currentLed >= NUM_LEDS) {
    completeCalibration();
    return;
  }
  
  // Interpolate positions for skipped LEDs
  if (ledStep > 1 && prevCalibratedLed < currentLed) {
    // Interpolate between the last calibrated LED and the next one to be calibrated
    Serial.println("Note: Will interpolate positions for LEDs " + 
                  String(prevCalibratedLed + 1) + " to " + 
                  String(currentLed - 1) + 
                  " after this LED is calibrated");
                  
    // We can't interpolate now because the current LED is not yet calibrated
    // The interpolation will happen after the next saveCurrentPosition() call
  }

  // Reset to positioning state
  calibState = WAIT_FOR_POSITION;
  stateStartTime = millis();
  sendLEDInfo();
}

// Set calibration to a specific LED index
void setCalibrationLED(int ledIndex) {
  if (ledIndex >= 0 && ledIndex < NUM_LEDS) {
    currentLed = ledIndex;

    // If we're currently calibrating, update the state
    if (calibState != IDLE && calibState != COMPLETED) {
      calibState = WAIT_FOR_POSITION;
      stateStartTime = millis();
    }

    Serial.println("Current LED set to: " + String(currentLed));
    sendLEDInfo();
  } else {
    Serial.println("ERROR: LED index out of range");
  }
}

// Reset all calibration data
void resetCalibrationData() {
  // Initialize all position data with forceClear=true
  initializeLEDMap(true);
  
  // Flash all LEDs red to indicate reset
  flashAllLEDs(CRGB::Red, 3, 300);
  
  Serial.println("Calibration data reset");
  Serial.println(getLEDSummary());
  
  // Reset current state if needed
  if (calibState != IDLE) {
    calibState = IDLE;
  }
  
  // Send updated mapping data
  sendMappingData();
}

// Complete calibration process
void completeCalibration() {
  // Find and interpolate any remaining uncalibrated LEDs
  int numInterpolated = 0;
  
  // First, collect all calibrated LED indices
  int calibratedLEDs[NUM_LEDS];
  int numCalibratedLEDs = 0;
  
  for (int i = 0; i < NUM_LEDS; i++) {
    if (sphereLedPositions[i].isCalibrated) {
      calibratedLEDs[numCalibratedLEDs++] = i;
    }
  }
  
  // If we have at least 2 calibrated LEDs, we can interpolate
  if (numCalibratedLEDs >= 2) {
    // Interpolate between consecutive calibrated LEDs
    for (int i = 0; i < numCalibratedLEDs - 1; i++) {
      int startIdx = calibratedLEDs[i];
      int endIdx = calibratedLEDs[i+1];
      
      if (endIdx > startIdx + 1) {
        bool success = interpolateBetweenLEDs(startIdx, endIdx);
        if (success) {
          numInterpolated += (endIdx - startIdx - 1);
        }
      }
    }
    
    // Handle wrap-around between last and first
    int startIdx = calibratedLEDs[numCalibratedLEDs-1];
    int endIdx = calibratedLEDs[0] + NUM_LEDS; // Add NUM_LEDS to handle wrap-around
    
    if (endIdx > startIdx + 1) {
      bool success = interpolateBetweenLEDs(startIdx, endIdx);
      if (success) {
        numInterpolated += (NUM_LEDS - startIdx - 1) + calibratedLEDs[0];
      }
    }
  }
  
  if (numInterpolated > 0) {
    Serial.println("Interpolated " + String(numInterpolated) + " uncalibrated LEDs during completion");
  }
  
  // Print final LED status
  Serial.println(getLEDSummary());
  
  // Update cardinal directions
  updateCardinalDirections();

  // Set state to completed
  calibState = COMPLETED;

  // Clear all LEDs
  fill_solid(leds, NUM_LEDS, CRGB::Black);
  FastLED.show();

  // Flash all LEDs green to indicate completion
  flashAllLEDs(CRGB::Green, 3, 300);

  // Send completion message
  Serial.println(MSG_CALIBRATION_COMPLETE);

  // Send the mapping data
  sendMappingData();
}

// Send mapping data over serial
void sendMappingData() {
  // Generate the JSON mapping data
  String mappingJson = generateMappingJSON();
  
  // Send with consistent message format
  Serial.println(String(MSG_MAPPING_DATA) + ":" + mappingJson);
}

// Process serial commands
void processSerialCommand() {
  if (Serial.available() > 0) {
    String command = Serial.readStringUntil('\n');
    command.trim();

    // Parse command
    int commaIndex = command.indexOf(',');
    String cmd = (commaIndex > 0) ? command.substring(0, commaIndex) : command;

    if (cmd.equals(MSG_START_CALIBRATION)) {
      // Extract parameters if available
      int startLed = 0;
      int step = CALIB_STEP;

      if (commaIndex > 0) {
        int secondComma = command.indexOf(',', commaIndex + 1);
        if (secondComma > 0) {
          startLed = command.substring(commaIndex + 1, secondComma).toInt();
          step = command.substring(secondComma + 1).toInt();
        } else {
          startLed = command.substring(commaIndex + 1).toInt();
        }
      }

      startCalibration(startLed, step);
    }
    else if (cmd.equals(MSG_SET_INDEX)) {
      if (commaIndex > 0) {
        int ledIndex = command.substring(commaIndex + 1).toInt();
        setCalibrationLED(ledIndex);
      } else {
        Serial.println("ERROR: Missing LED index");
      }
    }
    else if (cmd.equals(MSG_SAVE_POSITION)) {
      saveCurrentPosition();
    }
    else if (cmd.equals(MSG_GET_MAPPING)) {
      sendMappingData();
    }
    else if (cmd.equals(MSG_RESET_CALIBRATION)) {
      resetCalibrationData();
    }
    else {
      Serial.println("ERROR: Unknown command '" + cmd + "'");
    }
  }
}

// Update LEDs based on calibration state
void updateLEDs() {
  // First, clear all LEDs
  fill_solid(leds, NUM_LEDS, CRGB::Black);

  // Then update based on current state
  switch (calibState) {
    case WAIT_FOR_POSITION: {
      // Blink the current LED red
      unsigned long elapsed = millis() - stateStartTime;
      if ((elapsed / 500) % 2 == 0) { // Blink every 500ms
        leds[currentLed] = CRGB::Red;
      }
      break;
    }

    case CAPTURING_DATA:
      // Show solid yellow while capturing
      leds[currentLed] = CRGB::Yellow;
      break;

    case CONFIRMING_SAVED:
      // Show solid green for confirmation
      leds[currentLed] = CRGB::Green;
      break;

    case COMPLETED:
      // Rainbow pattern to show completion
      fill_rainbow(leds, NUM_LEDS, beatsin8(20, 0, 255), 255 / NUM_LEDS);
      break;

    case IDLE:
    default:
      // In idle mode, show all positioned LEDs with different colors
      int firstUncalibrated = -1;
      for (int i = 0; i < NUM_LEDS; i++) {
        if (sphereLedPositions[i].isCalibrated) {
          // Show directly calibrated LEDs in dim green
          leds[i] = CRGB(0, 32, 0); // Dim green
        } else if (sphereLedPositions[i].isSet) {
          // Show interpolated LEDs in dim blue
          leds[i] = CRGB(0, 0, 16); // Dim blue
        } else if (firstUncalibrated == -1) {
          // Remember the first uncalibrated LED
          firstUncalibrated = i;
        }
      }
      
      // Highlight the first uncalibrated LED if there is one
      if (firstUncalibrated != -1) {
        leds[firstUncalibrated] = CRGB::Blue;
      }
      break;
  }

  // Show the updated LEDs
  FastLED.show();
}

// Handle calibration state machine
void updateCalibrationState() {
  unsigned long currentTime = millis();
  unsigned long elapsed = currentTime - stateStartTime;

  switch (calibState) {
    case WAIT_FOR_POSITION:
      // Check if we exceed timeout
      if (elapsed > 30000) { // 30 second timeout
        Serial.println("WARNING: Positioning timeout. Press button or send save command to continue.");
        stateStartTime = currentTime; // Reset timer to prevent multiple warnings
      }
      break;

    case CONFIRMING_SAVED:
      // Move to next LED after confirmation time
      if (elapsed > CALIB_CONFIRM_TIME_MS) {
        moveToNextLED();
      }
      break;

    default:
      // Other states don't have automatic transitions
      break;
  }
}

// Setup function
void setup() {
  // Initialize serial communication
  Serial.begin(SERIAL_BAUD_RATE);
  Serial.println("Pixel Ball Calibration System");

  // Initialize button with pull-up
  pinMode(PIN_BUTTON, INPUT_PULLUP);

  // Initialize built-in LEDs
  pinMode(PIN_LED_R, OUTPUT);
  pinMode(PIN_LED_G, OUTPUT);
  pinMode(PIN_LED_B, OUTPUT);
  digitalWrite(PIN_LED_R, BUILTIN_LED_OFF);
  digitalWrite(PIN_LED_G, BUILTIN_LED_OFF);
  digitalWrite(PIN_LED_B, BUILTIN_LED_OFF);

  // Initialize LED strip
  FastLED.addLeds<LED_TYPE, LED_PIN, COLOR_ORDER>(leds, NUM_LEDS).setCorrection(TypicalLEDStrip);
  FastLED.setBrightness(BRIGHTNESS);

  // Clear all LEDs
  fill_solid(leds, NUM_LEDS, CRGB::Black);
  FastLED.show();

  // Initialize the accelerometer
  Serial.println("Initializing LSM6DS accelerometer...");
  if (!lsm6ds.begin_I2C()) {
    Serial.println("ERROR: Failed to find LSM6DS accelerometer!");
    digitalWrite(PIN_LED_R, BUILTIN_LED_ON); // Red LED for error
    while (1) yield();
  }

  // Set accelerometer range and data rate
  lsm6ds.setAccelRange(LSM6DS_ACCEL_RANGE_2_G);
  lsm6ds.setAccelDataRate(LSM6DS_RATE_104_HZ);
  lsm6ds.setGyroRange(LSM6DS_GYRO_RANGE_250_DPS);
  lsm6ds.setGyroDataRate(LSM6DS_RATE_104_HZ);

  // Initialize the LED mapping
  initializeLEDMap();

  // Flash all LEDs to show ready state
  flashAllLEDs(CRGB::Blue, 2, 250);

  // Check how many positions are calibrated
  int calibratedCount = 0;
  for (int i = 0; i < NUM_LEDS; i++) {
    if (sphereLedPositions[i].isCalibrated) {
      calibratedCount++;
    }
  }
  
  Serial.println("System ready for calibration.");
  Serial.println(getLEDSummary());
  Serial.println("Send '" + String(MSG_START_CALIBRATION) + "' to begin calibration.");
  Serial.println("Send '" + String(MSG_SET_INDEX) + "' to set specific LED.");
  Serial.println("Button press: Save current position when calibrating.");
}

// Main loop
void loop() {
  // Process serial commands
  processSerialCommand();

  // Read accelerometer data
  readAccelerometer();

  // Check for button press - saves current position in calibration mode
  if (digitalRead(PIN_BUTTON) == BUTTON_PRESSED) {
    if (millis() - lastButtonPressTime > debounceTime) {
      lastButtonPressTime = millis();

      if (calibState == WAIT_FOR_POSITION) {
        saveCurrentPosition();
      } else if (calibState == IDLE) {
        // Start calibration when button pressed in idle mode
        startCalibration();
      }
    }
  }

  // Update calibration state
  updateCalibrationState();

  // Update LEDs based on current state
  updateLEDs();

  // Small delay for stability
  delay(10);
}