import pandas as pd
import plotly.graph_objs as go
from plotly.subplots import make_subplots
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_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(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_level.set(100)
    rsv1.energy_value_input.set(10)

    da=shop.model.market.add_object('Day_Ahead')
    da.sale_price.set(pd.DataFrame([10,12,12,5,5,12,10,10.5,15,11,11.7,11.7,12,15,8,8.5,5,12,11,11,15,15,15,15],index=[starttime+ pd.Timedelta(hours=i) for i in range(0,24)]))
    da.buy_price.set(shop.model.market.Day_Ahead.sale_price.get()+0.002)
    da.max_sale.set(9999)
    da.max_buy.set(9999)

    #Add first needle_combination, valid from 7 to 20 MW
    n1 = shop.model.needle_combination.add_object("N1")
    n1.p_min.set(5)
    n1.p_max.set(20)
    n1.p_nom.set(40)

    q = [5.0, 30.0, 49.0]
    eff = [90, 95, 92]
    c1 = pd.Series(eff,index=q,name=45)

    n1.turb_eff_curves.set([c1])
    pelton.connect_to(n1)

    #Add second needle_combination, valid from 25 MW to 40 MW
    n2 = shop.model.needle_combination.add_object("N2")
    n2.p_min.set(25)
    n2.p_max.set(40)
    n2.p_nom.set(40)

    q = [50.0, 68.0, 110.0]
    eff = [82, 92, 87]
    c1 = pd.Series(eff,index=q,name=45)

    n2.turb_eff_curves.set([c1])
    pelton.connect_to(n2)

    return shop

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


def plot_production_and_frr_delivery(p, frr, p1min, p1max, p2min, p2max):

    # Create figure for plotting
    fig = make_subplots()
    fig.update_layout(title="Production and FRR up delivery")
    fig.update_yaxes(title_text="Production [MW]")
    fig.update_yaxes(range = [0,p2max+1])

    # Define arrays for plotting needle combination limits
    t = p.index
    p1min_arr = np.zeros((len(t))) + p1min
    p1max_arr = np.zeros((len(t))) + p1max
    p2min_arr = np.zeros((len(t))) + p2min
    p2max_arr = np.zeros((len(t))) + p2max

    # Color forbidden area between needle combinations
    fig.add_trace(go.Scatter(x=t, y=p1max_arr,line={'width': 0}, showlegend=False))
    fig.add_trace(go.Scatter(name="Forbidden zone",x=t, y=p2min_arr ,fill='tonexty',line={'color': "grey", 'width': 0},line_shape='hv'))

    # Color forbidden area between zero and first needle combination
    fig.add_trace(go.Scatter(x=p.index, y=(np.zeros((len(p.index))) + p1min) , fill = "tozeroy",mode="lines",showlegend=False, line = dict(color = 'grey', width = 0)))

    # Plot needle combination limits
    fig.add_trace(go.Scatter(name = "pmin1", x=t, y=p1min_arr, mode="lines", line = dict(color = 'steelblue', dash = "dash")))
    fig.add_trace(go.Scatter(name = "pmax1", x=t, y=p1max_arr ,mode="lines", line = dict(color = 'steelblue', dash = "dot")))
    fig.add_trace(go.Scatter(name = "pmin2", x=t, y=p2min_arr ,mode="lines", line = dict(color = 'midnightblue', dash = "dash")))
    fig.add_trace(go.Scatter(name = "pmax2",x=t, y=p2max_arr ,mode="lines", line = dict(color = 'midnightblue', dash = "dot")))

    # Plot vertical line between production and frr delivery
    for i, timestep in enumerate(t):
        fig.add_trace(go.Scatter(x=[timestep,timestep], y=[p.values[i], p.values[i] + frr.values[i]],mode="lines",showlegend=False, line = dict(color = 'grey')))

    # Add generator production and FRR up delivery results
    fig.add_trace(go.Scatter(x=p.index, y=p.values, name=f'Production', marker=dict( symbol = "line-ew", size = 10, line=dict(width=5, color = "indianred")), mode="markers"))
    fig.add_trace(go.Scatter(x=frr.index, y=p.values + frr.values, name=f'FRR up delivery', marker=dict( symbol = "line-ew", size = 10, line=dict(width=5, color = "seagreen")), mode="markers"))

    fig.show()