// // Nilm Whole-home Energy Metering Firmware /

// ============================================================
//  NILM Whole-Home Energy Metering Firmware
//  Board   : ESP32 DevKit v1
//  Sensors : PZEM-004T v3.0 (main) + optional 2nd PZEM
//            ADS1115 (optional high-rate SCT-013 sampling)
//  Output  : MQTT JSON  →  local MQTT broker (e.g. Mosquitto)
//            or HTTP POST → InfluxDB / CSV endpoint
//  Display : SSD1306 128×64 OLED (live readout)
//  Grid    : 230 V / 50 Hz (Nepal NEC standard)
// ============================================================

#include <Arduino.h>
#include <Wire.h>
#include <WiFi.h>
#include <HardwareSerial.h>
#include <PZEM004Tv30.h>
#include <Adafruit_ADS1X15.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <PubSubClient.h>
#include <ArduinoJson.h>

// ── Pin Definitions ────────────────────────────────────────
#define PZEM_RX       16   // ESP32 RX2 ← PZEM TX (via level shifter A2 / B2)
#define PZEM_TX       17   // ESP32 TX2 → PZEM RX (via level shifter A1 / B1)
// I2C: SDA=GPIO21, SCL=GPIO22 (ESP32 defaults, shared by ADS1115 + OLED)

// ── WiFi Credentials ───────────────────────────────────────
#define WIFI_SSID     "YOUR_SSID"
#define WIFI_PASSWORD "YOUR_PASSWORD"

// ── MQTT Broker ────────────────────────────────────────────
#define MQTT_HOST     "192.168.1.100"   // IP of your local Mosquitto broker
#define MQTT_PORT     1883
#define MQTT_CLIENT   "nilm_meter_01"
#define MQTT_TOPIC_MAIN   "nilm/main/power"
#define MQTT_TOPIC_SUB    "nilm/sub/power"
#define MQTT_TOPIC_ADS    "nilm/ads/current"

// ── Sampling Intervals ─────────────────────────────────────
#define PZEM_INTERVAL_MS    2000   // PZEM poll: ~2 s (PZEM internal settling ≥1.2 s)
#define ADS_INTERVAL_MS     100    // ADS1115 high-rate sampling: 100 ms (10 Hz)
#define OLED_INTERVAL_MS    1000   // OLED refresh: 1 s

// ── OLED ───────────────────────────────────────────────────
#define SCREEN_W    128
#define SCREEN_H     64
#define OLED_ADDR   0x3C

// ── ADS1115 ────────────────────────────────────────────────
#define ADS_ADDR    0x48   // ADDR pin → GND
// SCT-013 burden resistor 33 Ω + 2×10 kΩ voltage divider bias at 1.65 V
// ADS1115 ±2.048 V range (PGA = 2), 16-bit → 0.0625 mV/LSB
// Current scale: SCT-013-030 = 30A/1V ↔ 30A over 1V across 33Ω? Actually
// SCT-013-000 (0-100A, 50mA secondary) → I_sec = I_pri / 2000
// V_burden = I_sec × 33Ω  → I_pri = V_ads / 33 × 2000   (adjust to your CT ratio)
#define SCT_TURNS_RATIO   2000    // SCT-013-000: 100A primary / 50mA secondary
#define SCT_BURDEN_OHM    33.0f
#define ADS_PGA_SCALE     0.0000625f  // V per LSB at PGA=±2.048V (2.048/32768)

// ── PZEM Modbus addresses ──────────────────────────────────
#define PZEM_ADDR_MAIN  0xF8   // default broadcast (single device)
#define PZEM_ADDR_SUB   0x02   // set via PZEM-004T-v30 SetAddress sketch first

// ══════════════════════════════════════════════════════════
//  Globals
// ══════════════════════════════════════════════════════════

struct PzemData {
    float voltage   = 0;
    float current   = 0;
    float power     = 0;
    float energy    = 0;
    float frequency = 0;
    float pf        = 0;
    bool  valid     = false;
};


// Forward declarations
void connectWifi();
void ensureMqtt();
void publishPzem(const char* topic, const PzemData& d, uint8_t addr);
void publishAds();
float measureAdsRms(uint8_t channel, uint16_t samples);
void updateOled();

HardwareSerial pzemSerial(2);
PZEM004Tv30    pzem_main(pzemSerial, PZEM_RX, PZEM_TX, PZEM_ADDR_MAIN);
PZEM004Tv30    pzem_sub (pzemSerial, PZEM_RX, PZEM_TX, PZEM_ADDR_SUB);

Adafruit_ADS1115 ads;
Adafruit_SSD1306 display(SCREEN_W, SCREEN_H, &Wire);

WiFiClient   wifiClient;
PubSubClient mqtt(wifiClient);

// Latest PZEM readings — main circuit

PzemData mainData, subData;

// ADS1115 RMS buffer (compute RMS over one cycle at 50 Hz)
static const uint16_t ADS_SAMPLES_PER_CYCLE = 20; // 10 Hz × 2 cycles gives some averaging
float adsRms[4] = {0, 0, 0, 0};

unsigned long lastPzemMs  = 0;
unsigned long lastAdsMs   = 0;
unsigned long lastOledMs  = 0;
unsigned long lastMqttMs  = 0;

// ══════════════════════════════════════════════════════════
//  Helper: connect WiFi
// ══════════════════════════════════════════════════════════
void connectWifi() {
    Serial.printf("Connecting to WiFi: %s\n", WIFI_SSID);
    WiFi.mode(WIFI_STA);
    WiFi.begin(WIFI_SSID, WIFI_PASSWORD);
    uint8_t attempts = 0;
    while (WiFi.status() != WL_CONNECTED && attempts < 30) {
        delay(500);
        Serial.print('.');
        attempts++;
    }
    if (WiFi.status() == WL_CONNECTED) {
        Serial.printf("\nWiFi connected. IP: %s\n", WiFi.localIP().toString().c_str());
    } else {
        Serial.println("\nWiFi FAILED — running in offline mode.");
    }
}

// ══════════════════════════════════════════════════════════
//  Helper: MQTT reconnect (non-blocking attempt)
// ══════════════════════════════════════════════════════════
void ensureMqtt() {
    if (WiFi.status() != WL_CONNECTED) return;
    if (mqtt.connected()) return;
    Serial.print("MQTT connecting...");
    if (mqtt.connect(MQTT_CLIENT)) {
        Serial.println(" OK");
    } else {
        Serial.printf(" FAILED rc=%d\n", mqtt.state());
    }
}

// ══════════════════════════════════════════════════════════
//  Helper: publish a PzemData struct as JSON
// ══════════════════════════════════════════════════════════
void publishPzem(const char* topic, const PzemData& d, uint8_t addr) {
    if (!mqtt.connected()) return;
    StaticJsonDocument<256> doc;
    doc["ts"]   = millis();           // ms uptime; replace with NTP epoch in production
    doc["addr"] = addr;
    doc["V"]    = d.voltage;
    doc["A"]    = d.current;
    doc["W"]    = d.power;
    doc["VAr"]  = d.power * (d.pf > 0.001f ? sqrtf(1.0f - d.pf * d.pf) / d.pf : 0);
    doc["PF"]   = d.pf;
    doc["Hz"]   = d.frequency;
    doc["kWh"]  = d.energy;
    doc["valid"]= d.valid;
    char buf[256];
    serializeJson(doc, buf, sizeof(buf));
    mqtt.publish(topic, buf);
}

// ══════════════════════════════════════════════════════════
//  Helper: publish ADS1115 RMS currents
// ══════════════════════════════════════════════════════════
void publishAds() {
    if (!mqtt.connected()) return;
    StaticJsonDocument<192> doc;
    doc["ts"] = millis();
    for (uint8_t ch = 0; ch < 4; ch++) {
        char key[8];
        snprintf(key, sizeof(key), "A_ch%u", ch);
        doc[key] = adsRms[ch];
    }
    char buf[192];
    serializeJson(doc, buf, sizeof(buf));
    mqtt.publish(MQTT_TOPIC_ADS, buf);
}

// ══════════════════════════════════════════════════════════
//  Helper: compute RMS current from ADS1115 channel
//  Accumulates SAMPLE_N readings, returns RMS amp value
// ══════════════════════════════════════════════════════════
float measureAdsRms(uint8_t channel, uint16_t samples) {
    ads.setGain(GAIN_TWO);  // ±2.048 V
    double sum = 0;
    for (uint16_t i = 0; i < samples; i++) {
        int16_t raw = ads.readADC_SingleEnded(channel);
        float v = raw * ADS_PGA_SCALE;
        // Remove DC bias (1.65 V mid-rail): subtract ~26400 raw (1.65/0.0000625)
        // We use AC coupling assumption: mean removal in sliding window not shown here
        // For simplicity, square and accumulate (bias cancels in differential config)
        sum += (double)v * v;
        delayMicroseconds(200); // ~5 kHz effective sample rate
    }
    float vrms = sqrtf((float)(sum / samples));
    // Convert V_rms → I_primary_rms
    return (vrms / SCT_BURDEN_OHM) * SCT_TURNS_RATIO;
}

// ══════════════════════════════════════════════════════════
//  OLED display update
// ══════════════════════════════════════════════════════════
void updateOled() {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);

    if (mainData.valid) {
        display.printf("V:%.1fV  Hz:%.1f\n", mainData.voltage, mainData.frequency);
        display.printf("I:%.2fA  PF:%.2f\n", mainData.current, mainData.pf);
        display.printf("P:%.1fW\n",          mainData.power);
        display.printf("E:%.3fkWh\n",        mainData.energy);
    } else {
        display.println("Main: No data");
    }

    display.drawLine(0, 34, 127, 34, SSD1306_WHITE);

    display.setCursor(0, 37);
    if (subData.valid) {
        display.printf("Sub %.1fW %.3fkWh\n", subData.power, subData.energy);
    } else {
        display.println("Sub: No data");
    }

    display.setCursor(0, 50);
    display.printf("WiFi:%s MQTT:%s",
        WiFi.status() == WL_CONNECTED ? "OK" : "--",
        mqtt.connected()               ? "OK" : "--");

    display.display();
}

// ══════════════════════════════════════════════════════════
//  Setup
// ══════════════════════════════════════════════════════════
void setup() {
    Serial.begin(115200);
    delay(500);
    Serial.println("\n=== NILM Energy Meter Booting ===");

    // I2C (SDA=21, SCL=22 — ESP32 defaults)
    Wire.begin(21, 22);

    // OLED
    if (!display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDR)) {
        Serial.println("OLED init FAILED — continuing without display.");
    } else {
        display.clearDisplay();
        display.setTextSize(1);
        display.setTextColor(SSD1306_WHITE);
        display.setCursor(0, 0);
        display.println("NILM Meter v1.0");
        display.println("Booting...");
        display.display();
    }

    // ADS1115
    ads.setDataRate(RATE_ADS1115_860SPS);
    if (!ads.begin(ADS_ADDR)) {
        Serial.println("ADS1115 init FAILED — SCT-013 channels disabled.");
    } else {
        Serial.println("ADS1115 OK");
    }

    // PZEM UART2
    pzemSerial.begin(9600, SERIAL_8N1, PZEM_RX, PZEM_TX);
    delay(100);
    Serial.println("PZEM UART2 started.");

    // WiFi
    connectWifi();

    // MQTT
    mqtt.setServer(MQTT_HOST, MQTT_PORT);
    mqtt.setKeepAlive(30);
    mqtt.setSocketTimeout(5);
    ensureMqtt();

    Serial.println("Setup complete. Entering loop.");
}

// ══════════════════════════════════════════════════════════
//  Loop
// ══════════════════════════════════════════════════════════
void loop() {
    unsigned long now = millis();

    // 1. MQTT keep-alive
    if (WiFi.status() == WL_CONNECTED) {
        ensureMqtt();
        mqtt.loop();
    }

    // 2. Poll PZEM sensors every PZEM_INTERVAL_MS
    if (now - lastPzemMs >= PZEM_INTERVAL_MS) {
        lastPzemMs = now;

        // ── Main circuit ──────────────────────────────────
        mainData.voltage   = pzem_main.voltage();
        mainData.current   = pzem_main.current();
        mainData.power     = pzem_main.power();
        mainData.energy    = pzem_main.energy();
        mainData.frequency = pzem_main.frequency();
        mainData.pf        = pzem_main.pf();
        mainData.valid     = !isnan(mainData.voltage) && !isnan(mainData.current);

        if (mainData.valid) {
            Serial.printf("[MAIN] V=%.1f A=%.3f W=%.1f PF=%.2f Hz=%.1f kWh=%.3f\n",
                mainData.voltage, mainData.current, mainData.power,
                mainData.pf, mainData.frequency, mainData.energy);
            publishPzem(MQTT_TOPIC_MAIN, mainData, PZEM_ADDR_MAIN);
        } else {
            Serial.println("[MAIN] PZEM read error — check wiring and mains connection.");
        }

        // ── Sub-circuit (optional) ────────────────────────
        subData.voltage   = pzem_sub.voltage();
        subData.current   = pzem_sub.current();
        subData.power     = pzem_sub.power();
        subData.energy    = pzem_sub.energy();
        subData.frequency = pzem_sub.frequency();
        subData.pf        = pzem_sub.pf();
        subData.valid     = !isnan(subData.voltage) && !isnan(subData.current);

        if (subData.valid) {
            Serial.printf("[SUB]  V=%.1f A=%.3f W=%.1f PF=%.2f Hz=%.1f kWh=%.3f\n",
                subData.voltage, subData.current, subData.power,
                subData.pf, subData.frequency, subData.energy);
            publishPzem(MQTT_TOPIC_SUB, subData, PZEM_ADDR_SUB);
        }
        // No-error print for sub; it may simply be absent
    }

    // 3. ADS1115 high-rate RMS sampling every ADS_INTERVAL_MS
    if (now - lastAdsMs >= ADS_INTERVAL_MS) {
        lastAdsMs = now;
        // Sample each channel; 50 samples each at ~200µs → ~10 ms per channel
        for (uint8_t ch = 0; ch < 4; ch++) {
            adsRms[ch] = measureAdsRms(ch, 50);
        }
        publishAds();
    }

    // 4. OLED refresh every OLED_INTERVAL_MS
    if (now - lastOledMs >= OLED_INTERVAL_MS) {
        lastOledMs = now;
        updateOled();
    }

    // 5. WiFi reconnect watchdog (every 30 s)
    static unsigned long lastWifiCheck = 0;
    if (now - lastWifiCheck >= 30000) {
        lastWifiCheck = now;
        if (WiFi.status() != WL_CONNECTED) {
            Serial.println("WiFi lost — reconnecting...");
            WiFi.reconnect();
        }
    }
}