"""
The wntr.graphics.layer module includes methods plot
water network data layers.
"""
import logging
import networkx as nx
import numpy as np
import matplotlib.pyplot as plt

logger = logging.getLogger(__name__)

def plot_valve_layer(wn, valve_layer, valve_attribute=None, title=None,
               valve_size=15, valve_range=[None,None], valve_cmap=None, add_colorbar=True, 
               colorbar_label=None, include_network=True, ax=None, filename=None):
    """
    Plot valve layer
    
    Parameters
    ----------
    wn: wntr WaterNetworkModel
        A WaterNetworkModel object
    
    valve_layer: pd.Dataframe, optional
        Valve layer, defined by node and link pairs (for example, valve 0 is 
        on link A and protects node B). The valve_layer DataFrame is indexed by
        valve number, with columns named 'node' and 'link'.
    
    valve_attribute: pd.Series, optional
        Attribute for each valve
        
    title: str, optional
        Plot title 

    valve_size: int, optional
        Node size 
    
    valve_range: list, optional
        Value range used to scale colormap ([None,None] indicates autoscale)
        
    valve_cmap: matplotlib.pyplot.cm colormap or named color, optional
        Valve colormap 

    add_colorbar: bool, optional
        Add colorbar
    
    colorbar_label: str, optional
        Node colorbar label
    
    include_network: bool, options
        Include a plot of the water network 
        
    ax: matplotlib axes object, optional
        Axes for plotting (None indicates that a new figure with a single 
        axes will be used)
        
    filename : str, optional
        Filename used to save the figure
        
    Returns
    -------
    ax : matplotlib axes object  
    """

    if ax is None: # create a new figure
        plt.figure(facecolor='w', edgecolor='k')
        ax = plt.gca()
        
    # Graph, undirected
    G = wn.get_graph().to_undirected()

    # Position
    pos = nx.get_node_attributes(G,'pos')
    if len(pos) == 0:
        pos = None

    if title is not None:
        ax.set_title(title)
    
    if include_network:
        nx.draw_networkx_edges(G, pos, edge_color='grey', width=0.5, ax=ax)
        
    ax.axis('off')
    
    valve_coordinates = []
    for valve_name, (pipe_name, node_name) in valve_layer.iterrows():
        pipe = wn.get_link(pipe_name)
        if node_name == pipe.start_node_name:
            start_node = pipe.start_node
            end_node = pipe.end_node
        elif node_name == pipe.end_node_name:
            start_node = pipe.end_node
            end_node = pipe.start_node
        else:
            print("Not valid")
            continue
        x0 = start_node.coordinates[0]
        dx = end_node.coordinates[0] - x0
        y0 = start_node.coordinates[1]
        dy = end_node.coordinates[1] - y0
        valve_coordinates.append((x0 + dx * 0.1, y0 + dy * 0.1))
    
    valve_coordinates = np.array(valve_coordinates)
    
    if valve_attribute is not None: 
        sc = ax.scatter(valve_coordinates[:,0], valve_coordinates[:,1], s=valve_size, c=valve_attribute, marker='v', cmap=valve_cmap)  
    else:
        sc = ax.scatter(valve_coordinates[:,0], valve_coordinates[:,1], valve_size, 'k', 'v')   
 
    if add_colorbar:
        clb = plt.colorbar(sc, shrink=0.5, pad=0.05, ax=ax)
        clb.ax.set_title(colorbar_label, fontsize=10)
        clb.mappable.set_clim(valve_range[0],valve_range[1])
        
    if filename:
        plt.savefig(filename)
        
    return ax
