import pandas as pd
import numpy as np

def build_model(shop):
    starttime=pd.Timestamp("2024-01-01 00:00:00")
    endtime=pd.Timestamp("2024-01-02 00:00:00")
    shop.set_time_resolution(starttime=starttime, endtime=endtime, timeunit="hour")

    rsv1=shop.model.reservoir.add_object("Reservoir1")
    rsv1.hrl.set(100)
    rsv1.lrl.set(90)
    rsv1.max_vol.set(12)
    rsv1.vol_head.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(50)

    pelton=shop.model.generator.add_object("pelton")
    pelton.connect_to(plant1)
    pelton.penstock.set(1)
    pelton.startcost.set(500)
    pelton.gen_eff_curve.set(pd.Series([98,98], index=[0,100]))

    rsv1.connect_to(plant1)

    rsv1.start_head.set(95)
    rsv1.inflow.set(10)
    rsv1.energy_value_input.set(40)

    da=shop.model.market.add_object('Day_Ahead')
    t = [starttime+pd.Timedelta(hours=i) for i in range(24)]
    price = np.array([25+i for i in range(24)])
    da.sale_price.set(pd.Series(price,index=t))
    da.buy_price.set(pd.Series(price+0.1,index=t))
    da.max_sale.set(9999)
    da.max_buy.set(9999)
        
def run_model(shop):
    shop.start_sim([],['3'])
    shop.set_code(['inc'],[])
    shop.start_sim([],['3'])

def calc_p(curve:pd.Series):
    q = np.array(curve.index)
    eff = curve.values
    h = float(curve.name)
    p = 0.001*9.81*h*q*eff*0.01*0.98
    return p