"""
The tsnet.simulation.main module contains function to perform
the workflow of read, discretize, initial, and transient
simulation for the given .inp file.

"""


import folium
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

def plot_scatter(plot_data, x_col, y_col, color_col, figsize=(4, 3), color_map=None, s=0.5, xlabel='longitude', ylabel='latitude', colorbar_label='pf', save_path=None):
    """
    Plots a scatter plot from the given DataFrame with specified x, y, and color columns.

    Parameters:
    plot_data (pd.DataFrame): DataFrame containing the data
    x_col (str): Column name for x coordinates
    y_col (str): Column name for y coordinates
    color_col (str): Column name for coloring the scatter points
    figsize (tuple): Size of the figure (width, height)
    color_map: Colormap to use for the scatter plot (default is a predefined colormap)
    s (float): Marker size
    xlabel (str): Label for the x-axis
    ylabel (str): Label for the y-axis
    colorbar_label (str): Label for the color bar
    save_path (str): Path to save the figure (default is None, meaning not saving the figure)
    """
    if color_map is None:
        color_map = ListedColormap(['blue', 'green', 'yellow', 'orange', 'red'])

    fig = plt.figure(figsize=figsize)
    scatter = plt.scatter(plot_data[x_col], plot_data[y_col], c=plot_data[color_col], cmap=color_map, s=s)
    cbar = plt.colorbar(scatter)
    cbar.set_label(colorbar_label)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    if save_path:
        plt.savefig(save_path,bbox_inches='tight', dpi=300)
    plt.show()


"""
The tsnet.simulation.initialize contains functions to
1. Initialize the list containing numpy arrays for velocity and head.
2. Calculate initial conditions using Epanet engine.
3. Calculate D-W coefficients based on initial conditions.
4. Calculate demand coefficients based on initial conditions.

"""

import wntr
import numpy as np
import warnings
from tsnet.utils import calc_parabola_vertex

def Initializer(tm, t0, engine='DD'):
    """Initial Condition Calculation.

    Initialize the list containing numpy arrays for velocity and head.
    Calculate initial conditions using Epanet engine.
    Calculate D-W coefficients based on initial conditions.
    Calculate demand coefficients based on initial conditions.

    Parameters
    ----------
    tm : tsnet.network.geometry.TransientModel
        Simulated network
    t0 : float
        time to calculate initial condition
    engine : string
        steady state calculation engine:
        DD: demand driven;
        PDD: pressure dependent demand,
        by default DD


    Returns
    -------
    tm : tsnet.network.geometry.TransientModel
        Network with updated parameters
    """
    # adjust the time step and discretize each pipe

    tn = int(tm.simulation_period/tm.time_step) # Total time steps
    print ('Total Time Step in this simulation %s'  %tn)

    # create new attributes for each pipe to store head and velocity results
    # at its start and end node.
    for _, pipe in tm.pipes():
        pipe.start_node_head = np.zeros(tn)
        pipe.start_node_velocity = np.zeros(tn)
        pipe.start_node_flowrate = np.zeros(tn)
        pipe.end_node_head = np.zeros(tn)
        pipe.end_node_velocity = np.zeros(tn)
        pipe.end_node_flowrate = np.zeros(tn)

    # create new attributes for each node to store head and discharge results
    for _,node in tm.nodes():
        node.demand_discharge = np.zeros(tn)
        node.emitter_discharge = np.zeros(tn)

    # calculate initial conditions using EPAnet engine
    for _,node in tm.nodes():
        if node._leak_status == True:
            node.add_leak(tm, area=node.emitter_coeff/np.sqrt(2*9.81),
                    discharge_coeff = 1, start_time = t0)
    if engine.lower() == 'dd':
        tm.options.hydraulic.demand_model = 'DD'
        sim = wntr.sim.EpanetSimulator(tm)
        results = sim.run_sim()
    elif engine.lower() == 'pdd':
        tm.options.hydraulic.demand_model = 'PDD'
        sim = wntr.sim.EpanetSimulator(tm)
        results = sim.run_sim()
    else:
        raise Exception("Unknown initial calculation engine. \
            The engine can only be 'DD' or 'PDD'.")

    for _,node in tm.nodes():
        node.initial_head = results.node['head'].loc[t0, node.name]

    for _,link in tm.links():
        link.initial_flow = results.link['flowrate'].loc[t0, link.name]
    nu = 1.004e-6
    for _, pipe in tm.pipes():
        # assign the initial conditions to the latest result arrays

        V = np.sign(results.link['flowrate'].loc[t0, pipe.name])*\
                    results.link['velocity'].loc[t0, pipe.name]*\
                    np.ones(pipe.number_of_segments+1)

        H = [results.node['head'].loc[t0, pipe.start_node_name] +\
                      i* ((results.node['head'].loc[t0, pipe.end_node_name]-
                       results.node['head'].loc[t0, pipe.start_node_name])/
                       (pipe.number_of_segments))
                    for i in range(pipe.number_of_segments+1)]

        H = np.array(H)
        pipe.initial_head = H
        pipe.initial_velocity = V
        pipe.initial_Re =  np.abs(V[0]*pipe.diameter/nu)

        # assign the initial conditions to the results attributes
        pipe.start_node_velocity[0] = V[0]
        pipe.end_node_velocity[0] = V[-1]
        pipe.start_node_head[0] = H[0]
        pipe.end_node_head[0] = H[-1]
        pipe.start_node_flowrate[0] = V[0]*pipe.area
        pipe.end_node_flowrate[0] = V[-1]*pipe.area

        # calculate demand coefficient
        Hs = H[0]
        He = H[-1]
        demand = [0,0] # demand at start and end node
        try :
            demand[0] = results.node['demand'].loc[t0, pipe.start_node_name]
            # demand[0] = tm.nodes[pipe.start_node_name].demand_timeseries_list.at(t0)
        except:
            demand[0] = 0.
        try :
            demand[1] = results.node['demand'].loc[t0, pipe.end_node_name]
            # demand[1] = tm.nodes[pipe.end_node_name].demand_timeseries_list.at(t0)
        except :
            demand[1] = 0.
        try :
            Hsa = H[0] -pipe.start_node.elevation
        except:
            Hsa = 1.
        try :
            Hea = H[-1]- pipe.end_node.elevation
        except :
            Hea=1.
        pipe = cal_demand_coef(demand, pipe, Hsa, Hea, t0)

        # calculate demnad discharge and emitter discharge
        if pipe.start_node.node_type == 'Junction':
            pipe.start_node.emitter_discharge[0] = pipe.start_node.emitter_coeff * np.sqrt(Hsa)
            pipe.start_node.demand_discharge[0] = pipe.start_node.demand_coeff * np.sqrt(Hsa)
        if pipe.end_node.node_type == 'Junction':
            pipe.end_node.emitter_discharge[0] = pipe.end_node.emitter_coeff * np.sqrt(Hea)
            pipe.end_node.demand_discharge[0] =  pipe.end_node.demand_coeff * np.sqrt(Hea)

        # calculate roughness coefficient
        Vp = V[0]
        hl = abs(Hs - He)
        pipe = cal_roughness_coef(pipe, Vp, hl )

    # set initial conditions as a new attribute to TransientModel
    tm.initial_head = H
    tm.initial_velocity = V

    # add pump operation points
    tm = pump_operation_points(tm)

    return tm

def cal_demand_coef(demand, pipe, Hs, He, t0=0.):

    """Calculate the demand coefficient for the start and end node of the pipe.

    Parameters
    ----------
    demand : list
        Demand at the start (demand[0]) and end demand[1] node
    pipe : object
        Pipe object
    Hs : float
        Head at the start node
    He : float
        Head at the end node
    t0 : float, optional
        Time to start initial condition calculation, by default 0

    Returns
    -------
    pipe : object
        Pipe object with calculated demand coefficient
    """

    try:
        start_demand_coeff = demand[0]/ np.sqrt(Hs)
    except :
        start_demand_coeff = 0.

    try:
        end_demand_coeff = demand[1] / np.sqrt(He)
    except :
        end_demand_coeff = 0.
    pipe.start_node.demand_coeff = start_demand_coeff # [m^3/s/(m H20)^(1/2)]
    pipe.end_node.demand_coeff = end_demand_coeff   # [m^3/s/(m H20)^(1/2)]
    return pipe

def cal_roughness_coef(pipe, V, hl):
    """Calculate the D-W roughness coefficient based on initial conditions.

    Parameters
    ----------
    pipe : object
        Pipe object
    V : float
        Initial flow velocity in the pipe
    hl : float
        Initial head loss in the pipe

    Returns
    -------
    pipe : object
        Pipe object with calculated D-W roughness coefficient.
    """

    g = 9.8
    H_tol = 1e-3
    V_tol = 1e-5

    if abs(V) >= V_tol and hl >= H_tol:
        pipe.roughness = hl / (pipe.length/pipe.diameter) / (V**2/2/g)
    else:
        pipe.roughness = 0

    if pipe.roughness >0.08:
        warnings.warn("%s :the friction coefficient %.4f is too large. \
                        The D-W coeff has been set to 0.03 "
                        %(pipe.name, pipe.roughness))
        pipe.roughness = 0.03
    if pipe.roughness!= 0:
        pipe.roughness_height = max(10**(-1/1.8/np.sqrt(pipe.roughness)) - 6.9/pipe.initial_Re, 0)
    else:
        pipe.roughness_height = 0
    return pipe

def pump_operation_points(tm):
    #add operation points to the pump
    for _, pump in tm.pumps():
        opt_point = (pump.initial_flow, abs(pump.end_node.initial_head-pump.start_node.initial_head))
        def_points = pump.get_pump_curve().points
        # single-point pump curve
        if len(def_points) == 1:
            (flow, head) = def_points[0]
            def_points.append((0., 1.33*head))
            def_points.append((2*flow, 0.))
        elif len(def_points) != 3:
            raise Exception("TSNet only support one-point or three-point pump curve.")

        dist = []
        for (i,j) in def_points:
            dist.append(np.sqrt((i - opt_point[0])**2 + (j - opt_point[1])**2))

        pump.get_pump_curve().points.remove(def_points[dist.index(min(dist))])
        pump.get_pump_curve().points.append(opt_point)

        pump.curve_coef = calc_parabola_vertex(pump.get_pump_curve().points)

    return tm
