from gekko import GEKKO
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sqlalchemy import create_engine

"""
This model has a nuclear power plant, and uses Wind and gas
energy sources as the main contributers to the ERCOT grid.
It also uses 3 energy storage systems to optimize
a grid generation-load system. They have efficiencies and can
be limited by capacities.
"""


def model(plot=False):
    # Model Parameters
    nhrs = 120
    db_weather = '../../data/Ercot/Ercot Weather/weather_TX_west_2019.csv'
    db_load = '../../data/Ercot/Ercot Load/ERCOT_load_2019.csv'
    db_path = '../../data/ercot_data.db'

    tes_cp = 1530           # J/kg K, TES heat capacity
    tes_t0 = 300            # K,  initial temperture of TES
    tes_mass = 9e8          # kg, mass of TES
    tes_cost = 4.99         # $/lb, cost of TES unit
    tes_dTmax = 400         # deg K/hr, max rate of change for TES salt temperature
    tes_Tupper = 700        # K, Upper TES temperature limit
    tes_Tlower = 250        # K, Lower TES temperature limit
    tes_eff = .95

    bat_cap = 300            # MWhr, battry capacity
    bat_cost = 400*1000      # $/MWhr, Cost of the battery unit, Quickly guestimate
    bat_eff = .99

    npp_cap = 1280 * 20     # MW, capacity of South Texas Nuclear Generatingstation
    npp_ramprate = 0.01     # Max ramping as a percent of total capacity
    npp_cost = .021 * 1000  # $/MWh, M&O costs
    turb_efficiency = 0.8   # Efficieny of steam to electricity conversion

    chi_mass = 1e4          # kg of water in tank
    chi_t0 = 298            # Intitial Temp
    chi_cp = 4182           # J/Kg K, Heat capacity of water
    chi_eff = .7            # % of energy conserved when pulling energy
    chi_dTmax = 20          # K/hr, rate of chilled storage temperature
    chi_UA = .005           # Heatloss coefficient of system chiller is connected

    query = """
    SELECT Generation, Fuel, Date_Time from Generation
    WHERE Resolution = "Hourly" and
    Date_Time BETWEEN date("2019-01-01") AND date("2019-12-31")
    """
    engine = create_engine(f'sqlite:///{db_path}')

    data = pd.read_sql(query, con=engine)
    weather = pd.read_csv(db_weather)
    Loaddata = pd.read_csv(db_load)
    Winddata = data[data['Fuel'] == 'Wind']
    Solardata = data[data['Fuel'] == 'Solar']

    Wind = Winddata['Generation'].values[0:nhrs]
    Solar = Solardata['Generation'].values[0:nhrs]
    Load = .5*Loaddata['ERCOT'].values[0:nhrs]
    t = np.linspace(0, nhrs-1, nhrs)  # time in hours

    T_outside_F = weather['Temperature(F)'].values[4320:][0:nhrs]  # summer temps
    T_oustide_K = (T_outside_F-32)*5/9+273.15

    m = GEKKO(remote=True)
    m.time = t

    # Load from ERCOT
    load = m.Param(value=Load)
    wind = m.Param(value=Wind)
    solar = m.Param(value=Solar)
    Temp = m.Param(value=T_oustide_K)

    # A decision variable describing what percent of the generated heat is stored.
    # It can go negative to simulate retreiving energy from the TES.
    tes_store = m.Var(value=0, lb=-1, ub=1)
    bat_store = m.Var(value=0, lb=-1, ub=1)
    chi_cons = m.Var(value=0, lb=0, ub=1)  # % of energy entering storage

    # NPP operating characteristics
    npp_gen = m.MV(value=.8*npp_cap, lb=0.2*npp_cap, ub=npp_cap)
    npp_gen.STATUS = 1
    npp_gen.DMAX = npp_ramprate*npp_cap
    npp_gen.DCOST = 10

    # Tracking the temperature of the TES as a bounded state variable
    tes_T = m.SV(value=tes_t0, lb=tes_Tlower, ub=tes_Tupper)
    chi_T = m.SV(value=chi_t0, lb=273, ub=chi_t0)  # water freezes at 273K

    # Tracking the battery state of charge as a bounded state variable
    bat_SOC = m.SV(value=0, lb=0, ub=100)

    # Thermal side balance equation. Amount being stored affects TES temperature
    m.Equation(tes_T.dt() == 3.6e9*(tes_store*npp_gen)/(tes_mass*tes_cp))

    # Ensures that the TES does not charge/discharge too quickly
    m.Equation(tes_T.dt() <= tes_dTmax)
    m.Equation(tes_T.dt() >= -tes_dTmax)

    # Input to the turbine is the difference between generation and TES storage
    Eloss_tes = m.Intermediate((1 - tes_eff)*m.abs2(tes_store))
    turb_in = m.Intermediate(npp_gen - tes_store*npp_gen - Eloss_tes*npp_gen)

    # Make sure that the turbine never runs backwards...
    m.Equation(turb_in >= 0)

    # Slack variable for keeping the model feasible for infeasible unit sizings.
    # Very important when the model is being optimized in RAVEN.
    # slack = m.Var(value=0, lb=0)

    # Conversion of thermal to electrical energy
    turb_out = m.Intermediate(turb_in*turb_efficiency)

    # Electrical battery can buffer between load and electrical supply
    m.Equation(bat_SOC.dt() * bat_cap == bat_store*turb_out*100)
    m.Equation(chi_T.dt() * (chi_mass * chi_cp)/3.6e9 == chi_UA*(Temp-chi_t0) - (turb_out*chi_cons))
    m.Equation(chi_T.dt() <= chi_dTmax)
    m.Equation(chi_T.dt() >= -chi_dTmax)

    Eloss_bat = m.Intermediate((1 - bat_eff)*m.abs2(bat_store))
    net_load = m.Intermediate(load - wind - solar)
    m.Equation(net_load == (1-bat_store - chi_cons/chi_eff - Eloss_bat)*turb_out)

    # Minimize both capital and dynamic costs
    # m.Obj(npp_gen*npp_cost) <- better objective
    m.Obj(((1-bat_store - chi_cons/chi_eff - Eloss_bat)*turb_out)-net_load)

    m.options.IMODE = 5
    m.options.SOLVER = 3

    m.solve()

    if plot:
        #"""
        plt.figure(figsize=(10, 10))
        plt.subplot(5, 1, 1)
        plt.plot(t, Load, 'r-', label='Load')
        plt.plot(t, Wind, 'g-', label='Wind')
        plt.plot(t, Solar, 'y-', label='Solar')
        plt.plot(t, npp_gen.value, 'b--', label='Nuclear Generation')
        plt.ylabel('Energy (MW)')
        plt.yticks(np.arange(0, 24000+1, 6000))
        plt.legend(loc='upper right')

        plt.subplot(5, 1, 2)
        plt.plot(t, tes_T.value, '--', color='orange', label='TES Temp')
        plt.ylabel('Temp (K)')
        plt.yticks(np.arange(0, 800+1, 200))
        plt.legend(loc='upper right')

        plt.subplot(5, 1, 3)
        plt.plot(t, chi_T.value, label='Chilled Water Temp')
        plt.ylabel('Temp (K)')
        plt.yticks(np.arange(270, 310+1, 10))
        plt.legend(loc='upper right')

        plt.subplot(5, 1, 4)
        plt.plot(t, bat_SOC.value, '--', label='Battery State of Charge')
        plt.ylabel('SOC (%)')
        plt.legend(loc='upper right')

        gen_dt = np.zeros(len(t))
        for i in range(1, len(t)):
            gen_dt[i] = npp_gen.value[i] - npp_gen.value[i-1]
        plt.subplot(5, 1, 5)
        plt.plot(t, gen_dt, label='npp_ramp_rate')
        plt.legend(loc='upper right')
        plt.tight_layout()
        plt.show()
        """
        plt.figure(figsize=(10, 12))
        plt.subplot(5, 1, 1)
        plt.title('Advanced NHES')
        plt.plot(t, Load, 'r:', label='Load')
        plt.plot(t, npp_gen.value, 'b:', label='Nuclear Generation')
        plt.plot(t, Wind, 'b--', label='Wind')
        plt.plot(t, Solar, 'b-', label='Solar')
       
        plt.ylabel('Energy (MW)')
        plt.yticks(np.arange(0, 24000+1, 6000))
        plt.legend(loc='upper right')

        plt.subplot(5, 1, 2)
        plt.plot(t, tes_T.value, '--', color='orange', label='TES Temp')
        plt.ylabel('Temp (K)')
        plt.yticks(np.arange(0, 800+1, 200))
        plt.legend(loc='upper right')

        plt.subplot(5, 1, 3)
        plt.plot(t, chi_T.value, label='Chilled Water Temp')
        plt.ylabel('Temp (K)')
        plt.yticks(np.arange(270, 310+1, 10))
        plt.legend(loc='upper right')

        plt.subplot(5, 1, 4)
        plt.plot(t, bat_SOC.value, '--', label='Battery State of Charge')
        plt.ylabel('SOC (%)')
        plt.legend(loc='upper right')
        
        plt.subplot(5, 1, 5)
        gen_dt = np.zeros(len(t))
        for i in range(1, len(t)):
            gen_dt[i] = npp_gen.value[i] - npp_gen.value[i-1]
        plt.plot(t, gen_dt, 'k:', label='npp_ramp_rate')
        plt.legend(loc='upper right')
        plt.tight_layout()
        plt.xlabel("Time (hr)")
        plt.ylabel('$MW/hr$')
        plt.savefig('advanced_nhes.eps', format='eps')
        plt.show()
        #"""
        

    """
    GuessVal = pd.DataFrame({'Time': t.T})
    GuessVal['Gen'] = np.array(npp_gen.value).T
    GuessVal['tes_store'] = np.array(tes_store.value).T
    GuessVal['bat_store'] = np.array(bat_store.value).T
    GuessVal['chi_cons'] = np.array(chi_cons.value).T
    print(GuessVal.head())
    GuessVal.to_csv('../../data/GekkoSolutions.csv', index=False)
    #"""

    M = m.options
    #return [M.OBJFCNVAL, M.ITERATIONS, M.SOLVETIME, M.APPSTATUS]
    if M.APPSTATUS == 1:
        feasible = True
        message = 'Optimizaiton Terminated Successfully'
    else:
        feasible = False
    error1 = M.OBJFCNVAL
    return {
        'Model':'Gekko Simple-NHES',
        'time_steps':len(t),
        'fcalls':M.ITERATIONS,
        'gcalls':'NA',
        'f':M.OBJFCNVAL,
        'feasible':feasible,
        'ramp err':'NA',
        'total err':error1,
        'time (s)':M.SOLVETIME,
        'message':message,
        'status':M.APPSTATUS
    }

if __name__ == "__main__":
    sol = model(plot=True)
    print(sol)
