from pyshop import ShopSession
import pandas as pd
import numpy as np

def build_model() -> ShopSession:

    shop = ShopSession()

    shop.set_time_resolution(pd.Timestamp("2024-01-01"), pd.Timestamp("2024-01-02"), "hour")

    rsv = shop.model.reservoir.add_object("Reservoir")
    rsv.lrl.set(80)
    rsv.hrl.set(100)
    rsv.max_vol.set(50)
    rsv.vol_head.set(pd.Series([80,85,90,95,100,102], index=[0,5,12,20,50,60]))

    creek1 = shop.model.reservoir.add_object("Creek1")
    creek1.lrl.set(80)
    creek1.hrl.set(110)
    creek1.max_vol.set(0.001)
    creek1.vol_head.set(pd.Series([80,110,140], index=[0,0.001,0.002]))

    creek2 = shop.model.reservoir.add_object("Creek2")
    creek2.lrl.set(80)
    creek2.hrl.set(105)
    creek2.max_vol.set(0.001)
    creek2.vol_head.set(pd.Series([80,105,130], index=[0,0.001,0.002]))

    creek3 = shop.model.reservoir.add_object("Creek3")
    creek3.lrl.set(80)
    creek3.hrl.set(101)
    creek3.max_vol.set(0.001)
    creek3.vol_head.set(pd.Series([80,101,121], index=[0,0.001,0.002]))

    tunnel1 = shop.model.tunnel.add_object("Rsv_N1")
    tunnel1.start_height.set(80)
    tunnel1.end_height.set(80)
    tunnel1.loss_factor.set(0.00015)

    tunnel2 = shop.model.tunnel.add_object("N1_N2")
    tunnel2.start_height.set(80)
    tunnel2.end_height.set(80)
    tunnel2.loss_factor.set(0.000025)

    tunnel3 = shop.model.tunnel.add_object("N2_N3")
    tunnel3.start_height.set(80)
    tunnel3.end_height.set(80)
    tunnel3.loss_factor.set(0.000025)

    tunnel4 = shop.model.tunnel.add_object("N3_plant")
    tunnel4.start_height.set(80)
    tunnel4.end_height.set(0)
    tunnel4.loss_factor.set(0.0008)

    tunnelC1 = shop.model.tunnel.add_object("Creek1_N1")
    tunnelC1.start_height.set(80)
    tunnelC1.end_height.set(80)
    tunnelC1.loss_factor.set(1e-6)

    tunnelC2 = shop.model.tunnel.add_object("Creek2_N2")
    tunnelC2.start_height.set(80)
    tunnelC2.end_height.set(80)
    tunnelC2.loss_factor.set(1e-6)

    tunnelC3 = shop.model.tunnel.add_object("Creek3_N3")
    tunnelC3.start_height.set(80)
    tunnelC3.end_height.set(80)
    tunnelC3.loss_factor.set(1e-6)

    river1 = shop.model.river.add_object("Spill_rsv")
    river1.upstream_elevation.set(100)
    river1.up_head_flow_curve.set([pd.Series([0,50],index=[100,102], name=0)])
    river1.flow_cost.set(1000)

    river2 = shop.model.river.add_object("Spill_C1")
    river2.upstream_elevation.set(110)
    river2.up_head_flow_curve.set([pd.Series([0,50],index=[110,112], name=0)])

    river3 = shop.model.river.add_object("Spill_C2")
    river3.upstream_elevation.set(105)
    river3.up_head_flow_curve.set([pd.Series([0,50],index=[105,107], name=0)])

    river4 = shop.model.river.add_object("Spill_C3")
    river4.upstream_elevation.set(101)
    river4.up_head_flow_curve.set([pd.Series([0,50],index=[101,102], name=0)])

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

    q = [5,7.5,10,12.5,15,17.5,20,22.5,25,27.5,30,32.5,35,37.5,40,42.5,45,47.5,50]
    eff = [80.0, 82.15033783783784, 84.12612612612612, 85.92736486486487, 87.55405405405406, 89.00619369369369, 90.28378378378379, 91.38682432432432, 92.31531531531532, 93.06925675675676, 93.64864864864865, 94.053490990991, 94.28378378378379, 94.33952702702703, 94.22072072072072, 93.92736486486487, 93.45945945945945, 92.81700450450451, 92.0]

    gen1 = shop.model.generator.add_object("G1")
    gen1.penstock.set(1)
    gen1.p_min.set(5)
    gen1.p_max.set(35)
    gen1.p_nom.set(35)
    gen1.gen_eff_curve.set(pd.Series([95,95],index=[5,35]))
    gen1.turb_eff_curves.set([pd.Series(eff,index=q, name=70), pd.Series(np.array(eff)+2,index=q,name=100)])
    gen1.startcost.set(500)

    gen2 = shop.model.generator.add_object("G2")
    gen2.penstock.set(2)
    gen2.p_min.set(5)
    gen2.p_max.set(35)
    gen2.p_nom.set(35)
    gen2.gen_eff_curve.set(pd.Series([95,95],index=[5,35]))
    gen2.turb_eff_curves.set([pd.Series(eff,index=q, name=70), pd.Series(np.array(eff)+2,index=q,name=100)])
    gen2.startcost.set(500)

    rsv.connect_to(tunnel1)
    tunnel1.connect_to(tunnel2)
    tunnel2.connect_to(tunnel3)
    tunnel3.connect_to(tunnel4)
    tunnel4.connect_to(plant)

    creek1.connect_to(tunnelC1)
    tunnelC1.connect_to(tunnel2)

    creek2.connect_to(tunnelC2)
    tunnelC2.connect_to(tunnel3)

    creek3.connect_to(tunnelC3)
    tunnelC3.connect_to(tunnel4)

    rsv.connect_to(river1)
    creek1.connect_to(river2)
    creek2.connect_to(river3)
    creek3.connect_to(river4)


    plant.connect_to(gen1)
    plant.connect_to(gen2)

    market = shop.model.market.add_object("spot")

    t = [pd.Timestamp("2024-01-01") + pd.Timedelta(hours=i) for i in range(24)]
    p = [29.56,28.46,26.66,24.48,24.01,21.23,22.62,25.04,26.24,32.21,41.34,43.51,43.02,44.29,46.24,50.61,59.47,64.99,61.74,55.07,48.01,44.01,45.20,38.00]
    market.sale_price.set(pd.Series(p, index=t))
    market.max_sale.set(10000)

    for r in shop.model.reservoir:
        r.energy_value_input.set(40)

    return shop