import pandas as pd
from pyshop import ShopSession

def build_model():
    
    shop = ShopSession()
    starttime = pd.Timestamp('2022-04-16')
    endtime = pd.Timestamp('2022-04-19')
    shop.set_time_resolution(starttime=starttime, endtime=endtime, timeunit='hour')    
    
    rsv1 = shop.model.reservoir.add_object("rsv1")
    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))

    spill1 = shop.model.river.add_object("spill_rsv1")
    spill1.upstream_elevation.set(100)
    spill1.up_head_flow_curve.set([pd.Series([0,1000],index=[100,101], name=0)])
    spill1.flow_cost.set(100)
    rsv1.connect_to(spill1)
    
    rsv2 = shop.model.reservoir.add_object("rsv2")
    rsv2.max_vol.set(12)
    rsv2.lrl.set(90)
    rsv2.hrl.set(100)
    rsv2.vol_head.set(pd.Series([90, 100, 101], index=[0, 12, 14], name=0))

    spill2 = shop.model.river.add_object("spill_rsv2")
    spill2.upstream_elevation.set(100)
    spill2.up_head_flow_curve.set([pd.Series([0,1000],index=[100,101], name=0)])
    spill2.flow_cost.set(100)
    rsv2.connect_to(spill2)
    
    rsv3 = shop.model.reservoir.add_object("rsv3")
    rsv3.max_vol.set(12)
    rsv3.lrl.set(40)
    rsv3.hrl.set(50)
    rsv3.vol_head.set(pd.Series([40, 50, 51], index=[0, 12, 14], name=0))

    spill3 = shop.model.river.add_object("spill_rsv3")
    spill3.upstream_elevation.set(50)
    spill3.up_head_flow_curve.set([pd.Series([0,1000],index=[50,51], name=0)])
    spill3.flow_cost.set(100)
    rsv3.connect_to(spill3)
    
    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(10)
    p1g1.p_max.set(100)
    p1g1.p_nom.set(100)
    p1g1.startcost.set(500)
    p1g1.turb_eff_curves.set([pd.Series([77,91,94.75,95,94,93.4], index=[15,60,90,100,120,125], name=90),
                              pd.Series([78,92,95.75,96,95,94.4], index=[15,60,90,100,120,125], name=100)])

    p1g2 = shop.model.generator.add_object("plant1_g2")
    plant1.connect_to(p1g2)
    p1g2.penstock.set(1)
    p1g2.p_min.set(10)
    p1g2.p_max.set(100)
    p1g2.p_nom.set(100)
    p1g2.startcost.set(500)
    p1g2.turb_eff_curves.set([pd.Series([77,91,94.75,95,94,93.4], index=[15,60,90,100,120,125], name=90),
                              pd.Series([78,92,95.75,96,95,94.4], index=[15,60,90,100,120,125], name=100)])

    plant2 = shop.model.plant.add_object("plant2")
    plant2.outlet_line.set(40)
    plant2.main_loss.set([0.0002])
    plant2.penstock_loss.set([0.0001])

    p2p1 = shop.model.pump.add_object("plant2_p1")
    plant2.connect_to(p2p1)
    p2p1.penstock.set(1)
    p2p1.p_min.set(12)
    p2p1.p_max.set(50)
    p2p1.p_nom.set(50)
    p2p1.startcost.set(500)
    p2p1.turb_eff_curves.set([pd.Series([80, 95, 90], index=[25, 90, 100], name=40),
                              pd.Series([82, 97, 92], index=[25, 90, 100], name=50)])    
    
    t1 = shop.model.tunnel.add_object("rsv1_plant1")
    t1.start_height.set(89)
    t1.end_height.set(89)
    t1.loss_factor.set(0.002)

    t2 = shop.model.tunnel.add_object("node_plant1")
    t2.start_height.set(89)
    t2.end_height.set(89)
    t2.loss_factor.set(0.004)

    t3 = shop.model.tunnel.add_object("rsv2_node")
    t3.start_height.set(89)
    t3.end_height.set(89)
    t3.loss_factor.set(0.001)

    t4 = shop.model.tunnel.add_object("node_plant2")
    t4.start_height.set(89)
    t4.end_height.set(89)
    t4.loss_factor.set(0.001)
    
    price = pd.Series(
        [189.13, 187.4, 184.28, 179.79, 178.55, 184.04, 184.96, 185.19, 184.43, 178.41, 172.79, 170, 159.04, 151.02, 145.77, 140.71, 91.15, 151.32, 173.19, 184.14, 187.79, 187.2, 182.7, 172.4, 166.77, 158.86, 148.36, 144.25, 144.55, 144.58, 150.93, 148.77, 153.6, 159.3, 159.86, 161.08, 150.01, 127.22, 97.25, 70.03, 93.64, 107.47, 166.74, 182.45, 185.41, 185.42, 179.99, 163.3, 108.13, 83.68, 86.78, 83.68, 88.07, 95.26, 131.01, 148.91, 144.15, 151.97, 158.76, 153.21, 151.31, 130.02, 50.02, 77.61, 154.46, 175.8, 179.13, 186.71, 194.31, 202.31, 201.25, 195.54, 192.99, 187.38, 185.6, 185.41, 185.97, 192.53, 196.88, 204.7, 285.98, 247.53, 202.89, 193.53, 189.69, 185.28, 182.22, 181.75, 184.18, 189.16, 196.09, 203.57, 206.68, 206.65, 191.96, 185.45],
        index=[starttime+pd.Timedelta(hours=t) for t in range(24*4)]
    )
    price_mean = price.mean()
    
    rsv1.start_head.set(95)
    rsv2.start_head.set(95)
    rsv3.start_head.set(45)
    
    rsv1.energy_value_input.set(0.9*price_mean)
    rsv2.energy_value_input.set(0.8*price_mean)
    rsv3.energy_value_input.set(0.5*price_mean)
    
    spot = shop.model.market.add_object("spot")
    spot.buy_price.set(price + 0.01)
    spot.sale_price.set(price - 0.01)
    spot.max_sale.set(9999)
    spot.max_buy.set(9999)

    rsv3.inflow.set(20)    
    
    rsv1.connect_to(t1)
    t2.connect_to(plant1)
    t1.connect_to(plant1)

    rsv2.connect_to(t3)
    t3.connect_to(t2)
    t3.connect_to(t4)
    t4.connect_to(plant2)
    plant2.connect_to(rsv3)    
    
    return shop