"""
The wntr.metrics.hydraulic module contains hydraulic metrics.

.. rubric:: Contents

.. autosummary::

    expected_demand
    average_expected_demand
    water_service_availability
    todini_index
    entropy

"""
import wntr.network
from wntr.network.graph import _all_simple_paths
import numpy as np
import pandas as pd
import networkx as nx
import math
from collections import Counter
import sys
if sys.version_info >= (3,0):
    from functools import reduce
    
import logging

logger = logging.getLogger(__name__)

def expected_demand(wn, start_time=None, end_time=None, timestep=None):
    """
    Compute expected demand at each junction and time using base demands
    and demand patterns along with the demand multiplier
    
    Parameters
    -----------
    wn : wntr WaterNetworkModel
        Water network model
        
    Returns
    -------
    A pandas DataFrame that contains expected demand in m3/s
    (index = times, columns = junction names).
    """
    if start_time is None:
        start_time = 0
    if end_time is None:
        end_time = wn.options.time.duration
    if timestep is None:
        timestep = wn.options.time.report_timestep
        
    exp_demand = {}
    for name, junc in wn.junctions():
        exp_demand[name] = junc.demand_timeseries_list.get_values(start_time, 
                  end_time, timestep) * wn.options.hydraulic.demand_multiplier
    
    tsteps = np.arange(start_time, end_time+timestep, timestep)
    exp_demand = pd.DataFrame(index=tsteps, data=exp_demand)
    
    return exp_demand

def average_expected_demand(wn):
    """
    Compute average expected demand per day at each junction using base demands
    and demand patterns along with the demand multiplier
    
    Parameters
    -----------
    wn : wntr WaterNetworkModel
        Water network model
        
    Returns
    -------
    A pandas Series that contains average expected demand in m3/s
    (index = junction names).
    """
    L = [24*3600] # start with a 24 hour pattern
    for name, pattern in wn.patterns():
        L.append(len(pattern.multipliers)*wn.options.time.pattern_timestep)
    lcm = int(_lcml(L))
    
    start_time = wn.options.time.pattern_start
    end_time = start_time+lcm
    timestep = wn.options.time.pattern_timestep
        
    exp_demand = expected_demand(wn, start_time, end_time-timestep, timestep)
    ave_exp_demand = exp_demand.mean(axis=0)

    return ave_exp_demand

def _gcd(x,y):
  while y:
    if y<0:
      x,y=-x,-y
    x,y=y,x % y
    return x

def _gcdl(*list):
  return reduce(_gcd, *list)

def _lcm(x,y):
  return x*y / _gcd(x,y)

def _lcml(*list):
  return reduce(_lcm, *list)

def water_service_availability(expected_demand, demand):
    """
    Compute water service availability (WSA) at junctions, defined as follows:
        
    .. math:: WSA = \dfrac{demand}{expected\_demand}
        
    where 
    :math:`demand` is the actual demand computed from a hydraulic simulation, and 
    :math:`expected\_demand` is the expected demand computed from base demands and demand 
    patterns. Expected demand can be computed using the 
    :class:`~wntr.metrics.hydraulic.expected_demand` method.

    WSA can be averaged over times and/or nodes (see below).  If 
    expected demand is 0 for a particular junction, water service availability 
    will be set to NaN for that junction. 

    * To compute water service availability for each junction and timestep, 
      expected_demand and demand should be pandas DataFrames (index = times, columns = junction names). 
    
    * To compute an average water service availability for each junction (averaged over time), 
      expected_demand and demand should be a pandas Series, indexed by junction.  
      To convert a DataFrame (index = times, columns = junction names) to a 
      Series indexed by junction, use the following code:
    
        :math:`expected\_demand.sum(axis=0)`
		
        :math:`demand.sum(axis=0)`
    
    * To compute an average water service availability for each timestep (averaged over junctions), 
      expected_demand and demand should be a pandas Series, indexed by time.  
      To convert a DataFrame (index = times, columns = junction names) to a 
      Series indexed by time, use the following code:
        
        :math:`expected\_demand.sum(axis=1)`
		
        :math:`demand.sum(axis=1)`
        
    Parameters
    ----------
    expected_demand : pandas DataFrame or pandas Series (see note above)
        Expected demand.  
    
    demand : pandas DataFrame or pandas Series (see note above)
        Actual demand (generally from a PDD hydraulic simulation)

    Returns
    -------
    A pandas DataFrame or pandas Series that contains water service 
    availability.
    """

    wsa = demand.div(expected_demand) 
    
    return wsa

def todini_index(head, pressure, demand, flowrate, wn, Pstar):
    """
    Compute Todini index, equations from [Todi00]_.

    The Todini index is related to the capability of a system to overcome
    failures while still meeting demands and pressures at the nodes. The
    Todini index defines resilience at a specific time as a measure of surplus
    power at each node and measures relative energy redundancy.

    Parameters
    ----------
    head : pandas DataFrame
        A pandas Dataframe containing node head 
        (index = times, columns = node names).
        
    pressure : pandas DataFrame
        A pandas Dataframe containing node pressure 
        (index = times, columns = node names).
        
    demand : pandas DataFrame
        A pandas Dataframe containing node demand 
        (index = times, columns = node names).
        
    flowrate : pandas DataFrame
        A pandas Dataframe containing pump flowrates 
        (index = times, columns = pump names).

    wn : wntr WaterNetworkModel
        Water network model.  The water network model is needed to 
        find the start and end node to each pump.

    Pstar : float
        Pressure threshold.

    Returns
    -------
    A pandas Series that contains a time-series of Todini indexes
    """

    POut = {}
    PExp = {}
    PInRes = {}
    PInPump = {}

    time = head.index
    
    for name in wn.junction_name_list:
        h = np.array(head.loc[:,name]) # m
        p = np.array(pressure.loc[:,name])
        e = h - p # m
        q = np.array(demand.loc[:,name]) # m3/s
        POut[name] = q*h
        PExp[name] = q*(Pstar+e)

    for name, node in wn.nodes(wntr.network.Reservoir):
        H = np.array(head.loc[:,name]) # m
        Q = np.array(demand.loc[:,name]) # m3/s
        PInRes[name] = -Q*H # switch sign on Q.

    for name, link in wn.links(wntr.network.Pump):
        start_node = link.start_node_name
        end_node = link.end_node_name
        h_start = np.array(head.loc[:,start_node]) # (m)
        h_end = np.array(head.loc[:,end_node]) # (m)
        h = h_start - h_end # (m)
        q = np.array(flowrate.loc[:,name]) # (m^3/s)
        PInPump[name] = q*(abs(h)) # assumes that pumps always add energy to the system

    todini = (sum(POut.values()) - sum(PExp.values()))/  \
        (sum(PInRes.values()) + sum(PInPump.values()) - sum(PExp.values()))

    todini = pd.Series(data = todini.tolist(), index = time)

    return todini

def entropy(G, sources=None, sinks=None):
    """
    Compute entropy, equations from [AwGB90]_.

    Entropy is a measure of uncertainty in a random variable.
    In a water distribution network model, the random variable is
    flow in the pipes and entropy can be used to measure alternate flow paths
    when a network component fails.  A network that carries maximum entropy
    flow is considered reliable with multiple alternate paths.

    Parameters
    ----------
    G : NetworkX or WNTR graph
        Entropy is computed using a directed graph based on pipe flow direction.
        The 'weight' of each link is equal to the flow rate.

    sources : list of strings, optional (default = all reservoirs)
        List of node names to use as sources.

    sinks : list of strings, optional (default = all nodes)
        List of node names to use as sinks.

    Returns
    -------
    A tuple which includes:
        - A pandas Series that contains entropy for each node
        - System entropy (float)
    """

    if G.is_directed() == False:
        return

    if sources is None:
        sources = [key for key,value in nx.get_node_attributes(G,'type').items() if value == 'Reservoir' ]

    if sinks is None:
        sinks = G.nodes()

    S = {}
    Q = {}
    for nodej in sinks:
        if nodej in sources:
            S[nodej] = 0 # nodej is the source
            continue

        sp = [] # simple path
        if G.node[nodej]['type']  == 'Junction':
            for source in sources:
                if nx.has_path(G, source, nodej):
                    simple_paths = _all_simple_paths(G,source,target=nodej)
                    sp = sp + ([p for p in simple_paths])
                    # all_simple_paths was modified to check 'has_path' in the
                    # loop, but this is still slow for large networks
                    # what if the network was skeletonized based on series pipes
                    # that have the same flow direction?
                    # what about duplicating paths that have pipes in series?
                #print j, nodeid, len(sp)

        if len(sp) == 0:
            S[nodej] = np.nan # nodej is not connected to any sources
            continue

        sp = np.array(sp)

        # Uj = set of nodes on the upstream ends of links incident on node j
        Uj = G.predecessors(nodej)
        # qij = flow in link from node i to node j
        qij = []
        # aij = number of equivalnet independent paths through the link from node i to node j
        aij = []
        for nodei in Uj:
            mask = np.array([nodei in path for path in sp])
            # NDij = number of paths through the link from node i to node j
            NDij = sum(mask)
            if NDij == 0:
                continue
            temp = sp[mask]
            # MDij = links in the NDij path
            MDij = [(t[idx],t[idx+1]) for t in temp for idx in range(len(t)-1)]

            flow = 0
            for link in G[nodei][nodej].keys():
                flow = flow + G[nodei][nodej][link]['weight']
            qij.append(flow)

            # dk = degree of link k in MDij
            dk = Counter()
            for elem in MDij:
                # divide by the numnber of links between two nodes
                dk[elem] += 1/len(G[elem[0]][elem[1]].keys())
            V = np.array(list(dk.values()))
            aij.append(NDij*(1-float(sum(V - 1))/sum(V)))

        Q[nodej] = sum(qij) # Total flow into node j

        # Equation 7
        S[nodej] = 0
        for idx in range(len(qij)):
            if qij[idx]/Q[nodej] > 0:
                S[nodej] = S[nodej] - \
                    qij[idx]/Q[nodej]*math.log(qij[idx]/Q[nodej]) + \
                    qij[idx]/Q[nodej]*math.log(aij[idx])

    Q0 = sum(nx.get_edge_attributes(G, 'weight').values())

    # Equation 3
    S_ave = 0
    for nodej in sinks:
        if not np.isnan(S[nodej]):
            if nodej not in sources:
                if Q[nodej]/Q0 > 0:
                    S_ave = S_ave + \
                        (Q[nodej]*S[nodej])/Q0 - \
                        Q[nodej]/Q0*math.log(Q[nodej]/Q0)
                        
    S = pd.Series(S) # convert S to a series
    
    return [S, S_ave]
