import pandas as pd
import plotly.graph_objs as go


def build_model(shop):

    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_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(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_head.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.level_ramp_penalty_cost.set(1000000)
    

def run_model(shop):
    shop.set_universal_mip(['on'], [])
    shop.start_sim([],['3'])
    shop.set_code(['inc'],[])
    shop.start_sim([],['3'])


def plot_period_levels(starttime, rsv, offset, interval_length,lrl, title):

    fig = go.Figure()

    #Plot reservoir level
    head = rsv.head.get()
    fig.add_trace(go.Scatter(x=head.index, y=head.values, name="Reservoir level"))
    
    #plot historical reservoir level
    hist_head = rsv.historical_level.get()
    hist_head = hist_head[hist_head.index<=starttime]
    fig.add_trace(go.Scatter(x=hist_head.index, y=hist_head.values, mode = "lines", name="Historical level", line = dict(color = 'blue', dash = "dash")))

    y_min = lrl
    y_max = max(head)+1

    for i in range(7): 
        #The start and end time of each fixed  period
        if i == 0:
            period_start = hist_head.index[0]
        else: 
            period_start = starttime + pd.Timedelta(minutes=offset) + pd.Timedelta(minutes=(i-1)*interval_length) 
        period_end = period_start + pd.Timedelta(minutes=interval_length)

        #For first period, find max and min values taking historical level into account
        if i == 0:
            head_max = max(max(head[head.index <= period_end].values), max(hist_head.values))
            head_min = min(min(head[head.index <= period_end].values), min(hist_head.values))
        #Find period min and max level
        else:
            h = head[head.index >= period_start]
            h = h[h.index <= period_end]
            head_max = max(h.values)
            head_min = min(h.values)

        #Find max difference in level within period
        diff = head_max - head_min
        diff_txt = "{:.2f}".format(diff)

        #Plot min and max level within period
        t = [period_start,period_end]    
        max_period_level = [head_max, head_max]
        min_period_level = [head_min, head_min]
        fig.add_trace(go.Scatter(x=t, y=max_period_level ,mode="lines",showlegend=False, line = dict(color = 'darkred', dash = "dash")))
        fig.add_trace(go.Scatter(x=t, y=min_period_level ,mode="lines",showlegend=False, line = dict(color = 'darkred', dash = 'dot')))

        #Plot max difference
        fig.add_trace(go.Scatter(x=[period_start+pd.Timedelta(hours = 1),period_start+pd.Timedelta(hours = 1)], y=[head_min, head_max],mode="lines",showlegend=False, line = dict(color = 'red')))
        fig.add_trace(go.Scatter(x=[period_start+pd.Timedelta(minutes = interval_length/2)],y=[head_min - 0.25],mode="text",text=[f'Δh = {diff_txt}'],textposition="bottom center",textfont=dict(color="red"),showlegend=False))
        
        #Plot period start vertical line
        fig.add_trace(go.Scatter(x=[period_start,period_start], y=[y_min,y_max],mode="lines",showlegend=False, line = dict(color = 'black', dash = 'dot')))

    #Add legends  
    fig.add_trace(go.Scatter(x=t, y=max_period_level ,mode="lines", line = dict(color = 'darkred', dash = "dash"), name = "Max period level"))
    fig.add_trace(go.Scatter(x=t, y=min_period_level, mode="lines", line = dict(color = 'darkred', dash = 'dot'), name = "Min period level"))
    fig.add_trace(go.Scatter(x=[period_start,period_start], y=[y_min,y_max],mode="lines", line = dict(color = 'black', dash = 'dot'), name = "Period start"))

    #Plot optimization start vertical line
    fig.add_trace(go.Scatter(x=[starttime,starttime], y=[y_min,y_max],mode="lines", line = dict(color = 'grey', dash = 'dot'), name = "Optimization start"))

    fig.update_layout(title=title, xaxis_title="Time", yaxis_title="masl")
    fig.show()