import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_theme()
from gekko import GEKKO

nhrs = 60 # Got up to 1410 to work!

smr_plant_lifetime = 60 # years, expected lifetime of the nuc plant, used for LCOE
sol_plant_lifetime = 35 # years, expected lifetime of the plant, used for LCOE
wind_plant_lifetime = 30 # years, expected lifetime of the plant, used for LCOE

#Need to adjust LCOE based on minimum plant lifetime. This will adjust fixed costs by factor of sys_plant_lifetime/x_plant_lifetime
sys_plant_lifetime = wind_plant_lifetime #total system is expected to work for 30 years


#data_file = './caiso_steady.csv'
data_file = './more_data_scaled.csv'
tes_cp = 1530           # J/kg K, TES heat capacity
tes_t0 = 400            # K,  initial temperture of TES
tes_mass_0 = 1808000    # kg, initial 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

smr_cap = 550 # 6*250.0  # 6x250.0 MWth NuScale SMRs
smr_ramprate = 0.40     # Max ramping as a percent of total capacity

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]

# ---------------- Economic Parameters ----------------
# Sources:
# 1. EIA Capital Cost Estimates (https://www.eia.gov/analysis/studies/powerplants/capitalcost/pdf/capcost_assumption.pdf)
# 2. SMR economic viability (https://www.sciencedirect.com/science/article/pii/S1364032118308372)
# 3. Combined-Cycle Gas and Steam Turbine Power Plants
#   (https://books.google.com/books?id=OmnOG7vWfuQC&lpg=PP1&ots=VsvxReC5bN&dq=steam%20turbine%20economics&lr&pg=PA25#v=onepage&q=steam%20turbine%20economics&f=false)
# 4. Nuscale Electricity and water cogeneration (https://www.sciencedirect.com/science/article/pii/S0011916414000885)

# SMR and turbine capital costs
# Source 2 for capital cost
# Source 4 for operating and maintenence costs
smr_fix_cost = (3465.72 * smr_cap * 1000 + 185000000) * (smr_plant_lifetime/sys_plant_lifetime) # $

# TES capital cost
# Source: https://www.nrel.gov/docs/fy12osti/53066.pdf used $15/kWh as a cheap estimate
# Couldn't find a source for fixed operating costs... :( They should be pretty minimal though.
tes_fix_cost = 15*1000 * (smr_plant_lifetime/sys_plant_lifetime) # $/MWh with plant correction factor estimated as 

# Wind capital and fixed operating costs
# Source 1
wind_cap_cost = 1877  * max(Wind) * 1000 # $
wind_OM_cost = 39.7 * wind_plant_lifetime * max(Wind) * 1000 #$
wind_fix_cost = (wind_cap_cost + wind_OM_cost) * (sys_plant_lifetime/wind_plant_lifetime) #adjusts fixed costs

# Solar capital and fixed operating costs
# Source 1
solar_cap_cost = 2534  * max(Solar) * 1000 # $
solar_OM_cost = 21.8 * sol_plant_lifetime * max(Solar) * 1000 #$
solar_fix_cost = (solar_cap_cost + solar_OM_cost) * (sys_plant_lifetime/wind_plant_lifetime) #adjusts fixed costs

# Variable Operating costs
smr_vom_cost = 1.0 # $/MWh Approximated from Source 3, may be overlapping with OM cost in source 4
tes_vom_cost = 0.0 # $/MWh Current approximation. No sources found.
turb_vom_cost = 1.0 # $/MWh Approximated from Source 3
wind_vom_cost = 0.0 # Source 1
solar_vom_cost = 0.0 # Source 1

# Variable Operating costs (multiply to approximate 60 yr lifetime)
smr_om_cost = 1.0 # $/MWh Approximated from Source 3, may be overlapping with OM cost in source 4
tes_om_cost = 0.0 # $/MWh Current approximation. No sources found.
turb_om_cost = 1.0 # $/MWh Approximated from Source 3
wind_om_cost = 0.0 # Source 1
solar_om_cost = 0.0 # Source 1


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

# Data from CAISO
load = m.Param(value=Load)
wind = m.Param(value=Wind)
solar = m.Param(value=Solar)

#max amount of storage in MWh
tes_batt_max = m.FV(value=100, lb=0, ub=10000000)
tes_batt_max.STATUS = 1

# 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

smr_gen = m.MV(value=400, lb=200, ub=1000, fixed_initial=False)
smr_gen.STATUS = 1

tes_in = m.Var()
m.Equation(tes_in == tes_store * smr_gen)
turb_out = m.Intermediate(turb_efficiency*(smr_gen - tes_in))

tes_batt = m.Var(value = 0, lb = 0, ub = 10000)

m.Equation(tes_batt.dt() == tes_in)
m.Equation(tes_batt >= 0) #storage must stay positive
m.Equation(tes_batt <= tes_batt_max) #storage must stay below max stoarge used for cost calcs

# Limit the TES rate of change
#m.Equation(tes_in.dt() <= tes_max_ramp_up)

# Make sure the turbine does not run backwards
m.Equation(turb_out >= 0)

# Total electricity balance
m.Equation(turb_out + wind + solar == load)
m.Equation(smr_gen <= smr_cap)

# Not entirely sure why, but this variable and constraint are required to make
# the problem feasible
sum_load = m.Var(value = 10000)
m.Equation(sum_load == m.integral(load))

avg_MW = m.Intermediate(sum_load/nhrs) #MW/hr load, used to calculate sys lifetime energy produced
avg_MW_sol = m.Intermediate(m.integral(solar)/nhrs)
avg_MW_wind = m.Intermediate(m.integral(wind)/nhrs)
avg_MW_smr = m.Intermediate(m.integral(smr_gen)/nhrs)
hrs_per_yr = 8760 # number of hours in a year

final = np.zeros(nhrs)
final[-1] = 1
f = m.Param(final)

# LCOE is the sum of the costs over the sum of the electricity provided
LCOE = m.Intermediate(
    # Fixed Capital and OM Costs
    f*(smr_fix_cost + # SMR&Turb fixed capital and OM costs
    tes_fix_cost*tes_batt_max + # TES fixed capital and OM costs as a function of mass
    wind_cap_cost +  # wind fixed capital and OM costs
    solar_cap_cost + # solar fixed capital and OM costs
    # variable operating costs as a function of energy produced
    smr_vom_cost * avg_MW_smr * hrs_per_yr * sys_plant_lifetime  + # SMR&Turb VOM cost
    tes_vom_cost * 0 + #VOM costs of TES (need to update potentially)
    wind_vom_cost * avg_MW_wind * hrs_per_yr * sys_plant_lifetime + #VOM of wind
    solar_vom_cost * avg_MW_sol * hrs_per_yr * sys_plant_lifetime ) # VOM of solar
     / (avg_MW*hrs_per_yr*sys_plant_lifetime) ) #total energy produced of the entire system for the life time of the system

m.Obj(LCOE)

m.options.IMODE = 6
m.solve()

print('Optimized TES storage (MWh):', tes_batt_max.value[-1])
print('System LCOE ($/MW):', LCOE.value[-1])

plt.figure(figsize=(12, 7))
plt.subplot(3, 1, 1)
plt.plot(m.time, smr_gen.value, label='Nuclear')
plt.plot(m.time, tes_in.value, label='TES')
plt.plot(m.time, turb_out.value, label='Turbine')
plt.plot(m.time, wind.value, label='Wind')
plt.plot(m.time, solar.value, label='Solar')
plt.legend()

plt.subplot(3, 1, 2)
plt.plot(m.time, load.value, label='Load')
plt.legend()

plt.subplot(3, 1, 3)
plt.plot(m.time, tes_batt.value, label='TES Storage')
#plt.plot(m.time, tes_in.value, label='TES Useage')
plt.legend()

plt.xlabel('time (hrs)')
plt.savefig(f'solution_{nhrs}_hrs.pdf')
plt.show()
