#include <Wire.h>
#include <LiquidCrystal.h>
#include <math.h>

// LCD configuration
LiquidCrystal lcd(12, A2, 11, 6, 5, 13, A3, 10, 9, 8, 7);

// Analog/Digital pins 
const byte analogPinF = A0; // Forward power
const byte analogPinR = A1; // Reflected power
const byte analogPinPeak = A6; // Peak detector input
#define buttonMenuPin  4
#define buttonSelectPin 2

// Power measurement variables
int Vreference = 5080; // mV
const float adcToVoltage = Vreference / 1023.0f;
int adcValueF = 0;
int adcValueR = 0; 
int adcValuePeak = 0;
float pow_fwd = 0;
float pow_ref = 0;
float pow_peak = 0;
float peak_power = 0;
float peak_power_ref = 0;
float last_gamma = 0.0f;
float last_peak_power = 0.0f;
float swr = 1.0;
unsigned long last_peak_time = 0;
unsigned long last_swr_peak_reset = 0;
unsigned long peak_hold_time = 1000; // POWER/SWR Default hold time (ms)
const unsigned long peak_hold_options[] = {500, 1000, 2000, 3000}; // Available options of POWER/SWR Default hold time (ms)
byte current_hold_option = 1; // Default to 1000ms (index 1)

// Display mode options
enum DisplayMode {
  DISPLAY_FWD_POWER,
  DISPLAY_REF_POWER,
  DISPLAY_PEAK_POWER
};
byte currentDisplayMode = DISPLAY_FWD_POWER;

// Bar graph variables
static int last_lit_segments = -1;
static int last_peak_col = -1;
const int virtual_columns = 80;   // Each segment is devided by 5. For 1602A: 16x5=80.
float bar_level = 0;              // Do not change
const float bar_decay = 0.97;     // Decay rate (adjust in 0.95<bar_decay<0.99 range)
const float bar_threshold = 0.05; // Threshold when bars cleared out

const float power_scale_options[] = {100.0f, 500.0f, 1000.0f, 2000.0f}; // Available options (W)
byte current_power_scale_option = 0;    // Default to 100W (index 0)
float bar_graph = power_scale_options[current_power_scale_option] * 0.6f; // Calculated scale (60 for 100W)
float bar_graph_mult = 1.0f / bar_graph; // Pre-calculated multiplier

const int change_threshold = 1; //Minimum number of segment changes required before the lcd display is updated
// Peak Logic
float peak_bar_level = 0;             // Peak holding bar level
float peak_bar_position = 0;          // Tracks the highest point reached
const float peak_bar_decay = 0.995;    // Slower decay for peak bar
unsigned long last_peak_update = 0;    // When peak was last updated
const unsigned long peak_hold_duration = 1000; // Hold peak time before decay

// Custom characters
// Peak Bar
byte Pbar = 0x01;
byte bars[6][8] = {
    {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, Pbar, 0x00},  // bar0
    {0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x00},  // bar1
    {0x18, 0x18, 0x18, 0x18, 0x18, 0x18, 0x18, 0x00},  // bar2
    {0x1C, 0x1C, 0x1C, 0x1C, 0x1C, 0x1C, 0x1C, 0x00},  // bar3
    {0x1E, 0x1E, 0x1E, 0x1E, 0x1E, 0x1E, 0x1E, 0x00},  // bar4
    {0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x00}   // bar5
};  

byte peakMarker[8] = {Pbar, Pbar, Pbar, Pbar, Pbar, Pbar, Pbar, 0x00}; // A small horizontal line

// Band calibration
typedef struct {
  const char* name;
  float fwdA;
  float fwdB;
  float refA;
  float refB;
  float peakA;
  float peakB;
} BandInfo;

BandInfo bands[8] = {
  {"Def", 0.042766, 42.23601, 0.042766, 42.23601, 0.042766, 42.23601},
  {"160m", 0.042766, 42.23601, 0.042766, 42.23601, 0.042766, 42.23601},
  {"80m", 0.042766, 42.23601, 0.042766, 42.23601, 0.042766, 42.23601},
  {"60-40m", 0.042766, 42.23601, 0.042766, 42.23601, 0.042766, 42.23601},
  {"30-20m", 0.042766, 42.23601, 0.042766, 42.23601, 0.042766, 42.23601},
  {"17-15m", 0.042766, 42.23601, 0.042766, 42.23601, 0.042766, 42.23601},
  {"12-10m", 0.042766, 42.23601, 0.042766, 42.23601, 0.042766, 42.23601},
  {"6m", 0.042766, 42.23601, 0.042766, 42.23601, 0.042766, 42.23601}
};

int currentBand = 0;
bool menuMode = false;
byte menuState = 0; // 0=main, 1=band select, 2=display mode select, 3=hold time, 4=power scale
unsigned long lastButtonPress = 0;
const unsigned long debounceDelay = 200;
unsigned long menuStartTime = 0;
const int menu_timeout = 10000; // Menu Timeout in milliseconds (5 seconds)

float computePower(int adcValue, bool isForward) {
  if (adcValue < 5) return 0.0f;  // Deadzone for noise
  float voltage = adcValue * adcToVoltage;
  float slope = isForward ? bands[currentBand].fwdA : bands[currentBand].refA;
  float intercept = isForward ? bands[currentBand].fwdB : bands[currentBand].refB;
  float powerdBm = (slope * voltage) - intercept;
  return pow(10, (powerdBm - 30) * 0.1f);
}

float computePeakPower(int adcValue) {
  if (adcValue < 5) return 0.0f;  // Deadzone for noise
  float voltage = adcValue * adcToVoltage / 1000.0f; // Convert mV to V
  return 17.64f * voltage * voltage; // P = 17.64 * V^2
}

void setup() {
  lcd.begin(16, 2);
  // Initialize all custom characters (0 to 5)
  for (int i = 0; i < 6; i++) {
      lcd.createChar(i, bars[i]);
  }
  lcd.createChar(6, peakMarker);

  pinMode(buttonMenuPin, INPUT_PULLUP);
  pinMode(buttonSelectPin, INPUT_PULLUP);

  lcd.setCursor(0, 0);
  lcd.print("Po:0W");
  lcd.setCursor(10, 0);
  lcd.print("SW:1.00");

  analogReference(DEFAULT);
  ADCSRA = (ADCSRA & 0xF8) | 0x02;
}

void loop() {
  handleButtons();

  if (!menuMode) {
    readAndUpdatePower();
    calculateSWR();
    updateBarGraph();
    updateDisplay();
  }
}

int readAverage(byte pin, int samples = 100){
  long sum = 0;
  for (int i = 0; i < samples; i++){
    sum += analogRead(pin);
  }
  return sum / samples;
}

void readAndUpdatePower() {
  adcValueF = readAverage(analogPinF);
  adcValueR = readAverage(analogPinR);
  adcValuePeak = readAverage(analogPinPeak);

  pow_fwd = computePower(adcValueF, true);
  pow_ref = computePower(adcValueR, false);
  pow_peak = computePeakPower(adcValuePeak);

  // Always update both peaks for SWR calculation
  if (pow_fwd > peak_power) {
    peak_power = pow_fwd;
    last_peak_time = millis();
  } else if (millis() - last_peak_time > peak_hold_time) {
    peak_power = pow_fwd;
  }

  if (pow_ref > peak_power_ref) {
    peak_power_ref = pow_ref;
    last_swr_peak_reset = millis();
  }

  // For display purposes, track the selected power type
  float display_power = 0;
  switch(currentDisplayMode) {
    case DISPLAY_FWD_POWER:
      display_power = pow_fwd;
      break;
    case DISPLAY_REF_POWER:
      display_power = pow_ref;
      break;
    case DISPLAY_PEAK_POWER:
      display_power = pow_peak;
      break;
  }
  
  if (display_power > peak_power) {
    peak_power = display_power;
    last_peak_time = millis();
  } else if (millis() - last_peak_time > peak_hold_time) {
    peak_power = display_power;
  }
}

void calculateSWR() {
  // Update peak reflected power if current reading is higher
  if (pow_ref > peak_power_ref) {
    peak_power_ref = pow_ref;
    last_gamma = -1.0f; // Force SWR recalculation
  }

  // Reset peak reflected power after hold time or if current reading is 0
  if (millis() - last_swr_peak_reset > peak_hold_time || pow_ref == 0) {
    peak_power_ref = pow_ref;
    last_swr_peak_reset = millis();
    last_gamma = -1.0f; // Force SWR recalculation
  }

  // Only calculate SWR if we have sufficient forward power
  if (peak_power > 0.1f) {  // At least 0.1W forward power
    if (last_gamma < 0 || peak_power != last_peak_power) {
      // Recalculate gamma (reflection coefficient)
      if (peak_power_ref > 0) {
        float ratio = peak_power_ref / peak_power;
        last_gamma = sqrt(ratio);
      } else {
        last_gamma = 0.0f; // No reflected power = perfect match
      }
      last_peak_power = peak_power;
    }

    // Calculate SWR from gamma
    if (last_gamma >= 0 && last_gamma < 1.0f) {
      swr = (1.0f + last_gamma) / (1.0f - last_gamma);
    } else {
      swr = 99.9f; // Handle invalid cases
    }
  } else {
    // Insufficient forward power - assume perfect match
    swr = 1.0f;
    last_gamma = 0.0f;
  }

  // Ensure SWR is within valid range
  if (swr < 1.0f) swr = 1.0f;
  if (swr > 99.9f) swr = 99.9f;
}

void updateDisplay() {
  lcd.setCursor(0, 0);
  
  // Display the appropriate power label
  float display_power = 0;
  static float last_stable_power = 0; // Stores the last "stable" power value
  
  switch(currentDisplayMode) {
    case DISPLAY_FWD_POWER:
      lcd.print("Po:");
      display_power = pow_fwd;
      break;
    case DISPLAY_REF_POWER:
      lcd.print("Pr:");
      display_power = pow_ref;
      break;
    case DISPLAY_PEAK_POWER:
      lcd.print("Pk:");
      display_power = pow_peak;
      break;
  }

  // Only update if power changed by ±1.0W or more
  if (abs(display_power - last_stable_power) >= 2.0f || display_power < 1.0f) {
    int display_value = int(display_power + 0.5);
    
    if (display_value < 1) {
      lcd.print("0W   ");
    } else {
      lcd.print(display_value);
      lcd.print("W   ");
    }
    last_stable_power = display_power; // Update the stable value
  }

  // SWR display (unchanged, as it already uses 2 decimal places)
  lcd.setCursor(9, 0);
  lcd.print("SW:");
  lcd.print(swr, 2);
}

void handleButtons() {
  if (millis() - lastButtonPress < debounceDelay) return;

  if (digitalRead(buttonMenuPin) == LOW) {
    lastButtonPress = millis();
    menuStartTime = millis();  // Reset timeout timer on ANY menu button press
    
    if (!menuMode) {
      // Enter menu mode - Band Selection
      menuMode = true;
      menuState = 1;
      lcd.clear();
      lcd.setCursor(0, 0);
      lcd.print("Select Band:");
      lcd.setCursor(0, 1);
      lcd.print(bands[currentBand].name);
    } 
    else {
      // Already in menu mode - cycle through menu states
      menuState++;
      
      if (menuState == 1) {
        // Band Selection
        lcd.clear();
        lcd.setCursor(0, 0);
        lcd.print("Select Band:");
        lcd.setCursor(0, 1);
        lcd.print(bands[currentBand].name);
      }
      else if (menuState == 2) {
        // Display Mode Selection
        lcd.clear();
        lcd.setCursor(0, 0);
        lcd.print("Display Mode:");
        lcd.setCursor(0, 1);
        switch(currentDisplayMode) {
          case DISPLAY_FWD_POWER: lcd.print("FWD Power"); break;
          case DISPLAY_REF_POWER: lcd.print("REF Power"); break;
          case DISPLAY_PEAK_POWER: lcd.print("PEAK Power"); break;
        }
      }
      else if (menuState == 3) {
        // Peak Hold Time Selection
        lcd.clear();
        lcd.setCursor(0, 0);
        lcd.print("Peak Hold Time:");
        lcd.setCursor(0, 1);
        lcd.print(peak_hold_time);
        lcd.print(" ms");
      }
      else if (menuState == 4) {
        // Power Scale Selection
        lcd.clear();
        lcd.setCursor(0, 0);
        lcd.print("Power Scale:");
        lcd.setCursor(0, 1);
        lcd.print(power_scale_options[current_power_scale_option]);
        lcd.print(" W");
      }
      else {
        // Exit menu (state 5 or higher)
        menuMode = false;
        lcd.clear();
        updateDisplay();
      }
    }
  }

  if (menuMode && digitalRead(buttonSelectPin) == LOW) {
    lastButtonPress = millis();
    menuStartTime = millis();  // Reset timeout timer
    
    if (menuState == 1) {
      // Band selection
      currentBand = (currentBand + 1) % 8;
      lcd.setCursor(0, 1);
      lcd.print("                ");
      lcd.setCursor(0, 1);
      lcd.print(bands[currentBand].name);
    }
    else if (menuState == 2) {
      // Display mode toggle
      currentDisplayMode = (currentDisplayMode + 1) % 3;
      lcd.setCursor(0, 1);
      lcd.print("                ");
      lcd.setCursor(0, 1);
      switch(currentDisplayMode) {
        case DISPLAY_FWD_POWER: lcd.print("FWD Power"); break;
        case DISPLAY_REF_POWER: lcd.print("REF Power"); break;
        case DISPLAY_PEAK_POWER: lcd.print("PEAK Power"); break;
      }
    }
    else if (menuState == 3) {
      // Peak hold time selection
      current_hold_option = (current_hold_option + 1) % 4;
      peak_hold_time = peak_hold_options[current_hold_option];
      lcd.setCursor(0, 1);
      lcd.print("                ");
      lcd.setCursor(0, 1);
      lcd.print(peak_hold_time);
      lcd.print(" ms");
    }
    else if (menuState == 4) {
      // Power scale selection
      current_power_scale_option = (current_power_scale_option + 1) % 4;
      bar_graph = power_scale_options[current_power_scale_option] * 0.6f; // Maintain 60% scale factor
      bar_graph_mult = 1.0f / bar_graph; // Update multiplier
      lcd.setCursor(0, 1);
      lcd.print("                ");
      lcd.setCursor(0, 1);
      lcd.print(power_scale_options[current_power_scale_option]);
      lcd.print(" W");
    }
  }

  // AUTO EXIT after 10 seconds of inactivity
  if (menuMode && (millis() - menuStartTime > menu_timeout)) {
    menuMode = false;
    lcd.clear();
    updateDisplay();
  }
}

void updateBarGraph() {
  // Update current bar level based on selected display mode
  float power = 0;
  switch(currentDisplayMode) {
    case DISPLAY_FWD_POWER:
      power = pow_fwd;
      break;
    case DISPLAY_REF_POWER:
      power = pow_ref;
      break;
    case DISPLAY_PEAK_POWER:
      power = pow_peak;
      break;
  }

  float target_level = power * bar_graph_mult;
  bar_level = (target_level > bar_level) ? target_level : bar_level * bar_decay;
  
  // Update peak position
  if (target_level > peak_bar_position) {
    peak_bar_position = min(target_level, 16.0);  // Ensure it doesn't go beyond 16
    last_peak_update = millis();
  } else if (millis() - last_peak_update > peak_hold_duration) {
    peak_bar_position *= peak_bar_decay;
    peak_bar_position = constrain(peak_bar_position, 0, 16);  // Ensure it doesn't fall below 0
  }
  
  // Clear display if below threshold
  if (bar_level < bar_threshold && peak_bar_position < bar_threshold) {
    if (last_lit_segments != 0 || last_peak_col != -1) {
      last_lit_segments = 0;
      last_peak_col = -1;
      lcd.setCursor(0, 1);
      lcd.print("                ");  // Clear the line
    }
    return;
  }

  // Calculate current and peak positions
  int lit_segments = round(bar_level * virtual_columns);
  lit_segments = constrain(lit_segments, 0, virtual_columns);
  
  int peak_segments = round(peak_bar_position * virtual_columns);
  peak_segments = constrain(peak_segments, 0, virtual_columns);
  
  // Calculate the column where the peak marker should be placed
  int peak_col = peak_segments / 5;  // Which character column contains the peak
  peak_col = constrain(peak_col, 0, 15);  // Ensure it doesn't go beyond 15 (max column for LCD)

  // Only redraw if significant changes
  if (abs(lit_segments - last_lit_segments) >= change_threshold || 
      peak_col != last_peak_col ||
      lit_segments == virtual_columns || 
      lit_segments == 0) {
    
    last_lit_segments = lit_segments;
    last_peak_col = peak_col;
    
    // First draw the normal bar graph
    lcd.setCursor(0, 1);
    for (int i = 0; i < 16; i++) {
      int segments = lit_segments - (i * 5);
      byte char_num = 0;
      if (segments > 0) {
        char_num = (segments >= 5) ? 5 : segments;
      }
      lcd.write(char_num);
    }
    
    // Now overlay the peak marker if visible
    if (peak_bar_position >= bar_threshold && peak_col < 16) {
      lcd.setCursor(peak_col, 1);
      lcd.write(6); // Write the peak marker character
    }
  }
}