#include <avr/wdt.h>
#include <EEPROM.h>
#include <math.h>
#include <TM1637Display.h>

#define FIRMWARE_VERSION "v8.1 – Fixed Production Final"

#define SENSOR_TYPE 20
#if SENSOR_TYPE == 20
  #define ACS_SENS          0.100f
  #define INITIAL_NO_LOAD_A 1.28f
#endif

// ================= KONSTANTA =================
const uint32_t RMS_WINDOW_US          = 20000UL;
const uint32_t SAMPLE_INTERVAL_US     = 800UL;

const uint32_t STARTUP_IGNORE_MS      = 4500UL;
const uint32_t NOLOAD_ENTER_HYST      = 5200UL;
const uint32_t NOLOAD_EXIT_HYST       = 3200UL;
const uint32_t INPUT_DEBOUNCE_MS      = 2500UL;

const uint32_t OC_HYST_TIME           = 900UL;
const uint32_t STALL_HYST_TIME        = 3000UL;
const uint32_t START_DELAY_MS         = 1500UL;
const uint32_t MIN_RUN_TIME_MS        = 60000UL;
const uint32_t MIN_OFF_TIME_MS        = 10000UL;
const uint32_t LEARNING_MIN_STABLE_MS = 25000UL;

const uint32_t MAX_MUTE_DURATION_MS   = 300000UL;   // 5 menit maksimal mute

const float NO_LOAD_ENTER_RATIO       = 0.42f;
const float NO_LOAD_EXIT_RATIO        = 0.55f;
const float STALL_RATIO               = 1.70f;
const float OC_RATIO                  = 2.90f;

const float LEARNING_ADAPT_RATE       = 0.022f;
const float BASELINE_FIRST_LEARN      = 0.85f;
const float MIN_LEARNING_CURRENT      = 1.0f;
const float BASELINE_MAX_DELTA        = 2.8f;

// EEPROM
const uint16_t EEPROM_SIG_ADDR        = 0;
const uint16_t EEPROM_VER_ADDR        = 4;
const uint16_t EEPROM_BASELINE_ADDR   = 8;
const uint16_t EEPROM_LEARNED_ADDR    = 12;

const uint32_t EEPROM_SIGNATURE       = 0xDEADBEEF;
const uint8_t  EEPROM_VERSION         = 5;
const uint32_t EEPROM_WRITE_INTERVAL  = 60000UL;

// Calibration
const uint16_t CAL_SAMPLES            = 1100;        // ← FIXED
const uint8_t  CAL_TRIM_COUNT         = 70;

// ================= PINOUT =================
const uint8_t IO_S1_TANK    = 2;
const uint8_t IO_S2_DRY     = 4;
const uint8_t IO_RESET_BTN  = A0;
const uint8_t IO_CT_CURRENT = A6;
const uint8_t IO_R1_DRY     = 9;
const uint8_t IO_R2_PUMP    = 10;
const uint8_t IO_LED_PUMP   = 5;
const uint8_t IO_LED_SAFE   = 6;
const uint8_t IO_LED_STATUS = 7;
const uint8_t IO_BUZZER     = 3;

const uint8_t TM_CLK = 11;
const uint8_t TM_DIO = 12;

TM1637Display display(TM_CLK, TM_DIO);

#define LED_ON(p)  digitalWrite((p), LOW)
#define LED_OFF(p) digitalWrite((p), HIGH)

// ================= STATE =================
enum PumpState {
  ST_INIT, ST_IDLE, ST_REQUEST, ST_STARTING, ST_RUNNING,
  ST_DRY_FAULT, ST_NOLOAD_FAULT, ST_OVERCURRENT_FAULT,
  ST_STALL_FAULT, ST_LATCHED
};

PumpState state = ST_INIT;

// ================= SENSOR DATA =================
struct SensorData {
  float offset = 2.5f;
  float filteredCurrent = 0.0f;
  float baselineCurrent = 0.0f;
  
  bool tankNeedsWater = false;
  bool wellHasWater = true;
  bool noLoadDetected = false;
  
  bool learned = false;

  unsigned long tNoLoadEnter = 0;
  unsigned long tNoLoadExit  = 0;
  unsigned long pumpStartTime = 0;
  unsigned long lastStopTime = 0;
  unsigned long startDelayTimer = 0;
  unsigned long inputDebounceTimer = 0;
} sensor;

struct RMSEngine {
  float sumSq = 0.0f;
  uint32_t count = 0;
  unsigned long windowStart = 0;
  unsigned long lastSampleTime = 0;
} rms;

// ================= DISPLAY =================
float displayValue = 0.0f;
int lastDisplayVal = -1;
unsigned long hmiTimer = 0;

// Timers
unsigned long scanTimer = 0;
unsigned long ocTimer = 0;
unsigned long stallTimer = 0;
unsigned long lastEEPROMWrite = 0;

bool resetPressed = false;
unsigned long resetPressStart = 0;
bool buzzerMuted = false;
unsigned long muteStartTime = 0;

unsigned long ledPumpTimer = 0, ledStatusTimer = 0;
bool ledPumpState = false, ledStatusState = false;

// ================= EEPROM =================
void saveToEEPROM() {
  if (millis() - lastEEPROMWrite < EEPROM_WRITE_INTERVAL) return;

  EEPROM.put(EEPROM_SIG_ADDR, EEPROM_SIGNATURE);
  EEPROM.put(EEPROM_VER_ADDR, EEPROM_VERSION);
  EEPROM.put(EEPROM_BASELINE_ADDR, sensor.baselineCurrent);
  EEPROM.put(EEPROM_LEARNED_ADDR, sensor.learned);

  lastEEPROMWrite = millis();
}

bool loadFromEEPROM() {
  uint32_t sig;
  uint8_t ver;
  float base;
  bool learnedFlag;

  EEPROM.get(EEPROM_SIG_ADDR, sig);
  EEPROM.get(EEPROM_VER_ADDR, ver);
  EEPROM.get(EEPROM_BASELINE_ADDR, base);
  EEPROM.get(EEPROM_LEARNED_ADDR, learnedFlag);

  if (sig == EEPROM_SIGNATURE && ver == EEPROM_VERSION &&
      base >= 0.6f && base <= 55.0f) {
    sensor.baselineCurrent = base;
    sensor.learned = learnedFlag;
    return true;
  }
  return false;
}

void resetEEPROM() {
  EEPROM.put(EEPROM_SIG_ADDR, (uint32_t)0);
  sensor.learned = false;
  sensor.baselineCurrent = 0.0f;
}

// ================= SETUP =================
void setup() {
  wdt_disable();
  wdt_enable(WDTO_4S);
  analogReference(DEFAULT);

  pinMode(IO_R2_PUMP, OUTPUT);
  pinMode(IO_R1_DRY, OUTPUT);
  digitalWrite(IO_R2_PUMP, LOW);
  digitalWrite(IO_R1_DRY, LOW);

  display.setBrightness(7);
  display.clear();

  pinMode(IO_S1_TANK, INPUT_PULLUP);
  pinMode(IO_S2_DRY, INPUT_PULLUP);
  pinMode(IO_RESET_BTN, INPUT_PULLUP);
  pinMode(IO_CT_CURRENT, INPUT);
  pinMode(IO_LED_PUMP, OUTPUT);
  pinMode(IO_LED_SAFE, OUTPUT);
  pinMode(IO_LED_STATUS, OUTPUT);
  pinMode(IO_BUZZER, OUTPUT);

  loadFromEEPROM();

  // Calibration
  uint8_t calSeg[4] = {0x39, 0x77, 0x38, 0x00}; // CAL
  display.setSegments(calSeg);

  float sum = 0.0f;
  float minVal = 1023.0f, maxVal = 0.0f;
  uint16_t valid = 0;

  for (uint16_t i = 0; i < CAL_SAMPLES; i++) {
    uint16_t adc = analogRead(IO_CT_CURRENT);
    if (i >= 100) {
      sum += adc;
      valid++;
      if (adc < minVal) minVal = adc;
      if (adc > maxVal) maxVal = adc;
    }
    delay(2);
    wdt_reset();
  }

  float trimmed = sum - minVal - maxVal;
  sensor.offset = (trimmed / (valid - 2)) * 5.0f / 1023.0f;

  unsigned long nowMicro = micros();
  rms.windowStart = nowMicro;
  rms.lastSampleTime = nowMicro;

  sensor.lastStopTime = millis() - MIN_RUN_TIME_MS;
  lastEEPROMWrite = millis();

  uint8_t ready[4] = {0x50, 0x79, 0x77, 0x5E};
  display.setSegments(ready);
  delay(1100);
  display.clear();

  if (sensor.learned) {
    display.showNumberDecEx((int)(sensor.baselineCurrent * 10 + 0.5f), 0x40, true, 4, 0);
    delay(1400);
    display.clear();
  }

  state = ST_INIT;
}

// ================= RMS =================
void updateRMS() {
  unsigned long now = micros();
  if ((long)(now - rms.lastSampleTime) < 0) rms.lastSampleTime = now;

  while ((long)(now - rms.lastSampleTime) >= SAMPLE_INTERVAL_US) {
    rms.lastSampleTime += SAMPLE_INTERVAL_US;

    float curr = readCurrentRaw();
    static float lastValid = 0.0f;
    if (fabs(curr - lastValid) > 9.8f) curr = lastValid;
    else lastValid = curr;

    rms.sumSq += curr * curr;
    rms.count++;
  }

  if ((long)(now - rms.windowStart) >= RMS_WINDOW_US) {
    float rmsRaw = (rms.count > 12) ? sqrtf(rms.sumSq / rms.count) : sensor.filteredCurrent;

    if (rmsRaw < 0.12f) rmsRaw = 0.0f;
    if (rmsRaw > 60.0f) rmsRaw = 60.0f;

    bool isStartup = (state == ST_STARTING) || (millis() - sensor.pumpStartTime < 8000UL);
    float alpha = isStartup ? 0.32f : 0.12f;

    sensor.filteredCurrent = alpha * rmsRaw + (1.0f - alpha) * sensor.filteredCurrent;

    rms.windowStart += RMS_WINDOW_US;
    rms.sumSq = 0.0f;
    rms.count = 0;
  }
}

float readCurrentRaw() {
  uint16_t adc = analogRead(IO_CT_CURRENT);
  float v = adc * 5.0f / 1023.0f;
  float curr = (v - sensor.offset) / ACS_SENS;
  return fabs(curr) < 0.08f ? 0.0f : curr;
}

// ================= INPUTS =================
void readInputs() {
  unsigned long now = millis();

  // Float switch debounce
  static bool lastTank = false;
  bool currentTank = (digitalRead(IO_S1_TANK) == LOW);

  if (currentTank != lastTank) {
    sensor.inputDebounceTimer = now;
  }
  if (now - sensor.inputDebounceTimer >= INPUT_DEBOUNCE_MS) {
    sensor.tankNeedsWater = currentTank;
  }
  lastTank = currentTank;

  sensor.wellHasWater = (digitalRead(IO_S2_DRY) == HIGH);

  // Reset Button
  bool btn = (digitalRead(IO_RESET_BTN) == LOW);
  if (btn && !resetPressed) {
    resetPressed = true;
    resetPressStart = now;
  } else if (!btn && resetPressed) {
    unsigned long dur = now - resetPressStart;
    resetPressed = false;

    if (dur >= 3000) {
      resetEEPROM();
      tone(IO_BUZZER, 2500, 200);
      if (state == ST_LATCHED) state = ST_IDLE;
    } else if (dur >= 800) {
      buzzerMuted = !buzzerMuted;
      if (buzzerMuted) muteStartTime = now;
      tone(IO_BUZZER, 1200, 100);
    }
  }

  if (buzzerMuted && (now - muteStartTime >= MAX_MUTE_DURATION_MS)) buzzerMuted = false;

  // No-Load Detection
  if (state == ST_RUNNING && (now - sensor.pumpStartTime > STARTUP_IGNORE_MS)) {
    float enterThresh = sensor.learned ? sensor.baselineCurrent * NO_LOAD_ENTER_RATIO : INITIAL_NO_LOAD_A;
    float exitThresh  = sensor.learned ? sensor.baselineCurrent * NO_LOAD_EXIT_RATIO  : INITIAL_NO_LOAD_A * 1.25f;

    if (!sensor.noLoadDetected) {
      if (sensor.filteredCurrent < enterThresh) {
        if (sensor.tNoLoadEnter == 0) sensor.tNoLoadEnter = now;
        else if (now - sensor.tNoLoadEnter >= NOLOAD_ENTER_HYST)
          sensor.noLoadDetected = true;
      } else {
        sensor.tNoLoadEnter = 0;
      }
    } else {
      if (sensor.filteredCurrent > exitThresh) {
        if (sensor.tNoLoadExit == 0) sensor.tNoLoadExit = now;
        else if (now - sensor.tNoLoadExit >= NOLOAD_EXIT_HYST) {
          sensor.noLoadDetected = false;
          sensor.tNoLoadEnter = 0;
          sensor.tNoLoadExit = 0;
        }
      } else {
        sensor.tNoLoadExit = 0;
      }
    }
  } else {
    sensor.noLoadDetected = false;
    sensor.tNoLoadEnter = 0;
    sensor.tNoLoadExit = 0;
  }
}

// ================= FSM =================
void fsm() {
  unsigned long now = millis();

  switch (state) {
    case ST_INIT:
      state = ST_IDLE;
      break;

    case ST_IDLE:
      if ((now - sensor.lastStopTime >= MIN_OFF_TIME_MS) &&
          sensor.tankNeedsWater && sensor.wellHasWater) {
        state = ST_REQUEST;
      }
      break;

    case ST_REQUEST:
      state = ST_STARTING;
      sensor.startDelayTimer = now;
      break;

    case ST_STARTING:
      if (now - sensor.startDelayTimer >= START_DELAY_MS) {
        state = ST_RUNNING;
        sensor.pumpStartTime = now;
        sensor.tNoLoadEnter = 0;
        sensor.tNoLoadExit = 0;
        ocTimer = stallTimer = 0;
      }
      break;

    case ST_RUNNING:
      if (now - sensor.pumpStartTime < STARTUP_IGNORE_MS) break;

      // Relay Stuck Detection (Critical Safety)
      if (digitalRead(IO_R2_PUMP) == LOW && sensor.filteredCurrent > 0.8f) {
        state = ST_LATCHED;   // Relay weld → emergency latch
      }

      if (!sensor.wellHasWater) {
        state = ST_DRY_FAULT;
      }
      else if (sensor.noLoadDetected) {
        state = ST_NOLOAD_FAULT;
      }
      else if (sensor.learned && sensor.filteredCurrent > sensor.baselineCurrent * OC_RATIO) {
        if (ocTimer == 0) ocTimer = now;
        else if (now - ocTimer >= OC_HYST_TIME) state = ST_OVERCURRENT_FAULT;
      }
      else if (sensor.learned && sensor.filteredCurrent > sensor.baselineCurrent * STALL_RATIO) {
        ocTimer = 0;
        if (stallTimer == 0) stallTimer = now;
        else if (now - stallTimer >= STALL_HYST_TIME) state = ST_STALL_FAULT;
      }
      else {
        ocTimer = stallTimer = 0;

        if (!sensor.tankNeedsWater && (now - sensor.pumpStartTime >= MIN_RUN_TIME_MS)) {
          state = ST_IDLE;
          sensor.lastStopTime = now;
          saveToEEPROM();
        }
        else if (!sensor.learned && (now - sensor.pumpStartTime > LEARNING_MIN_STABLE_MS)) {
          if (sensor.filteredCurrent > MIN_LEARNING_CURRENT) {
            sensor.baselineCurrent = sensor.filteredCurrent * BASELINE_FIRST_LEARN;
            sensor.learned = true;
          }
        }
        else if (sensor.learned) {
          float error = sensor.filteredCurrent - sensor.baselineCurrent;
          if (fabs(error) < BASELINE_MAX_DELTA) {
            sensor.baselineCurrent += error * LEARNING_ADAPT_RATE;
          }
        }
      }
      break;

    case ST_DRY_FAULT:
      if (sensor.wellHasWater) state = ST_IDLE;
      break;

    case ST_NOLOAD_FAULT:
    case ST_OVERCURRENT_FAULT:
    case ST_STALL_FAULT:
      state = ST_LATCHED;
      break;

    case ST_LATCHED:
      break;
  }
}

// ================= OUTPUTS, LED, BUZZER, HMI =================
void outputs() {
  bool pumpCmd = (state == ST_RUNNING || state == ST_STARTING);
  bool safe = sensor.wellHasWater && state != ST_LATCHED;
  digitalWrite(IO_R2_PUMP, (pumpCmd && safe) ? HIGH : LOW);
  digitalWrite(IO_R1_DRY, (!sensor.wellHasWater || state == ST_DRY_FAULT) ? HIGH : LOW);
}

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

  if (state == ST_RUNNING) {
    if (now - ledPumpTimer >= 480) {
      ledPumpTimer = now;
      ledPumpState = !ledPumpState;
    }
    ledPumpState ? LED_ON(IO_LED_PUMP) : LED_OFF(IO_LED_PUMP);
  } else LED_OFF(IO_LED_PUMP);

  if (state == ST_IDLE && sensor.wellHasWater) LED_ON(IO_LED_SAFE);
  else LED_OFF(IO_LED_SAFE);

  bool faultOrDry = (state >= ST_DRY_FAULT) || (state == ST_IDLE && !sensor.wellHasWater);
  if (faultOrDry) {
    uint16_t interval = (state == ST_IDLE && !sensor.wellHasWater) ? 140 : 750;
    if (now - ledStatusTimer >= interval) {
      ledStatusTimer = now;
      ledStatusState = !ledStatusState;
    }
    ledStatusState ? LED_ON(IO_LED_STATUS) : LED_OFF(IO_LED_STATUS);
  } else LED_OFF(IO_LED_STATUS);
}

void updateBuzzer() {
  unsigned long now = millis();
  bool inFault = (state >= ST_DRY_FAULT && state <= ST_LATCHED);

  // Hard override untuk LATCHED state
  if (state == ST_LATCHED) {
    digitalWrite(IO_BUZZER, (now % 4600 < 1600) ? HIGH : LOW);
    return;
  }

  if (buzzerMuted || !inFault) {
    digitalWrite(IO_BUZZER, LOW);
    return;
  }

  digitalWrite(IO_BUZZER, (now % 580 < 190) ? HIGH : LOW);
}

void updateHMI() {
  unsigned long now = millis();
  if (now - hmiTimer < 120) return;
  hmiTimer = now;

  if (state == ST_NOLOAD_FAULT) {
    uint8_t seg[4] = {0x54, 0x38, 0x5E, 0x00};
    display.setSegments(seg);
  }
  else if (state == ST_DRY_FAULT) {
    uint8_t seg[4] = {0x5E, 0x50, 0x71, 0x00};
    display.setSegments(seg);
  }
  else if (state == ST_LATCHED) {
    uint8_t seg[4] = {0x38, 0x77, 0x77, 0x00};
    display.setSegments(seg);
  }
  else if (state == ST_RUNNING) {
    displayValue = 0.7f * displayValue + 0.3f * sensor.filteredCurrent;
    int val = (int)(displayValue * 10.0f + 0.5f);
    display.showNumberDecEx(val, 0x40, true, 4, 0);
  }
  else if (state == ST_IDLE && sensor.learned) {
    int val = (int)(sensor.baselineCurrent * 10.0f + 0.5f);
    display.showNumberDecEx(val, 0x40, true, 4, 0);
  }
}

// ================= MAIN =================
void loop() {
  wdt_reset();
  updateRMS();

  if (millis() - scanTimer >= 50) {
    scanTimer = millis();
    readInputs();
    fsm();
    outputs();
    updateHMI();
    updateLEDs();
  }
  updateBuzzer();
}