#include <math.h>  // Needed for the exp() function
#include <avr/interrupt.h>

volatile bool flagTimer1 = false;
volatile bool flagTimer2 = false;

const int buttonPin = 2;    // Pin for the pushbutton
const int ledPins[] = {3, 4, 5, 6};  // Pins for the LEDs
int buttonState = 0;        // Variable for reading the pushbutton status
int currentLed = 0;         // Variable to keep track of the current LED

const int pulsePin = 9;
const int servoPin = 11;
volatile int angle = 0;
volatile bool updateServo = false;

// Variables for pulse generation
float avgFrequency;   // Average frequency in Hz
float pulseLength; // Pulse length in milliseconds
float avgInterval; // Average interval between pulses in milliseconds
unsigned long pulseInterval; // Interval between pulses in milliseconds
unsigned long lastPulseTime = 0; // Last time a pulse was triggered

// Variables for button-controlled LEDs
unsigned long lastButtonPress = 0;
const unsigned long debounceDelay = 250; // Debounce delay in milliseconds
bool ledChanged = false;

// Variables for decay
float decayConstant; // Decay constant for exponential decay
unsigned long previousDecayTime = 0;
const unsigned long decayInterval = 1000; // Interval for decay calculation in milliseconds

// Configuration settings for each LED state
const int numConfigs = 4;
struct Config {
  float halfLife;
  float pulseFrequency;
  float pulseWidth;
} configs[numConfigs] = {
  {2.0, 1.0, 100.0},
  {4.0, 2.0, 200.0},
  {8.0, 0.5, 50.0},
  {16.0, 4.0, 400.0}
};

void setup() {
  Serial.begin(9600);
  pinMode(servoPin, OUTPUT);

  // Initialize the pushbutton pin as an input
  pinMode(buttonPin, INPUT);

  // Initialize the LED pins as outputs
  for (int i = 0; i < 4; i++) {
    pinMode(ledPins[i], OUTPUT);
  }

  // Initialize all LEDs to off
  for (int i = 0; i < 4; i++) {
    digitalWrite(ledPins[i], LOW);
  }

  // Turn on the first LED
  digitalWrite(ledPins[currentLed], HIGH);

  // Initialize the random pulse part
  pinMode(pulsePin, OUTPUT);
  randomSeed(analogRead(0)); // Initialize random number generator with a seed

  // Initialize the first configuration
  applyConfig(0);

  // Disable global interrupts
  cli();

  // Timer1 configuration (1 second interval)
  TCCR1A = 0;
  TCCR1B = 0;
  TCCR1B |= (1 << WGM12); // CTC mode
  OCR1A = 15624;          // (16 MHz / (1 Hz * 1024)) - 1
  TCCR1B |= (1 << CS12) | (1 << CS10); // Prescaler 1024
  TIMSK1 |= (1 << OCIE1A); // Enable Timer1 compare interrupt

  // Timer2 configuration (0.5 second interval)
  TCCR2A = 0;
  TCCR2B = 0;
  TCCR2A |= (1 << WGM21); // CTC mode
  OCR2A = 124;            // (16 MHz / (2 Hz * 128)) - 1
  TCCR2B |= (1 << CS22) | (1 << CS20); // Prescaler 128
  TIMSK2 |= (1 << OCIE2A); // Enable Timer2 compare interrupt

  // Enable global interrupts
  sei();

  // Configure Timer0 for custom use while maintaining millis() and delay()
  TCCR0A = (1 << WGM01) | (1 << WGM00);
  TCCR0B = (1 << WGM02) | (1 << CS00) | (1 << CS01); // Fast PWM mode with prescaler 64
  OCR0A = 255;  // Set TOP value to 255
  TIMSK0 |= (1 << OCIE0A); // Enable compare interrupt
}

void loop() {
  // Part 1: Handle button-controlled LEDs
  handleButtonControlledLEDs();

  // Part 2: Handle pulse generation
  handlePulseGeneration();
  
  // Part 3: Decay frequency using Timer2
  handleFrequencyDecay();

  if (updateServo) 
  {
    Serial.println("!!!");
    updateServo = false;
    setServoAngle(angle);
    angle = (angle == 0) ? 180 : 0; // Alternate between 0 and 180 degrees
    //delay(1000); // Wait for 1 second
  }
}

void handleButtonControlledLEDs() {
  unsigned long currentMillis = millis();
  
  if (currentMillis - lastButtonPress >= debounceDelay) {
    buttonState = digitalRead(buttonPin);
    if (buttonState == HIGH) {
      if (!ledChanged) {
        // Turn off the current LED
        digitalWrite(ledPins[currentLed], LOW);
        
        // Move to the next LED
        currentLed = (currentLed + 1) % 4;
        
        // Turn on the new current LED
        digitalWrite(ledPins[currentLed], HIGH);
        
        // Apply the new configuration
        applyConfig(currentLed);

        ledChanged = true;
        lastButtonPress = currentMillis;
      }
    } else {
      ledChanged = false;
    }
  }
}

void handlePulseGeneration() {
  unsigned long currentMillis = millis();
  static bool pulseState = LOW;

  if (flagTimer1) {
    flagTimer1 = false;

    // Check if enough time has passed to toggle the pulse state
    if (currentMillis - lastPulseTime >= pulseInterval) {
      pulseState = !pulseState;
      digitalWrite(pulsePin, pulseState);

      if (pulseState) {
        // Schedule the end of the pulse
        lastPulseTime = currentMillis + pulseLength;
      } else {
        // Calculate the next interval using uniform distribution
        float minInterval = avgInterval * 0.5;
        float maxInterval = avgInterval * 1.5;

        // Generate random interval and ensure it's not too small
        pulseInterval = minInterval + (random(0, 1000) / 1000.0) * (maxInterval - minInterval);

        // Schedule the next pulse start
        lastPulseTime = currentMillis;
      }
    }
  }
}

void handleFrequencyDecay() {
  unsigned long currentMillis = millis();

  if (flagTimer2) {
    flagTimer2 = false;

    // Update the frequency based on decay constant
    if (currentMillis - previousDecayTime >= decayInterval) {
      previousDecayTime = currentMillis;
      avgFrequency = avgFrequency * exp(-decayConstant);
      avgInterval = 1000.0 / avgFrequency; // Recalculate the average interval
      Serial.print("Current Frequency: ");
      Serial.println(avgFrequency);
    }
  }
}

void applyConfig(int configIndex) {
  // Apply the configuration settings
  avgFrequency = configs[configIndex].pulseFrequency;
  pulseLength = configs[configIndex].pulseWidth;
  decayConstant = log(2) / configs[configIndex].halfLife; // Calculate decay constant from half-life
  avgInterval = 1000.0 / avgFrequency; // Calculate the average interval in milliseconds

  // Reset the timers
  resetTimer1();
  resetTimer2();
}

void resetTimer1() {
  TCCR1B = 0; // Stop the timer
  TCNT1 = 0;  // Clear the timer counter
  // Reinitialize Timer1 settings
  TCCR1B |= (1 << WGM12); // CTC mode
  OCR1A = 15624;          // (16 MHz / (1 Hz * 1024)) - 1
  TCCR1B |= (1 << CS12) | (1 << CS10); // Prescaler 1024
}

void resetTimer2() {
  TCCR2B = 0; // Stop the timer
  TCNT2 = 0;  // Clear the timer counter
  // Reinitialize Timer2 settings
  TCCR2A |= (1 << WGM21); // CTC mode
  OCR2A = 124;            // (16 MHz / (2 Hz * 128)) - 1
  TCCR2B |= (1 << CS22) | (1 << CS20); // Prescaler 128
}

ISR(TIMER0_COMPA_vect) {
  static uint16_t counter = 0;
  counter++;
  if (counter >= 20000) { // 20 ms period for servo PWM
    updateServo = true;
    counter = 0;
  }
}

ISR(TIMER1_COMPA_vect) {
  flagTimer1 = true;
}

ISR(TIMER2_COMPA_vect) {
  flagTimer2 = true;
}

void setServoAngle(int angle) {
  uint16_t pulseWidth = map(angle, 0, 180, 544, 2400); // Map angle to pulse width (microseconds)
  digitalWrite(servoPin, HIGH);
  delayMicroseconds(pulseWidth);
  digitalWrite(servoPin, LOW);
}