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

def model(plot=False):
    # Model Parameters
    nhrs = 30

    data_file = './caiso_steady.csv'
    tes_cp = 1530           # J/kg K, TES heat capacity
    tes_t0 = 400            # K,  initial temperture of TES
    tes_mass = 9e8          # kg, mass of TES
    # tes_dTmax = 40          # deg K/hr, max rate of change for TES salt temp
    tes_Tupper = 700        # K, Upper TES temperature limit
    tes_Tlower = 250        # K, Lower TES temperature limit
    tes_eff = .95
    tes_max_ramp_up = 500   # FIXME: Need to update these with realistic values
    tes_max_ramp_down = 500

    npp_cap = 550 # 6*250.0  # 6x250.0 MWth NuScale SMRs
    npp_ramprate = 0.40     # Max ramping as a percent of total capacity
    npp_cost = .021 * 1000   # $/MWh, M&O costs

    turb_efficiency = 0.8   # Max efficiency of steam to electricity conversion
    turb_capacity = 7.5e4  # turbine capacity don't operate above 120% capacity

    t = np.linspace(0, nhrs-1, nhrs)  # time in hours

    data = pd.read_csv(data_file)
    Wind = data['Wind'].values[0:nhrs]
    Solar= data['Solar'].values[0:nhrs]
    Load = data['Load'].values[0:nhrs]

    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)

    # A decision variable describing what percent of the generated heat stored.
    # It can go negative to simulate retreiving energy from the TES.
    tes_store = m.MV(value=0, lb=-1, ub=1)
    tes_store.STATUS = 1

    # NPP operating characteristics
    npp_gen = m.MV(value=.7*npp_cap, lb=0.2*npp_cap, ub=npp_cap)
    npp_gen.STATUS = 1
    #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
    tlb = .2*(tes_Tupper - tes_Tlower) + tes_Tlower  # never run below 20%
    tes_T = m.Var(value=tes_t0, lb=tlb, ub=tes_Tupper)
    # Thermal side balance equation. Amount being stored affects TES temp
    m.Equation(tes_T.dt() == 3.6e9*(tes_store*npp_gen)/(tes_mass*tes_cp))

    # Percs model gives in %cap/s
    m.Equation(tes_T.dt() <= tes_max_ramp_up*3600*(tes_Tupper - tes_Tlower))
    m.Equation(tes_T.dt() >= tes_max_ramp_down*3600*(tes_Tupper - tes_Tlower))

    # 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)

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

    net_load = m.Intermediate(load - wind - solar)
    m.Equation(net_load == turb_out)

    # Minimize both capital and dynamic costs
    # m.Obj(npp_gen*npp_cost)

    m.options.IMODE = 5
    m.options.SOLVER = 3
    m.options.COLDSTART = 2
    m.solve()

    if plot:
        # """
        plt.figure(figsize=(10, 10))
        plt.subplot(4, 1, 1)
        plt.title('Simple NHES')
        plt.plot(t, Load, label='Demand 1:Load')
        plt.plot(t, npp_gen.value, label='Production 1:Nuclear')
        plt.plot(t, turb_out.value, label='Production 1.1:Turbine')
        plt.plot(t, Wind, label='Production 2:Wind')
        plt.plot(t, Solar, label='Production 3:Solar')
        plt.ylabel('Energy (MW)')
        plt.legend(loc='upper right')

        plt.subplot(4, 1, 2)
        plt.plot(t, tes_T.value, '--', color='orange', label='TES Temp')
        plt.ylabel('Temp (K)')
        plt.legend(loc='upper right')

        plt.subplot(4, 1, 3)
        ar = np.array
        Ebal = ar(npp_gen.value)*(1-ar(tes_store.value) -
                                  ar(Eloss_tes.value)) * \
            turb_efficiency - net_load
        plt.plot(t[1:], Ebal[1:], label='Energy Balance')
        plt.legend(loc='upper right')

        plt.subplot(4, 1, 4)
        plt.plot(t, tes_store.value, label='TES_store')
        plt.legend(loc='upper right')
        plt.xlabel('Time (Hr)')
        plt.show()

    M = m.options

    return {
        'time_steps': len(t),
        'fcalls': M.ITERATIONS,
        'f': M.OBJFCNVAL,
        'time (s)': M.SOLVETIME,
        'message': message,
        'status': M.APPSTATUS
    }


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