#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <EmonLib.h>  // Energy monitoring library

// OLED Display Configuration
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET    -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Define GPIO Pins
#define GREEN_LED 12
#define YELLOW_LED 14
#define RED_LED 15
#define RELAY_PIN 27
#define BUZZER_PIN 32
#define BUTTON_PIN 33

EnergyMonitor emon;  // Energy monitoring object
bool manualOverride = false;  // Manual relay control
float voltage = 0.0, current = 0.0, powerConsumption = 0.0;

void displayMessage(String line1, String line2);  // Function prototype

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

    // Initialize OLED Display
    if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
        Serial.println(F("SSD1306 allocation failed"));
        while (1);
    }
    display.clearDisplay();

    // Setup GPIO Pins
    pinMode(GREEN_LED, OUTPUT);
    pinMode(YELLOW_LED, OUTPUT);
    pinMode(RED_LED, OUTPUT);
    pinMode(RELAY_PIN, OUTPUT);
    pinMode(BUZZER_PIN, OUTPUT);
    pinMode(BUTTON_PIN, INPUT_PULLUP); // Enable internal pull-up resistor

    // Initialize Energy Monitor
    emon.voltage(35, 234.26, 1.7);  // Voltage Sensor: pin 35, calibration values
    emon.current(34, 111.1);        // Current Sensor: pin 34, calibration value

    // Display startup message
    displayMessage("Smart Grid System", "Initializing...");
    delay(2000);
    display.clearDisplay();
}

void loop() {
    // Read real-time voltage, current, and power
    emon.calcVI(20, 2000);  // Read 20 half cycles, timeout 2000ms
    voltage = emon.Vrms;
    current = emon.Irms;
    powerConsumption = voltage * current;

    // Print on Serial Monitor
    Serial.print("Voltage: "); Serial.print(voltage); Serial.println(" V");
    Serial.print("Current: "); Serial.print(current); Serial.println(" A");
    Serial.print("Power: "); Serial.print(powerConsumption); Serial.println(" W");
    
    // Display on OLED
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 10);
    display.print("Voltage: "); display.print(voltage); display.println(" V");
    display.setCursor(0, 25);
    display.print("Current: "); display.print(current); display.println(" A");
    display.setCursor(0, 40);
    display.print("Power: "); display.print(powerConsumption); display.println(" W");
    display.display();

    // Check manual override button
    if (digitalRead(BUTTON_PIN) == LOW) {
        manualOverride = !manualOverride;
        delay(500);  // Debounce delay
    }

    // Reset LEDs and Buzzer
    digitalWrite(GREEN_LED, LOW);
    digitalWrite(YELLOW_LED, LOW);
    digitalWrite(RED_LED, LOW);
    digitalWrite(BUZZER_PIN, LOW);

    // Load Balancing Logic with LED Indication
    if (!manualOverride) {
        if (powerConsumption < 200) {
            digitalWrite(GREEN_LED, HIGH);
            digitalWrite(RELAY_PIN, HIGH);  // Load ON (Safe)
        } 
        else if (powerConsumption < 350) {
            digitalWrite(YELLOW_LED, HIGH);
            digitalWrite(RELAY_PIN, HIGH);  // Load ON (Medium Usage)
        } 
        else {
            digitalWrite(RED_LED, HIGH);
            digitalWrite(RELAY_PIN, LOW);  // Load OFF (High Usage)
            digitalWrite(BUZZER_PIN, HIGH);  // Alert
            delay(500);
            digitalWrite(BUZZER_PIN, LOW);
        }
    } 
    else {
        Serial.println("Manual Override Enabled");
    }

    delay(3000);  // Update every 3 seconds
}

// Function to display messages on OLED
void displayMessage(String line1, String line2) {
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 10);
    display.print(line1);
    display.setCursor(0, 30);
    display.print(line2);
    display.display();
}