#include "RTClib.h" // SPI and Wire are set for OLED
#include <SPI.h>
#include <Wire.h>

// For INA219
#include <Adafruit_INA219.h>
Adafruit_INA219 ina219;

// For OLED SETUP
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#define i2c_Address 0x3C
#define SCREEN_WIDTH 128 // OLED display width, in pixels
#define SCREEN_HEIGHT 64 // OLED display height, in pixels
#define OLED_RESET    -1 // Reset pin # (or -1 if sharing Arduino reset pin)
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

RTC_DS1307 rtc;

// Set for Channels Relays
#define CH1 15
#define CH2 2
#define CH3 0
#define CH4 4

// For auto or manual signal
#define Auto 26

// Set for Manual Mode Switches
#define MANUAL_SWITCH_1 25
#define MANUAL_SWITCH_2 33
#define MANUAL_SWITCH_3 32
#define MANUAL_SWITCH_4 13

// Set for buttons to replace rotary encoder
#define BUTTON_UP    27
#define BUTTON_DOWN  14
#define BUTTON_SELECT 12

// Stored Alarm Value
int ResHr[11];
int ResMin[11];
int ResSec[11];
int Channel_StateReal[41];
bool HaveAlarm[11] = {false, false, false, false, false, false, false, false, false, false, false};
int CheckTime = 1;

int timeHr = 0;
int timeMin = 0;
int timeSec = 0;
int Channel_State[5];
int ChanMult = 1;
int Alarm = 0; // Set for Alarm Arrays
bool AlarmChoose[6]; // Set for Options in Alarm execution
int ChoiceChannels = 0;
bool ChanChoose = false;
int CheckChoice = 1;

// Math Equations
// ((i + 1) * 4) + a = n for intervals of 4 such as timer

// INA219 Variables
float shuntvoltage = 0;
float busvoltage = 0;
float current_mA = 0;
float loadvoltage = 0;
float power_mW = 0;
float finalW = 0;
float CheckWatt = 9.5; // Set For Default Wattage Limit

// Function to read button values
int readButton() {
    if (digitalRead(BUTTON_UP) == LOW) {
        return 1; // UP
    } else if (digitalRead(BUTTON_DOWN) == LOW) {
        return -1; // DOWN
    } else if (digitalRead(BUTTON_SELECT) == LOW) {
        return 2; // SELECT
    } else {
        return 0; // NO BUTTON PRESSED
    }
}

// RTC Ini then Select 
void setup() {
    Serial.begin(115200);

    pinMode(Auto, INPUT_PULLUP);
    pinMode(CH1, OUTPUT);
    pinMode(CH2, OUTPUT);
    pinMode(CH3, OUTPUT);
    pinMode(CH4, OUTPUT);

    pinMode(MANUAL_SWITCH_1, INPUT_PULLUP);
    pinMode(MANUAL_SWITCH_2, INPUT_PULLUP);
    pinMode(MANUAL_SWITCH_3, INPUT_PULLUP);
    pinMode(MANUAL_SWITCH_4, INPUT_PULLUP);

    pinMode(BUTTON_UP, INPUT_PULLUP);
    pinMode(BUTTON_DOWN, INPUT_PULLUP);
    pinMode(BUTTON_SELECT, INPUT_PULLUP);
    
    digitalWrite(CH1, HIGH);
    digitalWrite(CH2, HIGH);
    digitalWrite(CH3, HIGH);
    digitalWrite(CH4, HIGH);

    delay(250); // wait for the OLED to power up

    if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) { // Address 0x3C for 128x64
        Serial.println(F("SSD1306 allocation failed"));
        for (;;);
    }
    
    delay(1000);
    display.clearDisplay();
    
    // RTC Set and Find
    if (!rtc.begin()) {
        display.println("Couldn't find RTC");
        Serial.flush();
        while (1) delay(10);
    }

    if (!rtc.isrunning()) {
        display.println("RTC is NOT running, let's set the time!");
        rtc.adjust(DateTime(2014, 1, 21, 22, 55, 0));
    }

    // Set Text
    display.setTextSize(1);
    display.setCursor(0, 0);
    display.setTextColor(SSD1306_WHITE);

    // DebugData
    ResHr[1] = 11;
    ResMin[1] = 35;
    ResSec[1] = 0;
    HaveAlarm[1] = true;
    Channel_StateReal[1] = false;
    Channel_StateReal[2] = false;
    Channel_StateReal[3] = false;
    Channel_StateReal[4] = false;

    ResHr[2] = 19;
    ResMin[2] = 35;
    ResSec[2] = 5;
    HaveAlarm[2] = true;
    Channel_StateReal[5] = false;
    Channel_StateReal[6] = true;
    Channel_StateReal[7] = false;
    Channel_StateReal[8] = true;

    ResHr[3] = 19;
    ResMin[3] = 35;
    ResSec[3] = 10;
    HaveAlarm[3] = true;
    Channel_StateReal[9] = true;
    Channel_StateReal[10] = true;
    Channel_StateReal[11] = true;
    Channel_StateReal[12] = true;
}

void loop() {
    display.clearDisplay();
    RTCTick_Power();
    
    if (digitalRead(Auto) == HIGH) {
        // Automatic Mode
        switch (CheckChoice) {
            case 1:
                SelectOpt();
                break;
            case 2:
                SelectTimer();
                break;
            case 3:
                SelectAlarm();
                break;
            case 4:
                SelectLimit();
                break;
        }
    } else {
        // Manual Mode
        display.setCursor(0, 16);
        display.print("System set to Manual");

        digitalWrite(CH1, digitalRead(MANUAL_SWITCH_1) == HIGH ? LOW : HIGH);
        digitalWrite(CH2, digitalRead(MANUAL_SWITCH_2) == HIGH ? LOW : HIGH);
        digitalWrite(CH3, digitalRead(MANUAL_SWITCH_3) == HIGH ? LOW : HIGH);
        digitalWrite(CH4, digitalRead(MANUAL_SWITCH_4) == HIGH ? LOW : HIGH);
    }
    
    display.display();
}

void RTCTick_Power() {
    // Check Power
    shuntvoltage = ina219.getShuntVoltage_mV();
    busvoltage = ina219.getBusVoltage_V();
    current_mA = ina219.getCurrent_mA();
    power_mW = ina219.getPower_mW();
    loadvoltage = busvoltage + (shuntvoltage / 1000);
    // Watt-Hour
    finalW = ((power_mW * 1000) * 60) * 60;

    display.setTextSize(1);
    display.setCursor(0, 0);
    display.setTextColor(SSD1306_WHITE);
    DateTime now = rtc.now();
    if (CheckWatt >= power_mW * 1000) {
        for (CheckTime = 1; CheckTime <= 10; CheckTime++) {
            if (HaveAlarm[CheckTime] == true && digitalRead(Auto) == HIGH) {
                if (ResHr[CheckTime] == now.hour() && ResMin[CheckTime] == now.minute() && ResSec[CheckTime] == now.second()) {
                    digitalWrite(CH1, Channel_StateReal[((CheckTime - 1) * 4) + 1]);
                    digitalWrite(CH2, Channel_StateReal[((CheckTime - 1) * 4) + 2]);
                    digitalWrite(CH3, Channel_StateReal[((CheckTime - 1) * 4) + 3]);
                    digitalWrite(CH4, Channel_StateReal[((CheckTime - 1) * 4) + 4]);
                }
            }
        }

        display.print(now.hour(), DEC);
        display.print(":");
        display.print(now.minute(), DEC);
        display.print(":");
        display.print(now.second(), DEC);

        // Display WattHour
        display.print(" [");
        display.print(finalW);
        display.print("Wh]");

    } else {
        // Turn off all relays if power limit exceeded
        digitalWrite(CH1, HIGH);
        digitalWrite(CH2, HIGH);
        digitalWrite(CH3, HIGH);
        digitalWrite(CH4, HIGH);

        display.setCursor(0, 0);
        display.print("WARNING! LIMIT");
        display.setCursor(0, 8);
        display.print("HAS BEEN");
        display.setCursor(0, 16);
        display.print("SURPASSED!!!");
        CheckChoice = 0;
    }
}

void SelectOpt() {
    static int option = 1; // Default option
    int btnState = readButton();
    if (btnState == 1) {
        option = min(option + 1, 3); // Go up
    } else if (btnState == -1) {
        option = max(option - 1, 1); // Go down
    } else if (btnState == 2) {
        // Select option
        if (option == 1) {
            CheckChoice = 2;
        } else if (option == 2) {
            CheckChoice = 3;
        } else if (option == 3) {
            CheckChoice = 4;
        }
    }

    // Display current option
    display.setCursor(0, 24);
    if (option == 1) {
        display.print("Set Timer");
    } else if (option == 2) {
        display.print("Alarms");
    } else if (option == 3) {
        display.print("Set Watt Limit");
    }

    delay(150);
    return;
}

void SelectTimer() {
    static int step = 0;
    int btnState = readButton();
    switch (step) {
        case 0: // Select Alarm
            display.setCursor(0, 32);
            display.print("Choose Timer:");
            display.setCursor(0, 40);
            Alarm = min(max(Alarm + btnState, 1), 10); // Adjust Alarm number
            display.print("Alarm: ");
            display.print(Alarm);
            if (btnState == 2) {
                step++;
            }
            break;
        case 1: // Set Hour
            display.setCursor(0, 40);
            display.print("Set Hour: ");
            timeHr = min(max(timeHr + btnState, 0), 23);
            display.print(timeHr);
            if (btnState == 2) {
                step++;
            }
            break;
        case 2: // Set Minute
            display.setCursor(0, 40);
            display.print("Set Minute: ");
            timeMin = min(max(timeMin + btnState, 0), 59);
            display.print(timeMin);
            if (btnState == 2) {
                step++;
            }
            break;
        case 3: // Set Second
            display.setCursor(0, 40);
            display.print("Set Second: ");
            timeSec = min(max(timeSec + btnState, 0), 59);
            display.print(timeSec);
            if (btnState == 2) {
                step++;
            }
            break;
        case 4: // Set Channels
            display.setCursor(0, 32);
            display.print("Set Lights:");
            display.setCursor(0, 40);
            ChoiceChannels = min(max(ChoiceChannels + btnState, 1), 5);
            if (ChoiceChannels == 5) {
                display.print("Finish");
                if (btnState == 2) {
                    step++;
                }
            } else {
                display.print("Channel ");
                display.print(ChoiceChannels);
                if (btnState == 2) {
                    ChanChoose = true;
                    step++;
                }
            }
            break;
        case 5: // Set Channel States
            if (ChanChoose) {
                display.setCursor(0, 40);
                display.print("Channel ");
                display.print(ChoiceChannels);
                display.setCursor(0, 48);
                display.print("Setting to: ");
                int state = min(max(Channel_State[ChoiceChannels] + btnState, 0), 1);
                Channel_State[ChoiceChannels] = state;
                if (state == 1) {
                    display.print("ON");
                } else {
                    display.print("OFF");
                }
                if (btnState == 2) {
                    ChanChoose = false;
                    step--;
                }
            }
            break;
        case 6: // Confirm Timer
            display.setCursor(0, 48);
            display.print("Confirm Settings");
            if (btnState == 1) {
                display.setCursor(0, 56);
                display.print("YES");
                if (readButton() == 2) {
                    // Save settings
                    ResHr[Alarm] = timeHr;
                    ResMin[Alarm] = timeMin;
                    ResSec[Alarm] = timeSec;
                    for (ChanMult = 1; ChanMult <= 4; ChanMult++) {
                        Channel_StateReal[((Alarm - 1) * 4) + ChanMult] = Channel_State[ChanMult];
                    }
                    HaveAlarm[Alarm] = true;
                    step = 0;
                    CheckChoice = 1; // Return to main menu
                }
            } else if (btnState == -1) {
                display.setCursor(0, 56);
                display.print("NO");
                if (readButton() == 2) {
                    step = 0;
                    CheckChoice = 1; // Return to main menu
                }
            }
            break;
    }
}

void SelectAlarm() {
    display.setCursor(0, 16);
    display.print("Press Button to Exit");
    display.setCursor(0, 24);
    display.print("Alarms (Up to 10): ");
    int btnState = readButton();
    Alarm = min(max(Alarm + btnState, 1), 10); // Adjust Alarm number
    display.print(Alarm);
    display.setCursor(0, 32);
    display.print(ResHr[Alarm]);
    display.print(":");
    display.print(ResMin[Alarm]);
    display.print(":");
    display.print(ResSec[Alarm]);
    display.setCursor(0, 40);
    display.print("Relay Status: ");
    display.setCursor(0, 48);
    display.print("CH1: ");
    display.print(Channel_StateReal[((Alarm - 1) * 4) + 1]);
    display.print(" CH2: ");
    display.print(Channel_StateReal[((Alarm - 1) * 4) + 2]);
    display.setCursor(0, 56);
    display.print("CH3: ");
    display.print(Channel_StateReal[((Alarm - 1) * 4) + 3]);
    display.print(" CH4: ");
    display.print(Channel_StateReal[((Alarm - 1) * 4) + 4]);
    if (btnState == 2) {
        CheckChoice = 1;
    }
}

void SelectLimit() {
    display.setCursor(0, 24);
    display.print("Limit for Watts: ");
    int btnState = readButton();
    CheckWatt = min(max(CheckWatt + (btnState * 10), 0)), 999); // Adjust limit
    display.setCursor(0, 32);
    display.print(CheckWatt);
    display.print("W");
    display.setCursor(0, 40);
    display.print("Press Button to Exit");
    if (btnState == 2) {
        CheckChoice = 1;
    }
}