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

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 32
#define OLED_RESET -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

#define VOLTAGE_PIN A0
#define RED_LED_PIN 2
#define GREEN_LED_PIN 3

// Voltage calibration constants
#define VREF 4.2
#define V_MIN 3.0
#define V_MAX 4.2
#define ADC_MAX 1023.0

// Battery icon parameters
#define BATTERY_WIDTH 110
#define BATTERY_HEIGHT 26
#define BATTERY_X 0
#define BATTERY_Y 6
#define BATTERY_MARGIN 4
#define BATTERY_SEGMENT_HEIGHT 18

// LED threshold with hysteresis
#define LOW_BATTERY_THRESHOLD 20
#define HIGH_BATTERY_THRESHOLD 30

// Smoothing parameters
#define ALPHA 0.15  // Smoothing factor (lower = smoother)
float smoothedVoltage = 0;
bool lowBatteryState = false;

// Function prototypes
void drawBatteryIcon(int percentage);
int voltageToPercentage(float voltage);
float readSmoothedVoltage();

void setup() {
  Serial.begin(9600);
  pinMode(RED_LED_PIN, OUTPUT);
  pinMode(GREEN_LED_PIN, OUTPUT);
  
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("OLED initialization failed"));
    // Blink red LED to indicate error
    while(1) {
      digitalWrite(RED_LED_PIN, !digitalRead(RED_LED_PIN));
      delay(200);
    }
  }
  
  display.clearDisplay();
  display.display();
  delay(500);
  
  // Initialize smoothed voltage with first reading
  int sensorValue = analogRead(VOLTAGE_PIN);
  smoothedVoltage = sensorValue * (VREF / ADC_MAX);
  
  Serial.println(F("Battery Monitor Initialized"));
}

void loop() {
  // Read and smooth voltage
  float voltage = readSmoothedVoltage();
  
  // Convert to percentage using non-linear curve
  int percentage = voltageToPercentage(voltage);
  
  // Constrain to valid range
  percentage = constrain(percentage, 0, 100);
  
  // LED control with hysteresis to prevent flickering
  if (percentage <= LOW_BATTERY_THRESHOLD) {
    lowBatteryState = true;
  } else if (percentage >= HIGH_BATTERY_THRESHOLD) {
    lowBatteryState = false;
  }
  
  digitalWrite(RED_LED_PIN, lowBatteryState ? HIGH : LOW);
  digitalWrite(GREEN_LED_PIN, lowBatteryState ? LOW : HIGH);
  
  // Update display
  display.clearDisplay();
  drawBatteryIcon(percentage);
  
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setRotation(1);
  display.setCursor(6, 0);
  display.print(percentage);
  display.print("%");
  
  display.display();
  
  // Debug output
  Serial.print("Voltage: ");
  Serial.print(voltage, 2);
  Serial.print("V | Percentage: ");
  Serial.print(percentage);
  Serial.println("%");
  
  delay(1000); // Reduced update frequency for power saving
}

// Exponential moving average for voltage smoothing
float readSmoothedVoltage() {
  int sensorValue = analogRead(VOLTAGE_PIN);
  float instantVoltage = sensorValue * (VREF / ADC_MAX);
  smoothedVoltage = (ALPHA * instantVoltage) + ((1 - ALPHA) * smoothedVoltage);
  return smoothedVoltage;
}

// Non-linear voltage-to-percentage conversion for Li-ion batteries
int voltageToPercentage(float voltage) {
  // Lookup table based on typical Li-ion discharge curve
  const float voltageTable[] = {3.0, 3.3, 3.5, 3.6, 3.7, 3.8, 3.9, 4.0, 4.1, 4.2};
  const int percentTable[] = {0, 5, 10, 20, 35, 55, 75, 85, 95, 100};
  const int tableSize = 10;
  
  // Handle boundary conditions
  if (voltage <= voltageTable[0]) return 0;
  if (voltage >= voltageTable[tableSize - 1]) return 100;
  
  // Find the appropriate segment and interpolate
  for (int i = 0; i < tableSize - 1; i++) {
    if (voltage >= voltageTable[i] && voltage < voltageTable[i + 1]) {
      // Linear interpolation between two points
      float voltageRange = voltageTable[i + 1] - voltageTable[i];
      float percentRange = percentTable[i + 1] - percentTable[i];
      float voltageOffset = voltage - voltageTable[i];
      
      return percentTable[i] + (int)((voltageOffset / voltageRange) * percentRange);
    }
  }
  
  return 0;
}

// Draw battery icon with capacity bars
void drawBatteryIcon(int percentage) {
  display.setRotation(0);
  
  // Draw battery outline
  display.drawRect(BATTERY_X, BATTERY_Y, BATTERY_WIDTH, BATTERY_HEIGHT, SSD1306_WHITE);
  
  // Draw battery terminal (positive tip)
  display.fillRect(BATTERY_X + BATTERY_WIDTH, BATTERY_Y + 8, 3, 10, SSD1306_WHITE);
  
  // Calculate number of bars (10 segments for smooth visualization)
  int numBars = map(percentage, 0, 100, 0, 10);
  
  // Draw capacity bars
  int barWidth = (BATTERY_WIDTH - 2 * BATTERY_MARGIN) / 10 - 1;
  for (int i = 0; i < numBars; i++) {
    int barX = BATTERY_X + BATTERY_MARGIN + i * (barWidth + 1);
    int barY = BATTERY_Y + BATTERY_MARGIN;
    display.fillRect(barX, barY, barWidth, BATTERY_SEGMENT_HEIGHT, SSD1306_WHITE);
  }
}