"""
The graphics module contains graphic functions.

.. rubric:: Contents

.. autosummary::

    signal_convexhull
    signal_xsection
    animate_puffs
    sensor_locations
"""
from __future__ import print_function, division
try:
    import matplotlib.pyplot as plt
    from matplotlib import ticker
    from matplotlib.patches import Circle
    from matplotlib.collections import PatchCollection
    from matplotlib.animation import FuncAnimation
except:
    plt = None

from scipy.spatial import ConvexHull
import numpy as np
from chama.sensors import Mobile


def signal_convexhull(signal, scenarios, threshold, timesteps=None,
                      colormap=None,
                      x_range=(None, None), y_range=(None, None),
                      z_range=(None, None)):
    """
    Plots of the signal's 3D convex hull.

    Parameters 
    -------------- 
    signal: pandas DataFrame
        Signal data from the simulation.  
        The DataFrame contains columns 'X', 'Y', 'Z', 'T', and one column 
        for each scenario.
    scenarios: list
        Column names for the scenarios to be plotted
    threshold: float
        The minimum value of the signal to be included
    timesteps: list (optional)
        List of the time steps to include in the plot
    colormap: matplotlib.pyplot ColorMap (optional)
        A ColorMap object sent to the contourf function
    x_range: tuple (optional)
        The x-axis limits for the plot
    y_range: tuple (optional)
        The y-axis limits for the plot
    z_range: tuple (optional)
        The z-axis limits for the plot
    """
    if plt is None:
        raise ImportError('matplotlib is required for graphics')
    if colormap is None:
        colormap = plt.get_cmap('viridis')
        
    t_col = 'T'
    x_col = 'X'
    y_col = 'Y'
    z_col = 'Z'

    if timesteps is None:
        timesteps = sorted(set(signal.loc[:, t_col]))

    fig = plt.figure()
    plt.set_cmap(colormap)
    ax = fig.add_subplot(111, projection='3d')

    for scenario in scenarios:
        i = 0
        for timestep in timesteps:
            color = colormap(i)
            i += 1 / float(len(timesteps))
            
            signal_t = signal[signal[t_col] == timestep]
            conc_filter = signal_t[scenario] > threshold
            
            data = signal_t[[x_col, y_col, z_col]][conc_filter]
            data = data.values
            try:
                hull = ConvexHull(data)
                ax.plot_trisurf(data[:, 0], data[:, 1], data[:, 2],
                                triangles=hull.simplices,
                                edgecolor='none', 
                                shade=False,
                                color=color)
            except:
                print('Convex Hull not created for scenario %s timestep %d' % (scenario, timestep))
                      
    ax.set_xlabel(x_col)
    ax.set_ylabel(y_col)
    ax.set_zlabel(z_col)
    
    ax.set_xlim3d(x_range[0], x_range[1])
    ax.set_ylim3d(y_range[0], y_range[1])
    ax.set_zlim3d(z_range[0], z_range[1])
    fig.show()


def signal_xsection(signal, signal_name, threshold=None, timesteps=None, 
                    x_value=None, y_value=None, z_value=None, log_flag=False,
                    colormap=None, alpha=0.7, N=5,
                    x_range=(None, None), y_range=(None, None), 
                    z_range=(None, None)):
    """
    Plots x-y, x-z, and y-z cross sections. The signal is
    summed over all desired time steps and summed across the axis not
    included in the plot unless a value for the third access is specified.

    Parameters 
    -------------- 
    signal: pandas DataFrame
        Signal data from the simulation.  
        The DataFrame contains columns 'X', 'Y', 'Z', 'T', and one column 
        for each scenario.
    signal_name: string
        Column name for the signal to be plotted
    threshold: float (optional)
        The minimum value of the signal to be plotted
    timesteps: list (optional)
        List of the time steps to include in the plot
    x_value: list (optional)
        List of the x locations to include in the plot
    y_value: list (optional)
        List of the y locations to include in the plot
    z_value: list (optional)
        List of the z locations to include in the plot
    log_flag: boolean (optional)
        Flag specifying whether the signal should be plotted on a log scale
    colormap: matplotlib.pyplot ColorMap (optional)
        A ColorMap object sent to the contourf function
    alpha: float (optional)
        Value between 0 and 1 representing the alpha blending value
        passed to the contourf function
    N: int (optional)
        The number of levels to include in the plot, passed to the
        contourf function
    x_range: tuple (optional)
        The x-axis limits for the plot
    y_range: tuple (optional)
        The y-axis limits for the plot
    z_range: tuple (optional)
        The z-axis limits for the plot
    """
    if plt is None:
        raise ImportError('matplotlib is required for graphics')
    if colormap is None:
        colormap = plt.get_cmap('viridis')
        
    t_col = 'T'
    x_col = 'X'
    y_col = 'Y'
    z_col = 'Z'
    
    if timesteps is None:
        timesteps = sorted(set(signal.loc[:, t_col]))

    if log_flag:
        log_flag = ticker.LogLocator()
    else:
        log_flag = ticker.MaxNLocator()
      
    fig = plt.figure(figsize=(20, 5))
    plt.set_cmap(colormap)
    ax1 = fig.add_subplot(1, 3, 1)
    ax2 = fig.add_subplot(1, 3, 2)
    ax3 = fig.add_subplot(1, 3, 3)
    
    signal_t = signal[signal[t_col].isin(timesteps)]
    signal_t = signal_t.groupby([x_col, y_col, z_col]).sum()
    data = signal_t[signal_name]

    def contour_data(temp, threshold, log_flag):
        temp = temp.unstack()
        X = temp.index.values
        Y = temp.columns.values
        Z = temp.values

        if threshold:
            Z[Z <= threshold] = 0
        Z = np.transpose(Z)
        return X, Y, Z
    
    temp = data.reset_index()
    if z_value:
        if type(z_value) is not list:
            z_value = [z_value]
        temp = temp[temp[z_col].isin(z_value)]
    temp = temp.groupby([x_col, y_col])[signal_name].sum()
    
    Xi, Yi, Z = contour_data(temp, threshold, log_flag)

    cplot1 = ax1.contourf(Xi, Yi, Z, alpha=alpha, cmap=colormap,
                          locator=log_flag)
    ax1.set_xlim(x_range[0], x_range[1])
    ax1.set_ylim(y_range[0], y_range[1])
    ax1.set_xlabel(x_col)
    ax1.set_ylabel(y_col)
    ax1.set_title(signal_name)

    temp = data.reset_index()
    if y_value:
        if type(y_value) is not list:
            y_value = [y_value]
        temp = temp[temp[y_col].isin(y_value)]
    temp = temp.groupby([x_col, z_col])[signal_name].sum()
    
    Xi, Yi, Z = contour_data(temp, threshold, log_flag)
    cplot2 = ax2.contourf(Xi, Yi, Z, N, alpha=alpha, cmap=colormap,
                          locator=log_flag)
    ax2.set_xlim(x_range[0], x_range[1])
    ax2.set_ylim(z_range[0], z_range[1])
    ax2.set_xlabel(x_col)
    ax2.set_ylabel(z_col)
    ax2.set_title(signal_name)

    temp = data.reset_index()
    if x_value:
        if type(x_value) is not list:
            x_value = [x_value]
        temp = temp[temp[x_col].isin(x_value)]
    temp = temp.groupby([y_col, z_col])[signal_name].sum()
    
    Xi, Yi, Z = contour_data(temp, threshold, log_flag)
    cplot3 = ax3.contourf(Xi, Yi, Z, N, alpha=alpha, cmap=colormap,
                          locator=log_flag)
    ax3.set_xlim(y_range[0], y_range[1])
    ax3.set_ylim(z_range[0], z_range[1])
    ax3.set_xlabel(y_col)
    ax3.set_ylabel(z_col)
    ax3.set_title(signal_name)

    fig.colorbar(cplot1, ax=ax1)
    fig.colorbar(cplot2, ax=ax2)
    fig.colorbar(cplot3, ax=ax3)

    fig.show()


def animate_puffs(puff, x_range=(None, None), y_range=(None, None), repeat=True):
    """
    Plots the horizontal movement of puffs from a GaussianPuff simulation
    over time. Each puff is represented as a circle centered at the puff
    center location with radius equal to the standard deviation in the
    horizontal direction (sigmaY).

    Parameters
    ------------------
    puff: pandas DataFrame
        The puff DataFrame created by a GaussianPuff object
    x_range: tuple (xmin, xmax) (optional)
        The x-axis limits for the plot
    y_range: tuple (ymin, ymax) (optional)
        The y-axis limits for the plot
    repeat : bool, optional
        If True, the animation will repeat
    """

    def circles(x, y, s, c='b', vmin=None, vmax=None, **kwargs):
        """
        Make a scatter plot of circles. 
        Similar to plt.scatter, but the size of circles are in data scale.

        Parameters
        ----------
        x, y : scalar or array_like, shape (n, )
            Input data
        s : scalar or array_like, shape (n, ) 
            Radius of circles.
        c : color or sequence of color, optional, default : 'b'
            `c` can be a single color format string, or a sequence of color
            specifications of length `N`, or a sequence of `N` numbers to be
            mapped to colors using the `cmap` and `norm` specified via kwargs.
            Note that `c` should not be a single numeric RGB or RGBA sequence 
            because that is indistinguishable from an array of values
            to be colormapped. (If you insist, use `color` instead.)  
            `c` can be a 2-D array in which the rows are RGB or RGBA, however. 
        vmin, vmax : scalar, optional, default: None
            `vmin` and `vmax` are used in conjunction with `norm` to normalize
            luminance data.  If either are `None`, the min and max of the
            color array is used.
        kwargs : `~matplotlib.collections.Collection` properties
            Eg. alpha, edgecolor(ec), facecolor(fc), linewidth(lw),
            linestyle(ls), norm, cmap, transform, etc.

        Returns
        -------
        paths : `~matplotlib.collections.PathCollection`
        """

        if np.isscalar(c):
            kwargs.setdefault('color', c)
            c = None

        if 'fc' in kwargs:
            kwargs.setdefault('facecolor', kwargs.pop('fc'))
        if 'ec' in kwargs:
            kwargs.setdefault('edgecolor', kwargs.pop('ec'))
        if 'ls' in kwargs:
            kwargs.setdefault('linestyle', kwargs.pop('ls'))
        if 'lw' in kwargs:
            kwargs.setdefault('linewidth', kwargs.pop('lw'))
        # You can set `facecolor` with an array for each patch,
        # while you can only set `facecolors` with a value for all.

        zipped = np.broadcast(x, y, s)
        patches = [Circle((x_, y_), s_)
                    for x_, y_, s_ in zipped]
        collection = PatchCollection(patches, **kwargs)
        if c is not None:
            c = np.broadcast_to(c, zipped.shape).ravel()
            collection.set_array(c)
            collection.set_clim(vmin, vmax)

        ax = plt.gca()
        ax.add_collection(collection)
        ax.autoscale_view()
        plt.draw_if_interactive()
        if c is not None:
            plt.sci(collection)
        return collection
    
    if plt is None:
        raise ImportError('matplotlib is required for graphics')
    
    buffer = puff['sigmaY'].max() + puff['sigmaY'].max()/5
    xmin = min([v for v in [puff['X'].min() - buffer, x_range[0]] if v is not None])
    xmax = max([v for v in [puff['X'].max() + buffer, x_range[1]] if v is not None])
    ymin = min([v for v in [puff['Y'].min() - buffer, y_range[0]] if v is not None])
    ymax = max([v for v in [puff['Y'].max() + buffer, y_range[1]] if v is not None])
    x_range = (xmin, xmax)
    y_range = (ymin, ymax)

    fig, ax = plt.subplots()
    # ln, = plt.plot([],[],animated=True)

    def update(time):
        plt.cla()
        ax.set_xlim(x_range[0], x_range[1])
        ax.set_ylim(y_range[0], y_range[1])
        ax.set_title('T = %6.2f' % time)
        ax.set_xlabel('X')
        ax.set_ylabel('Y')

        temp = puff.loc[puff['T'] == time]
        out = circles(temp['X'], temp['Y'], temp['sigmaY'], alpha=0.5,
                      edgecolor='none')
        return out
    
    anim = FuncAnimation(fig, update, frames=puff['T'].unique(), interval=50, 
                         blit=False, repeat=repeat)

    plt.show()
    
    return anim


def sensor_locations(sensors, x_range=(None, None), y_range=(None, None), 
                     z_range=(None, None), legend=False,
                     colors=None, markers=None):
    """
    Plots sensor locations.
    
    Parameters
    -------------
    sensors : dict
        A dictionary of sensors with key:value pairs containing
        {'sensor name': chama Sensor object}
    x_range: tuple (optional)
        The x-axis limits for the plot
    y_range: tuple (optional)
        The y-axis limits for the plot
    z_range: tuple (optional)
        The z-axis limits for the plot
    legend : Boolean (optional)
        Indicates if legend should be added to the plot
    colors : dict (optional)
        A dictionary containing the color string to be used when plotting each
        sensor. The key:value pairs are {'sensor name' : String
        representing the color to be passed to the plot function)
    markers : dict (optional)
        A dictionary containing the marker to be used when plotting each
        sensor. The key:value pairs are {'sensor name' : String
        representing the marker to be passed to the plot function)
    """
    if plt is None:
        raise ImportError('matplotlib is required for graphics')
        
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    for name, sensor in sensors.items():

        plot_options = {}
        if colors:
            plot_options['color'] = colors[name]
        if markers:
            plot_options['marker'] = markers[name]

        position = sensor.position
        if isinstance(position, Mobile):
            x = [val[0] for val in position.location]
            y = [val[1] for val in position.location]
            z = [val[2] for val in position.location]
            ax.plot(x, y, z, label=name)
        else:
            x = position.location[0]
            y = position.location[1]
            z = position.location[2]
            ax.scatter(x, y, z, label=name, **plot_options)

    ax.set_xlim3d(x_range[0], x_range[1])
    ax.set_ylim3d(y_range[0], y_range[1])
    ax.set_zlim3d(z_range[0], z_range[1])
    if legend:
        ax.legend()
    fig.show()
