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('2018-02-27')
    endtime = pd.Timestamp('2018-02-28')
    starttime = pd.Timestamp('2018-02-27')
    endtime = pd.Timestamp('2018-02-28')
    shop.set_time_resolution(starttime=starttime, endtime=endtime, timeunit='hour')

    # Add topology
    rsv1 = shop.model.reservoir.add_object('Reservoir1')
    rsv1.max_vol.set(12)
    rsv1.lrl.set(90)
    rsv1.hrl.set(100)
    rsv1.vol_level.set(pd.Series([90, 100, 101], index=[0, 12, 14], name=0))
    rsv1.flow_descr.set(pd.Series([0, 1000], index=[100, 101], name=0))

    river1 = shop.model.river.add_object('River1')
    river1.upstream_elevation.set(90)
    
    rsv2 = shop.model.reservoir.add_object('Reservoir2')
    rsv2.max_vol.set(1)
    rsv2.lrl.set(85)
    rsv2.hrl.set(86)
    rsv2.vol_level.set(pd.Series([85, 86, 87], index=[0, 1, 2]))
    rsv2.flow_descr.set(pd.Series([0, 1000], index=[86, 87]))

    plant1 = shop.model.plant.add_object('Plant1')
    plant1.outlet_line.set(40)
    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(15)
    p1g1.p_max.set(100)
    p1g1.p_nom.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([70, 80, 95, 90], index=[10, 25, 90, 100], name=45),
                            pd.Series([72, 82, 98, 92], index=[10, 25, 90, 100], name=55)])

    # Connect objects
    rsv1.connect_to(river1)
    river1.connect_to(rsv2)
    rsv2.connect_to(plant1)

    rsv1.start_level.set(94)
    rsv2.start_level.set(85.95)
    rsv1.energy_value_input.set(38.6)
    rsv2.energy_value_input.set(38.6)

    shop.model.market.add_object('Day_ahead')
    da = shop.model.market.Day_ahead
    da.max_buy.set(9999)
    da.max_sale.set(9999)
    shop.model.market.Day_ahead.sale_price.set(pd.DataFrame([30,80,80,30,30,30,30,30,40,40,30,30,40,30,30,30,30,40,40,30,30,30,30,30],index=[starttime+ pd.Timedelta(hours=i) for i in range(0,24)]))
    shop.model.market.Day_ahead.buy_price.set(shop.model.market.Day_ahead.sale_price.get()+0.002)

    # add gate
    river1.width_depth_curve.set(pd.Series([0,5],index=[2,2],name=0))
    river1.gate_opening_curve.set(pd.Series([0,10],index=[0,1],name=0))
    river1.mip_flag.set(1)

    rsv1.inflow.set(pd.DataFrame([101, 50], index=[starttime, starttime + pd.Timedelta(hours=1)]))
    return shop

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