#include <LiquidCrystal_I2C.h>
#include <Wire.h>

#ifdef STM32
  #include <STM32ADC.h>
  STM32ADC myADC(ADC1);
  const int MODE_BUTTON_PIN = PA0;
  const int DOWN_BUTTON_PIN = PA1;
  const int UP_BUTTON_PIN = PA2;
  const int POT_PIN = PA3;
  const int PWM_PIN = PA8; // Example PWM pin for STM32

  const long ADC_MIN = 0;
  const long ADC_MAX = 4095; // 12-bit ADC for STM32

  unsigned int prescaler = 72; // STM32 prescaler example
#else
  const int MODE_BUTTON_PIN = 8;
  const int DOWN_BUTTON_PIN = 7;
  const int UP_BUTTON_PIN = 6;
  const int POT_PIN = A0;
  const int PWM_PIN = 10;

  const long ADC_MIN = 0;
  const long ADC_MAX = 1023; // 10-bit ADC for Arduino Nano

  unsigned int prescaler = 8;  // Arduino prescaler
#endif

LiquidCrystal_I2C lcd(0x27, 16, 2);

const unsigned int MIN_FREQUENCY = 1;
const unsigned int MAX_DUTY_CYCLE = 100;
const unsigned int MIN_DUTY_CYCLE = 0;

enum Mode { MODE_FREQUENCY, MODE_MIN_DUTY_CYCLE, MODE_MAX_DUTY_CYCLE, NUM_MODES };
Mode currentMode = MODE_FREQUENCY;

template <typename T>
class ValuePair {
public:
    T current;
    T old;

    ValuePair(T initialCurrent, T initialOld) : current(initialCurrent), old(initialOld) {}

    bool hasChanged() const {
        return current != old;
    }

    void update() {
        old = current;
    }

    void set(T newValue) {
        current = newValue;
    }

    void add(T value) {
        current += value;
    }
};

ValuePair<long> frequency(10000, -1);
ValuePair<long> minDutyCycle(0, -1);
ValuePair<long> maxDutyCycle(100, -1);
ValuePair<float> dutyCycle(50.0, -1.0); // Store the old and current duty cycle as float
ValuePair<int> potValue(0, 0); // Store the old and current potentiometer values

unsigned int freqStep = 50;

bool displayNeedsUpdate = true; // Boolean to track if display update is needed

// Debouncing variables
unsigned long lastModeButtonPress = 0;
unsigned long lastDownButtonPress = 0;
unsigned long lastUpButtonPress = 0;
const unsigned long debounceDelay = 200; // 50 milliseconds debounce delay

/**
 * Initializes the setup by configuring pins, starting the LCD, and setting the initial display.
 */
void setup() {
  Wire.begin();
  lcd.begin(16, 2);
  lcd.backlight();
  updateDisplay();

#ifdef STM32
  pinMode(MODE_BUTTON_PIN, INPUT_PULLUP);
  pinMode(DOWN_BUTTON_PIN, INPUT_PULLUP);
  pinMode(UP_BUTTON_PIN, INPUT_PULLUP);
  pinMode(POT_PIN, INPUT_ANALOG);
#else
  pinMode(MODE_BUTTON_PIN, INPUT);
  pinMode(DOWN_BUTTON_PIN, INPUT);
  pinMode(UP_BUTTON_PIN, INPUT);
  pinMode(POT_PIN, INPUT);
#endif
  pinMode(PWM_PIN, OUTPUT);

  Serial.begin(9600);
  Serial.println(F("\nStarting PWM_Waveform"));
}

/**
 * Main loop that handles button presses, potentiometer changes, and updates the display if needed.
 */
void loop() {
  handleButtons();
  handlePotentiometer();
  updatePWM();

  if (displayNeedsUpdate) {
    updateDisplay();
    displayNeedsUpdate = false;
  }
}

/**
 * Handles the button presses to change modes or adjust values.
 */
void handleButtons() {
  unsigned long currentMillis = millis();

  if (digitalRead(MODE_BUTTON_PIN) == HIGH) {
    if (currentMillis - lastModeButtonPress > debounceDelay) {
      currentMode = static_cast<Mode>((currentMode + 1) % NUM_MODES);
      displayNeedsUpdate = true;
      lastModeButtonPress = currentMillis;
    }
  }

  if (digitalRead(DOWN_BUTTON_PIN) == HIGH) {
    if (currentMillis - lastDownButtonPress > debounceDelay) {
      changeValue(-1);
      lastDownButtonPress = currentMillis;
    }
  }

  if (digitalRead(UP_BUTTON_PIN) == HIGH) {
    if (currentMillis - lastUpButtonPress > debounceDelay) {
      changeValue(1);
      lastUpButtonPress = currentMillis;
    }
  }
}

/**
 * Changes the value based on the current mode and the direction specified.
 * 
 * @param direction The direction to change the value (positive or negative).
 */
void changeValue(int direction) {
  switch (currentMode) {
    case MODE_FREQUENCY: 
      frequency.add(direction * freqStep);
      frequency.set(max(frequency.current, (long)MIN_FREQUENCY)); // Ensure frequency is not below 1Hz
      break;
    case MODE_MIN_DUTY_CYCLE: 
      minDutyCycle.add(direction);
      minDutyCycle.set(max(minDutyCycle.current, (long)MIN_DUTY_CYCLE)); // Ensure minDutyCycle is not below 0
      break;
    case MODE_MAX_DUTY_CYCLE: 
      maxDutyCycle.add(direction);
      maxDutyCycle.set(min(max(maxDutyCycle.current, minDutyCycle.current), (long)MAX_DUTY_CYCLE)); // Ensure maxDutyCycle is not below minDutyCycle and not above 100
      break;
  }
}

/**
 * Handles changes in the potentiometer value and updates the duty cycle accordingly.
 */
void handlePotentiometer() {
#ifdef STM32
  potValue.set(myADC.read(POT_PIN)); // STM32 ADC read
#else
  potValue.set(analogRead(POT_PIN)); // Arduino Nano ADC read
#endif
}

/**
 * Updates the PWM configuration if any of the values have changed.
 */
void updatePWM() {
  if (minDutyCycle.hasChanged() || maxDutyCycle.hasChanged() || potValue.hasChanged()) {
    minDutyCycle.update();
    maxDutyCycle.update();
    potValue.update();
    displayNeedsUpdate = true;
    float computedDuty = mapFloat(potValue.current, ADC_MIN, ADC_MAX, minDutyCycle.current, maxDutyCycle.current);
    dutyCycle.set(computedDuty);
  }

  if (frequency.hasChanged() || dutyCycle.hasChanged()) {
    frequency.update();
    dutyCycle.update();
    displayNeedsUpdate = true;
    setupPWM(PWM_PIN, frequency.current, dutyCycle.current, prescaler);
  }
}

/**
 * Updates the LCD display with the current frequency and wave type in the first row,
 * and the current mode setting in the second row.
 */
void updateDisplay() {
  lcd.clear();
  
  // First row: current frequency and wave type
  lcd.setCursor(0, 0);
  displayFrequency(frequency.current);
  
  lcd.print(" ");
  lcd.print(dutyCycle.current);
  lcd.print("%");

  // Second row: current setting based on mode
  lcd.setCursor(0, 1);
  switch (currentMode) {
    case MODE_FREQUENCY:
      lcd.print("Set Freq: ");
      displayFrequency(frequency.current);
      break;
    case MODE_MIN_DUTY_CYCLE:
      lcd.print("Min Duty: ");
      lcd.print(minDutyCycle.current);
      lcd.print("% ");
      break;
    case MODE_MAX_DUTY_CYCLE:
      lcd.print("Max Duty: ");
      lcd.print(maxDutyCycle.current);
      lcd.print("% ");
      break;
  }
}

/**
 * Displays the frequency with appropriate units (Hz, kHz, MHz).
 * 
 * @param freq The frequency to display.
 */
void displayFrequency(long freq) {
  if (freq < 1000) {
    lcd.print(freq);
    lcd.print("Hz");
  } else if (freq < 1000000) {
    lcd.print(freq / 1000.0, 0);
    lcd.print("KHz");
  } else {
    lcd.print(freq / 1000000.0, 0);
    lcd.print("MHz");
  }
}

/**
 * Maps a float value from one range to another.
 * 
 * @param x The input value.
 * @param in_min The minimum value of the input range.
 * @param in_max The maximum value of the input range.
 * @param out_min The minimum value of the output range.
 * @param out_max The maximum value of the output range.
 * @return The mapped value as a float.
 */
float mapFloat(float x, float in_min, float in_max, float out_min, float out_max) {
  return (x - in_min) * (out_max - out_min) / (in_max - in_min) + out_min;
}

/**
 * Sets up the PWM configuration for the specified pin.
 * 
 * @param pin The PWM pin.
 * @param freq The frequency in Hz.
 * @param duty The duty cycle as a percentage.
 * @param prescaler The prescaler value.
 */
void setupPWM(int pin, unsigned long freq, float duty, unsigned int prescaler) {
#ifdef STM32
  // STM32 specific PWM setup
  if (pin != PA8) return; // Example for STM32 pin

  pinMode(pin, PWM);

  Timer1.setPeriod(1000000 / freq); // in microseconds
  Timer1.setPrescaler(prescaler);
  Timer1.setCompare(TIM_CH1, duty);
  Timer1.resume();
#else
  // Arduino Nano specific PWM setup
  if (pin != 10) return;

  pinMode(pin, OUTPUT);

  unsigned long timerFrequency = 16000000 / prescaler;
  unsigned long topValue = (timerFrequency / freq) - 1;

  TCCR1A = 0;
  TCCR1B = 0;
  TCNT1 = 0;

  TCCR1A = _BV(WGM11) | _BV(WGM10) | _BV(COM1B1);
  TCCR1B = _BV(WGM13) | _BV(WGM12);

  setPrescalerArduino(prescaler);

  OCR1A = topValue;
  OCR1B = (int)(((topValue + 1) * duty) / 100 - 1);
#endif
}

/**
 * Sets the prescaler for the Arduino PWM configuration.
 * 
 * @param prescaler The prescaler value.
 */
void setPrescalerArduino(unsigned int prescaler) {
  switch (prescaler) {
    case 1: TCCR1B |= _BV(CS10); break;
    case 8: TCCR1B |= _BV(CS11); break;
    case 64: TCCR1B |= _BV(CS11) | _BV(CS10); break;
    case 256: TCCR1B |= _BV(CS12); break;
    case 1024: TCCR1B |= _BV(CS12) | _BV(CS10); break;
    default: Serial.println("Invalid prescaler value"); break;
  }
}