from pyshop import ShopSession
import pandas as pd
from pandas import Timestamp
from numpy import nan

def build_model() -> ShopSession:

    #Initialize a new ShopSession
    shop = ShopSession()

    #Set the time resolution of the optimization
    start_time = Timestamp('2018-02-27 00:00:00')
    end_time = Timestamp('2018-03-05 00:00:00')
    t = [Timestamp('2018-02-27 00:00:00')]
    y = [15.0]
    step_length = pd.Series(y,index=t)
    shop.set_time_resolution(start_time, end_time, 'minute', step_length)

    #Add all objects and set all attributes
    river_w_flow_curve = shop.model.river.add_object('river_w_flow_curve')
    river_w_flow_curve.upstream_elevation.set(100.0)

    river_w_delay_curve = shop.model.river.add_object('river_w_delay_curve')
    river_w_delay_curve.upstream_elevation.set(60.0)

    river_w_min_flow = shop.model.river.add_object('river_w_min_flow')
    river_w_min_flow.upstream_elevation.set(50.0)

    river_w_inflow = shop.model.river.add_object('river_w_inflow')
    river_w_inflow.upstream_elevation.set(60.0)

    rsv_unreg = shop.model.reservoir.add_object('rsv_unreg')
    rsv_unreg.max_vol.set(12.0)
    rsv_unreg.lrl.set(100.0)
    rsv_unreg.hrl.set(110.0)
    x = [0.0, 12.0, 14.0]
    y = [100.0, 110.0, 111.0]
    rsv_unreg.vol_head.set(pd.Series(y,index=x,name=0.0))
    rsv_unreg.start_head.set(102.0)
    rsv_unreg.inflow.set(50.0)
    rsv_unreg.energy_value_input.set(40.0)

    rsv_reg = shop.model.reservoir.add_object('rsv_reg')
    rsv_reg.max_vol.set(120.0)
    rsv_reg.lrl.set(100.0)
    rsv_reg.hrl.set(110.0)
    x = [0.0, 120.0, 124.0]
    y = [100.0, 110.0, 111.0]
    rsv_reg.vol_head.set(pd.Series(y,index=x,name=0.0))
    rsv_reg.start_head.set(102.0)
    rsv_reg.inflow.set(0.0)
    rsv_reg.energy_value_input.set(41.0)

    plant = shop.model.plant.add_object('plant')
    plant.prod_area.set(1.0)
    plant.outlet_line.set(60.0)
    plant.main_loss.set([0.0002])
    plant.penstock_loss.set([0.0001])

    market_day_ahead = shop.model.market.add_object('Day_ahead')
    market_day_ahead.market_type.set('ENERGY')
    market_day_ahead.max_buy.set(9999.0)
    market_day_ahead.max_sale.set(9999.0)
    market_day_ahead.buy_price.set(40.01)
    market_day_ahead.sale_price.set(39.99)

    g1 = shop.model.generator.add_object('plant_G1')
    g1.penstock.set(1)
    g1.p_min.set(5.0)
    g1.p_max.set(20.0)
    g1.p_nom.set(20.0)
    x = [0.0, 20.0]
    y = [95.0, 98.0]
    g1.gen_eff_curve.set(pd.Series(y,index=x,name=0.0))
    xy_curves = []
    x = [15.0, 48.3, 57.0]
    y = [80.0, 95.0, 90.0]
    xy_curves.append(pd.Series(y,index=x,name=40.0))
    x = [12.0, 37.5, 45.0]
    y = [82.0, 98.0, 92.0]
    xy_curves.append(pd.Series(y,index=x,name=50.0))
    g1.turb_eff_curves.set(xy_curves)
    g1.startcost.set(500.0)


    #Connect all objects
    river_w_flow_curve.connect_to(river_w_min_flow)

    river_w_delay_curve.connect_to(river_w_min_flow)

    river_w_inflow.connect_to(river_w_min_flow)

    rsv_unreg.connect_to(river_w_flow_curve)

    rsv_reg.connect_to(plant)

    plant.connect_to(river_w_delay_curve)
    plant.connect_to(g1)

    return shop

def run_commands(shop:ShopSession) -> None:
    
    shop.set_universal_mip("off",[])
    shop.start_sim([],4)
    shop.set_universal_mip("on",[])
    shop.start_sim([],2)
    shop.set_code("incremental",[])
    shop.start_sim([],5)