#include <Wire.h>
#include <LiquidCrystal.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
#define buttonMenuPin  4
#define buttonSelectPin 2

// Power measurement variables
int Vreference = 5000; // mV
const float adcToMillivolts = Vreference / 1023.0f; // Now gives mV directly
int adcValueF = 0;
int adcValueR = 0; 
float pow_fwd = 0;
float pow_ref = 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;
const unsigned long peak_hold_time = 1000; // Power And SWR Hold time

// Power lookup table (millivolts, power)
typedef struct {
  int millivolts;
  float power;
} PowerLookup;

const PowerLookup powerTable[] = {
  {1500, 0.0},    // 0mV = 0W
  {1857, 5},
  {1911, 10},
  {1948, 20},
  {1961, 30},
  {1969, 40},
  {1975, 50},
  {1978, 60},
  {1980, 70},
  {1983, 80},
  {1984, 90},
  {1985, 100}
};
const int tableSize = sizeof(powerTable)/sizeof(powerTable[0]);

// Bar graph variables (unchanged)
static int last_lit_segments = -1;
static int last_peak_col = -1;
const int virtual_columns = 80;
float bar_level = 0;
const float bar_decay = 0.975;
const float bar_threshold = 0.05;
float bar_graph = 60;
float bar_graph_mult = 1.0f / bar_graph;
const int change_threshold = 1;

// Peak Logic (unchanged)
float peak_bar_level = 0;
float peak_bar_position = 0;
const float peak_bar_decay = 0.995;
unsigned long last_peak_update = 0;
const unsigned long peak_hold_duration = 2000;

// Custom characters (unchanged)
byte Pbar = 0x01;
byte bars[6][8] = {
    {0x00, 0x00, 0x00, 0x00, 0x00, 0x00, Pbar, 0x00},
    {0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x10, 0x00},
    {0x18, 0x18, 0x18, 0x18, 0x18, 0x18, 0x18, 0x00},
    {0x1C, 0x1C, 0x1C, 0x1C, 0x1C, 0x1C, 0x1C, 0x00},
    {0x1E, 0x1E, 0x1E, 0x1E, 0x1E, 0x1E, 0x1E, 0x00},
    {0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x1F, 0x00}
};
byte peakMarker[8] = {Pbar, Pbar, Pbar, Pbar, Pbar, Pbar, Pbar, 0x00};

// Band names only
const char* bands[8] = {
  "Def", "160m", "80m", "60-40m", 
  "30-20m", "17-15m", "12-10m", "6m"
};

int currentBand = 0;
bool menuMode = false;
bool showReflectedPower = false;
byte menuState = 0;
unsigned long lastButtonPress = 0;
const unsigned long debounceDelay = 200;
unsigned long menuStartTime = 0;
const int menu_timeout = 5000;

float lookupPower(int millivolts) {
  // Handle voltages below minimum (including 0mV)
  if (millivolts <= powerTable[0].millivolts) return 0.0f;
  
  // Handle voltages above maximum
  if (millivolts >= powerTable[tableSize-1].millivolts) 
    return powerTable[tableSize-1].power;
  
  // Find the interval where the voltage lies
  for (int i = 0; i < tableSize-1; i++) {
    if (millivolts >= powerTable[i].millivolts && 
        millivolts <= powerTable[i+1].millivolts) {
      // Linear interpolation
      float fraction = (float)(millivolts - powerTable[i].millivolts) / 
                      (powerTable[i+1].millivolts - powerTable[i].millivolts);
      return powerTable[i].power + fraction * 
             (powerTable[i+1].power - powerTable[i].power);
    }
  }
  return 0.0f;
}

float computePower(int adcValue) {
  if (adcValue < 5) return 0.0f;  // Deadzone for noise
  int millivolts = adcValue * adcToMillivolts;
  return lookupPower(millivolts);
}

void setup() {
  lcd.begin(16, 2);
  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(INTERNAL);
  ADCSRA = (ADCSRA & 0xF8) | 0x02;
  Serial.begin(9600);  // Initialize serial communication
Serial.println("Power Meter Debug - Starting...");
}

void debugVoltages() {
  static unsigned long lastDebugTime = 0;
  const unsigned long debugInterval = 1000; // Print every 1000ms
  
  if (millis() - lastDebugTime > debugInterval) {
    lastDebugTime = millis();
    
    // Take fresh readings
    int rawF = analogRead(analogPinF);
    int rawR = analogRead(analogPinR);
    int mV_F = rawF * (Vreference / 1023.0);
    int mV_R = rawR * (Vreference / 1023.0);
    
    Serial.print("Forward - Raw: ");
    Serial.print(rawF);
    Serial.print(" ADC (");
    Serial.print(mV_F);
    Serial.print(" mV) | Reflected - Raw: ");
    Serial.print(rawR);
    Serial.print(" ADC (");
    Serial.print(mV_R);
    Serial.println(" mV)");
  }
}

// ALL OTHER FUNCTIONS REMAIN EXACTLY THE SAME AS BEFORE
// (readAndUpdatePower, calculateSWR, updateDisplay, handleButtons, updateBarGraph, loop)

void readAndUpdatePower() {
    adcValueF = 0;
    adcValueR = 0;

    int adcValueFF = 0;
    int adcValueRR = 0;
    for(int i = 0; i <=20; i++) {
        adcValueFF = analogRead(analogPinF);
        adcValueRR = analogRead(analogPinR);
        adcValueF = max(adcValueFF , adcValueF);
        adcValueR = max(adcValueRR , adcValueR);
    }
  
    pow_fwd = computePower(adcValueF);
    pow_ref = computePower(adcValueR);

    // 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 either forward or reflected power
    float display_power = showReflectedPower ? pow_ref : pow_fwd;
    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);
  lcd.print(showReflectedPower ? "Pr:" : "Po:");
  
  int display_power = int(peak_power + 0.5);
  if (display_power < 1) {
    lcd.print("0W   ");
  } else {
    lcd.print(display_power);
    lcd.print("W   ");
  }

  lcd.setCursor(9, 0);
  lcd.print("SW:");
  lcd.print(swr, 2);
}

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

  if (digitalRead(buttonMenuPin) == LOW) {
    lastButtonPress = millis();
    
    if (!menuMode) {
      // Enter menu mode - Band Selection
      menuMode = true;
      menuState = 1;
      menuStartTime = millis();
      lcd.clear();
      lcd.setCursor(0, 0);
      lcd.print("Select Band:");
      lcd.setCursor(0, 1);
      lcd.print(bands[currentBand]);
    } 
    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]);
      }
      else if (menuState == 2) {
        // FWD/REF Selection
        lcd.clear();
        lcd.setCursor(0, 0);
        lcd.print("Select Mode:");
        lcd.setCursor(0, 1);
        lcd.print(showReflectedPower ? "REF Power" : "FWD Power");
      }
      else {
        // Exit menu (state 3 or higher)
        menuMode = false;
        lcd.clear();
        lcd.setCursor(0, 0);
        lcd.print(showReflectedPower ? "Pr:0W" : "Po:0W");
        lcd.setCursor(9, 0);
        lcd.print("SW:");
        lcd.print(swr, 2);
      }
    }
  }

  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]);
    }
    else if (menuState == 2) {
      // FWD/REF toggle
      showReflectedPower = !showReflectedPower;
      lcd.setCursor(0, 1);
      lcd.print("                ");
      lcd.setCursor(0, 1);
      lcd.print(showReflectedPower ? "REF Power" : "FWD Power");
    }
  }

  // AUTO EXIT after 5 seconds of inactivity
  if (menuMode && (millis() - menuStartTime > menu_timeout)) {
    menuMode = false;
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print(showReflectedPower ? "Pr:0W" : "Po:0W");
    lcd.setCursor(9, 0);
    lcd.print("SW:");
    lcd.print(swr, 2);
  }
}

void updateBarGraph() {
  // Update current bar level
  float power = showReflectedPower ? pow_ref : pow_fwd;

  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
    }
  }
}

void loop() {
  handleButtons();
    debugVoltages();  // Add this line

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