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

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

#define VOLTAGE_PIN PA0
#define CURRENT_PIN PA1
#define BTN_RESET PB0
#define BTN_MODE PB1
#define BUZZER_PIN PB10

float batteryCapacity_mAh = 2000.0;
float chargeUsed_mAh = 0.0;
float socPercent = 100.0;
unsigned long lastSampleTime = 0;
const unsigned long sampleInterval_ms = 500;
float prevCurrent_mA = 0.0;

// --- Rest detection ---
bool isResting = false;
unsigned long restStartTime = 0;
const float REST_CURRENT_THRESHOLD_mA = 50.0;
const unsigned long REST_DURATION_ms = 5000;

// --- Self-calibrating OCV table ---
#define NUM_BINS 6
float binAvgSoc[NUM_BINS];
int binCount[NUM_BINS];
const float BIN_MIN_VOLTAGE = 3.0;
const float BIN_MAX_VOLTAGE = 4.2;
const int CONFIDENCE_THRESHOLD = 2; // Requires 2 rest samples to calibrate
bool restEventApplied = false;

// --- Novelty 3: Current sensor offset calibration ---
float currentOffset_mA = 0.0;

// --- Novelty 2: Confidence-timer prompting ---
unsigned long lastConfidentCorrection = 0;
const unsigned long CONFIDENCE_TIMEOUT_ms = 15000; // 15 seconds for live demo
bool promptBuzzerPlayed = false;

// Gentle, pleasant chime function (E5 -> G5 soft alert)
void playPleasantChime() {
  tone(BUZZER_PIN, 659); // Play E5
  delay(150);
  tone(BUZZER_PIN, 784); // Play G5
  delay(200);
  noTone(BUZZER_PIN);    // Turn off
}

float readVoltage() {
  int raw = analogRead(VOLTAGE_PIN);
  return (raw / 2048.0) * 4.2;
}

float readCurrent_mA() {
  if (digitalRead(BTN_MODE) == LOW) {
    return 0.0; // Force 0.0 mA when holding BTN_MODE
  }
  int raw = analogRead(CURRENT_PIN);
  float measured = ((raw / 4095.0) * 1000.0) - 500.0;
  return measured - currentOffset_mA; // Applies offset correction
}

bool checkRestState(float current_mA) {
  if (abs(current_mA) < REST_CURRENT_THRESHOLD_mA) {
    if (!isResting) {
      restStartTime = millis();
      isResting = true;
    }
    if (millis() - restStartTime >= REST_DURATION_ms) {
      return true;
    }
  } else {
    isResting = false;
  }
  return false;
}

int voltageToBinIndex(float voltage) {
  float range = BIN_MAX_VOLTAGE - BIN_MIN_VOLTAGE;
  float clamped = constrain(voltage, BIN_MIN_VOLTAGE, BIN_MAX_VOLTAGE);
  int index = (int)(((clamped - BIN_MIN_VOLTAGE) / range) * (NUM_BINS - 1));
  return constrain(index, 0, NUM_BINS - 1);
}

void updateOcvTable(float voltage, float currentSoc) {
  int idx = voltageToBinIndex(voltage);
  binAvgSoc[idx] += (currentSoc - binAvgSoc[idx]) / (binCount[idx] + 1);
  binCount[idx]++;
}

bool applyOcvCorrection(float voltage, float &soc) {
  int idx = voltageToBinIndex(voltage);
  if (binCount[idx] >= CONFIDENCE_THRESHOLD) {
    soc = (soc + binAvgSoc[idx]) / 2.0; // Gentle correction blend
    return true;
  }
  return false;
}

void setup() {
  Serial.begin(115200);
  Wire.begin();

  pinMode(BTN_RESET, INPUT_PULLUP);
  pinMode(BTN_MODE, INPUT_PULLUP);
  pinMode(BUZZER_PIN, OUTPUT);

  // Initialize OCV curve seed
  for (int i = 0; i < NUM_BINS; i++) {
    binAvgSoc[i] = map(i, 0, NUM_BINS - 1, 0, 100);
    binCount[i] = 0;
  }

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("OLED init failed");
    while (1);
  }

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.println("Battery SoC Estimator");
  display.println("Starting...");
  display.display();
  delay(1000);

  prevCurrent_mA = readCurrent_mA();
  lastConfidentCorrection = millis();
}

void loop() {
  unsigned long now = millis();

  if (now - lastSampleTime >= sampleInterval_ms) {
    float dt_hours = (now - lastSampleTime) / 3600000.0;
    lastSampleTime = now;

    float voltage = readVoltage();
    float current_mA = readCurrent_mA();

    float avgCurrent_mA = (prevCurrent_mA + current_mA) / 2.0;
    chargeUsed_mAh += (avgCurrent_mA * dt_hours);
    prevCurrent_mA = current_mA;

    socPercent = 100.0 - (chargeUsed_mAh / batteryCapacity_mAh) * 100.0;
    socPercent = constrain(socPercent, 0, 100);

    bool atRest = checkRestState(current_mA);
    static bool ocvApplied = false;

    if (atRest && !restEventApplied) {
      // Novelty 3: Zero-offset self-calibration
      float rawRestReading = ((analogRead(CURRENT_PIN) / 4095.0) * 1000.0) - 500.0;
      currentOffset_mA += (rawRestReading - currentOffset_mA) * 0.3; // Gentle blend update

      // Novelty 1: OCV table update and correction
      updateOcvTable(voltage, socPercent);
      ocvApplied = applyOcvCorrection(voltage, socPercent);
      if (ocvApplied) {
        lastConfidentCorrection = millis(); // Reset confidence timer
        promptBuzzerPlayed = false;
      }
      restEventApplied = true;
    } else if (!atRest) {
      restEventApplied = false;
      ocvApplied = false;
    }

    if (digitalRead(BTN_RESET) == LOW) {
      chargeUsed_mAh = 0;
      socPercent = 100.0;
    }

    // Novelty 2: Check confidence timer timeout
    bool needsPrompt = (millis() - lastConfidentCorrection) > CONFIDENCE_TIMEOUT_ms;
    
    // Play pleasant dual-tone chime once when rest is needed
    if (needsPrompt && !promptBuzzerPlayed) {
      playPleasantChime();
      promptBuzzerPlayed = true;
    }

    // --- Display Section ---
    display.clearDisplay();
    display.setCursor(0, 0);
    display.setTextSize(1);
    display.print("Voltage: ");
    display.print(voltage, 2);
    display.println(" V");

    display.print("Current: ");
    display.print(current_mA, 1);
    display.println(" mA");

    display.setCursor(0, 25);
    display.setTextSize(2);
    display.print("SoC: ");
    display.print(socPercent, 1);
    display.println("%");

    display.setTextSize(1);
    display.setCursor(0, 50);
    if (atRest) {
      display.print(ocvApplied ? "OCV corrected!" : "RESTING-learning");
    } else if (needsPrompt) {
      display.print("Please let battery rest");
    }

    display.display();

    Serial.print("V=");
    Serial.print(voltage);
    Serial.print(" I=");
    Serial.print(current_mA);
    Serial.print(" SoC=");
    Serial.println(socPercent);
  }
}