import pandas as pd
from pyshop import ShopSession

def build_model(shop:ShopSession) -> None:
    starttime = pd.Timestamp('2018-02-26')
    endtime = pd.Timestamp('2018-02-27 13:00:00')
    shop.set_time_resolution(starttime=starttime, endtime=endtime, timeunit='hour')

    rsv1 = shop.model.reservoir.add_object('Reservoir1')
    rsv1.max_vol.set(12)
    rsv1.lrl.set(90)
    rsv1.hrl.set(100)
    rsv1.vol_head.set(pd.Series([90, 100, 101], index=[0, 12, 14], name=0))
    rsv1.flow_descr.set(pd.Series([0, 1000], index=[100, 101], name=0))

    plant1 = shop.model.plant.add_object('Plant1')
    plant1.outlet_line.set(0)
    plant1.main_loss.set([0.0002])
    plant1.penstock_loss.set([0.0001])

    p1g1 = shop.model.generator.add_object('Plant1_G1')
    plant1.connect_to(p1g1)
    p1g1.penstock.set(1)
    p1g1.p_min.set(0)
    p1g1.p_max.set(100)
    p1g1.p_nom.set(100)
    p1g1.startcost.set(500)
    p1g1.gen_eff_curve.set(pd.Series([95, 98], index=[0, 100]))
    p1g1.turb_eff_curves.set([pd.Series([65, 75, 90, 93, 89], index=[10, 25, 70, 90, 100], name=85),
                              pd.Series([72, 81, 94, 96, 93], index=[10, 25, 70, 90, 100], name=90),
                              pd.Series([74, 85, 97, 99, 96], index=[10, 25, 70, 90, 100], name=95),
                            pd.Series([73, 83, 96, 97.5, 95], index=[10, 25, 70, 90, 100], name=100)])
    # Connect objects
    rsv1.connect_to(plant1)

    rsv1.energy_value_input.set(39.7)
    rsv1.start_head.set(100)

    #Add market
    shop.model.market.add_object('Day_ahead')
    da = shop.model.market.Day_ahead
    da.sale_price.set(49.99)
    da.buy_price.set(da.sale_price.get()+0.02)
    da.max_buy.set(9999)
    da.max_sale.set(9999)

    starttime = shop.get_time_resolution()['starttime']
    rsv1.inflow.set(pd.DataFrame([0, 0], index=[starttime, starttime + pd.Timedelta(hours=1)]))

def run_model(shop:ShopSession) -> None:
    shop.save_pq_curves(['on'], [])
    shop.start_sim([], [3])
    shop.set_code(['incremental'], [])
    shop.start_sim([], [3])