// Tactile Labyrinth - Axis Direction Test
// Lights up LEDs closest to each of the 6 cardinal directions

#include <FastLED.h>

// Pin Definitions
#define LED_PIN     D6
#define PIN_BUTTON  D10
#define PIN_LED_R   17  // Built-in LEDs on XIAO RP2040
#define PIN_LED_G   16
#define PIN_LED_B   25

// Constants
#define NUM_LEDS    250
#define BRIGHTNESS  64
#define LED_TYPE    WS2812B
#define COLOR_ORDER GRB
#define BUTTON_PRESSED LOW  // Pull-up means button goes LOW when pressed
#define BUILTIN_LED_ON  LOW  // Built-in LEDs are active LOW
#define BUILTIN_LED_OFF HIGH

// LED array
CRGB leds[NUM_LEDS];

// LED position storage - xyz coordinates for each LED
struct Vector3 {
  float x, y, z;
};
Vector3 ledPositions[NUM_LEDS];

// Store indices of LEDs closest to cardinal directions
int cardinalLEDs[6]; // +X, -X, +Y, -Y, +Z, -Z

// Direction names for serial output
const char* directionNames[6] = {"+X", "-X", "+Y", "-Y", "+Z", "-Z"};

// Colors for each cardinal direction
const CRGB directionColors[6] = {
  CRGB::Red,      // +X
  CRGB::DarkRed,  // -X
  CRGB::Green,    // +Y
  CRGB::DarkGreen,// -Y
  CRGB::Blue,     // +Z
  CRGB::Purple    // -Z
};

// State variables
int displayMode = 0;  // 0=all directions, 1=only +X, 2=only -X, etc.
unsigned long lastButtonPress = 0;  // For debouncing

// Function to calculate position for LED along the double coil
void getPosition(float s, float& x, float& y, float& z) {
  const float a = 7.5;  // Controls number of turns

  // Map s to angle around the sphere (0 to 2π)
  float angle = a * PI * s;

  // Create a pattern that goes from South to North to South
  float height = -cos(2 * PI * s);  // -1 at s=0, +1 at s=0.5, -1 at s=1

  // Calculate position
  x = sqrt(1.0 - height*height) * cos(angle);
  y = height;  // Vertical axis
  z = sqrt(1.0 - height*height) * sin(angle);
}

// Normalize a vector to unit length
void normalizeVector(float& x, float& y, float& z) {
  float length = sqrt(x*x + y*y + z*z);
  if (length > 0) {
    x /= length;
    y /= length;
    z /= length;
  }
}

// Find LED closest to a given direction
int findClosestLED(float x, float y, float z) {
  int closestLED = 0;
  float bestMatch = -1.0;  // Dot product will be in range [-1,1], with 1 being perfect match

  for (int i = 0; i < NUM_LEDS; i++) {
    // Dot product measures how aligned two unit vectors are
    float dotProduct = ledPositions[i].x * x +
                       ledPositions[i].y * y +
                       ledPositions[i].z * z;

    if (dotProduct > bestMatch) {
      bestMatch = dotProduct;
      closestLED = i;
    }
  }

  return closestLED;
}

// Setup function to calculate all LED positions and find cardinal directions
void calculateLedPositions() {
  // Calculate positions for all LEDs
  for (int i = 0; i < NUM_LEDS; i++) {
    float s = (float)i / (NUM_LEDS - 1);
    getPosition(s, ledPositions[i].x, ledPositions[i].y, ledPositions[i].z);
    normalizeVector(ledPositions[i].x, ledPositions[i].y, ledPositions[i].z);
  }

  // Find LEDs closest to cardinal directions
  cardinalLEDs[0] = findClosestLED( 1.0,  0.0,  0.0);  // +X
  cardinalLEDs[1] = findClosestLED(-1.0,  0.0,  0.0);  // -X
  cardinalLEDs[2] = findClosestLED( 0.0,  1.0,  0.0);  // +Y
  cardinalLEDs[3] = findClosestLED( 0.0, -1.0,  0.0);  // -Y
  cardinalLEDs[4] = findClosestLED( 0.0,  0.0,  1.0);  // +Z
  cardinalLEDs[5] = findClosestLED( 0.0,  0.0, -1.0);  // -Z

  // Print out the indices of cardinal LEDs
  Serial.println("Cardinal direction LEDs:");
  for (int i = 0; i < 6; i++) {
    Serial.print(directionNames[i]);
    Serial.print(": LED #");
    Serial.print(cardinalLEDs[i]);
    Serial.print(" at position (");
    Serial.print(ledPositions[cardinalLEDs[i]].x, 3);
    Serial.print(", ");
    Serial.print(ledPositions[cardinalLEDs[i]].y, 3);
    Serial.print(", ");
    Serial.print(ledPositions[cardinalLEDs[i]].z, 3);
    Serial.println(")");
  }
}

// Update display based on current mode
void updateDisplay() {
  // First clear all LEDs
  FastLED.clear();

  if (displayMode == 0) {
    // Show all cardinal directions at once
    for (int i = 0; i < 6; i++) {
      leds[cardinalLEDs[i]] = directionColors[i];
    }
  } else {
    // Show just one direction (displayMode 1-6 corresponds to directions 0-5)
    int dirIndex = displayMode - 1;
    leds[cardinalLEDs[dirIndex]] = directionColors[dirIndex];
  }

  FastLED.show();
}

// Update built-in LED to indicate current display mode
void updateBuiltinLED() {
  digitalWrite(PIN_LED_R, BUILTIN_LED_OFF);
  digitalWrite(PIN_LED_G, BUILTIN_LED_OFF);
  digitalWrite(PIN_LED_B, BUILTIN_LED_OFF);

  if (displayMode == 0) {
    // All directions mode - all LEDs on
    digitalWrite(PIN_LED_R, BUILTIN_LED_ON);
    digitalWrite(PIN_LED_G, BUILTIN_LED_ON);
    digitalWrite(PIN_LED_B, BUILTIN_LED_ON);
  } else {
    // Single direction mode - match the direction color
    int dirIndex = displayMode - 1;
    CRGB color = directionColors[dirIndex];

    if (color.r > 100) digitalWrite(PIN_LED_R, BUILTIN_LED_ON);
    if (color.g > 100) digitalWrite(PIN_LED_G, BUILTIN_LED_ON);
    if (color.b > 100) digitalWrite(PIN_LED_B, BUILTIN_LED_ON);
  }
}

// Handle button press and change display mode
void checkButton() {
  if (digitalRead(PIN_BUTTON) == BUTTON_PRESSED) {
    // Debounce
    if (millis() - lastButtonPress > 300) {
      lastButtonPress = millis();

      // Cycle through modes (0 = all directions, 1-6 = individual directions)
      displayMode = (displayMode + 1) % 7;

      // Update display
      updateBuiltinLED();
      updateDisplay();

      // Print current mode
      Serial.print("Display mode: ");
      if (displayMode == 0) {
        Serial.println("All cardinal directions");
      } else {
        Serial.print("Direction ");
        Serial.println(directionNames[displayMode-1]);
      }
    }
  }
}

void setup() {
  Serial.begin(115200);
  Serial.println("Tactile Labyrinth - Axis Direction Test");

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

  // Calculate LED positions and find cardinal points
  calculateLedPositions();

  Serial.println("Setup complete. Press button to cycle display modes.");

  // Show initial display
  updateBuiltinLED();
  updateDisplay();
}

void loop() {
  // Check for button press
  checkButton();

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