from pyshop import ShopSession
import pandas as pd

def build_model() -> ShopSession:

    shop = ShopSession()
    starttime=pd.Timestamp("2023-01-01 00:00:00")
    endtime=pd.Timestamp("2023-01-02 00:00:00")
    shop.set_time_resolution(starttime=starttime, endtime=endtime, timeunit="minute", timeresolution=pd.Series([60],index=[starttime]))

    rsv1=shop.model.reservoir.add_object("Reservoir1")
    rsv1.hrl.set(100)
    rsv1.lrl.set(90)
    rsv1.max_vol.set(12)
    rsv1.vol_level.set(pd.Series([90,100,101],index=[0,12,14]))

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

    p1g1=shop.model.generator.add_object("P1G1")
    p1g1.connect_to(plant1)
    p1g1.penstock.set(1)
    p1g1.p_min.set(10)
    p1g1.p_nom.set(100)
    p1g1.p_max.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([80,95,90],index=[25,90,100],name=90),pd.Series([82,98,92],index=[25,90,100],name=100)])

    rsv1.connect_to(plant1)
   
    rsv1.start_level.set(95)
    rsv1.inflow.set(10)
    rsv1.energy_value_input.set(40)

    da=shop.model.market.add_object('Day_Ahead')
    da.sale_price.set(pd.Series([40 + i for i in range(24)],index=[starttime+pd.Timedelta(hours=i) for i in range(24)]))
    da.buy_price.set(da.sale_price.get()+0.1)
    da.max_sale.set(9999)
    da.max_buy.set(9999)

    gs = shop.model.global_settings.global_settings
    gs.plant_discharge_ramp_penalty_cost.set(1000000)

    return shop
 
def run_model(shop:ShopSession) -> None:
    
    shop.set_universal_mip("on",[])
    shop.start_sim([],3)
    shop.set_code("incremental",[])
    shop.start_sim([],5)