#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <Arduino.h>

// Define OLED display size
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

// Define I2C address for SSD1306
#define OLED_I2C_ADDRESS 0x3C

// Define OLED pins
#define OLED_SDA 21 // GPIO for I2C Data Line
#define OLED_SCL 22 // GPIO for I2C Clock Line

// Create an OLED object with the I2C address 0x3C
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// Define sensor pins
const int sensorPins[5] = {34, 35, 32, 33, 25};

// Define ADC thresholds for colors
const int whiteThreshold = 3500;
const int greenThreshold = 2500;
const int redThreshold = 1500;
const int blueThreshold = 500;
const int blackThreshold = 10;

// Define variables for both timing mechanisms
const float speed = 10.0; // Speed of MARV in mm/s
const float distance = 50.0; // Distance between sensors in mm

// Timing variables for the first mechanism
unsigned long startTime1 = 0; // Start time for timing
unsigned long endTime1 = 0; // End time for timing
bool timerRunning1 = false; // Flag to indicate if the timer is running

// Timing variables for the second mechanism
unsigned long startTime2 = 0; // Start time for timing
unsigned long endTime2 = 0; // End time for timing
bool timerRunning2 = false; // Flag to indicate if the timer is running

float timeDifference1 = 0.0; // Time difference for the first mechanism
float timeDifference2 = 0.0; // Time difference for the second mechanism
float angle1 = 0.0; // Angle of incidence for the first timing mechanism
float angle2 = 0.0; // Angle of incidence for the second timing mechanism

void setup() {
  // Initialize the I2C communication with SDA and SCL pins
  Wire.begin(OLED_SDA, OLED_SCL); // Specify the SDA and SCL pins

  // Initialize the OLED with dimensions: 128 columns and 64 rows
  if (!display.begin(SSD1306_SWITCHCAPVCC, OLED_I2C_ADDRESS)) { // Use 0x3C as the address
    Serial.println(F("SSD1306 allocation failed"));
    for (;;); // Infinite loop if allocation fails
  }
  display.clearDisplay();
  display.setTextSize(1); // Set text size to 1
  display.setTextColor(SSD1306_WHITE); // Set text color to white
  display.display();
  
  // Initialize sensor pins
  for (int i = 0; i < 5; i++) {
    pinMode(sensorPins[i], INPUT);
  }

  // Start serial communication for input
  Serial.begin(9600);
}

void loop() {
  // Read ADC values from sensors and determine colors
  String sensorColors = "";
  for (int i = 0; i < 5; i++) {
    int adcValue = analogRead(sensorPins[i]);

    // Determine color based on thresholds
    if (adcValue >= whiteThreshold) {
      sensorColors += "W "; // White
    } else if (adcValue >= greenThreshold) {
      sensorColors += "G "; // Green
    } else if (adcValue >= redThreshold) {
      sensorColors += "R "; // Red
    } else if (adcValue >= blueThreshold) {
      sensorColors += "B "; // Blue
    } else if (adcValue >= blackThreshold) {
      sensorColors += "K "; // Black
    } else {
      sensorColors += "? "; // Unknown
    }
  }

  // First Timing Mechanism
  int sensor1Value = analogRead(sensorPins[0]); // Read sensor 1
  if (sensor1Value < whiteThreshold) { // Not white
    if (!timerRunning1) {
      startTime1 = millis(); // Start the timer
      timerRunning1 = true;
    }
  }

  int sensor2Value = analogRead(sensorPins[1]); // Read sensor 2
  if (sensor2Value < whiteThreshold) { // Not white
    if (timerRunning1) {
      endTime1 = millis(); // Stop the timer
      timeDifference1 = (endTime1 - startTime1) / 1000.0; // Convert to seconds
      angle1 = atan((speed * timeDifference1) / distance) * (180.0 / PI); // Calculate angle
      timerRunning1 = false;
    }
  }

  // Second Timing Mechanism
  int sensor5Value = analogRead(sensorPins[4]); // Read sensor 5
  if (sensor5Value < whiteThreshold) { // Not white
    if (!timerRunning2) {
      startTime2 = millis(); // Start the timer
      timerRunning2 = true;
    }
  }

  int sensor4Value = analogRead(sensorPins[3]); // Read sensor 4
  if (sensor4Value < whiteThreshold) { // Not white
    if (timerRunning2) {
      endTime2 = millis(); // Stop the timer
      timeDifference2 = (endTime2 - startTime2) / 1000.0; // Convert to seconds
      angle2 = atan((speed * timeDifference2) / distance) * (180.0 / PI); // Calculate angle
      timerRunning2 = false;
    }
  }

  // Display sensor colors and angle on the OLED
  display.clearDisplay();
  display.setCursor(0, 0);
  display.print(sensorColors);
  display.setCursor(0, 20);
  display.print("AOI1: ");
  display.print(angle1, 2); // Display angle for the first timing mechanism with 2 decimal places
  display.print(" deg");

  display.setCursor(0, 40);
  display.print("AOI2: ");
  display.print(angle2, 2); // Display angle for the second timing mechanism with 2 decimal places
  display.print(" deg");
  
  display.display();

  // Add a delay for readability
  delay(500); // Update every 100 milliseconds
}