#!/usr/bin/env python3.8
# coding: latin-1

# (c) Massachusetts Institute of Technology 2015-2018
# (c) Brian Teague 2018-2022
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
cytoflow.views.bar_chart
------------------------

Plot a bar chart from a statistic.

`BarChartView` -- the `IView` class that makes the plot.
"""

from traits.api import provides, Constant

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

import cytoflow.utility as util
from .i_view import IView
from .base_views import Base1DStatisticsView

@provides(IView)
class BarChartView(Base1DStatisticsView):
    """
    Plots a bar chart of some summary statistic
    
    Attributes
    ----------
    
    Examples
    --------
    
    Make a little data set.
    
    .. plot::
        :context: close-figs
            
        >>> import cytoflow as flow
        >>> import_op = flow.ImportOp()
        >>> import_op.tubes = [flow.Tube(file = "Plate01/RFP_Well_A3.fcs",
        ...                              conditions = {'Dox' : 10.0}),
        ...                    flow.Tube(file = "Plate01/CFP_Well_A4.fcs",
        ...                              conditions = {'Dox' : 1.0})]
        >>> import_op.conditions = {'Dox' : 'float'}
        >>> ex = import_op.apply()
        
    Add a threshold gate
    
    .. plot::
        :context: close-figs
    
        >>> ex2 = flow.ThresholdOp(name = 'Threshold',
        ...                        channel = 'Y2-A',
        ...                        threshold = 2000).apply(ex)
        
    Add a statistic
    
    .. plot::
        :context: close-figs

        >>> ex3 = flow.ChannelStatisticOp(name = "ByDox",
        ...                               channel = "Y2-A",
        ...                               by = ['Dox', 'Threshold'],
        ...                               function = len).apply(ex2) 
    
    Plot the bar chart
    
    .. plot::
        :context: close-figs
        
        >>> flow.BarChartView(statistic = ("ByDox", "len"),
        ...                   variable = "Dox",
        ...                   huefacet = "Threshold").plot(ex3)
        
    """
    
    # traits   
    id = Constant("edu.mit.synbio.cytoflow.view.barchart")
    friendly_id = Constant("Bar Chart") 
        
    def enum_plots(self, experiment):
        """
        Returns an iterator over the possible plots that this View can
        produce.  The values returned can be passed to "plot".
        """
                
        return super().enum_plots(experiment)
        
        
    def plot(self, experiment, plot_name = None, **kwargs):
        """
        Plot a bar chart
        
        Parameters
        ----------
            
        color : a matplotlib color
            Sets the colors of all the bars, even if there is a hue facet
            
        errwidth : scalar
            The width of the error bars, in points
            
        errcolor : a matplotlib color
            The color of the error bars
            
        capsize : scalar
            The size of the error bar caps, in points
            
        Notes
        -----
        
        Other ``kwargs`` are passed to `matplotlib.axes.Axes.bar <https://matplotlib.org/devdocs/api/_as_gen/matplotlib.axes.Axes.bar.html>`_

        """
        
        super().plot(experiment, plot_name, **kwargs)
        
    def _grid_plot(self, experiment, grid, **kwargs):
                 
        # because the bottom of a bar chart is "0", masking out bad
        # values on a log scale doesn't work.  we must clip instead.
        orientation = kwargs.pop('orientation', 'vertical')
        
        # statistic scale
        scale = kwargs.pop('scale')
        
        if scale.name == "log":
            scale.mode = "clip"
            
        # limits
        lim = kwargs.pop('lim', None)
                      
        stat = experiment.statistics[self.statistic]
        if orientation == 'vertical':
            map_args = [self.variable, stat.name]
        else:
            map_args = [stat.name, self.variable]
        
        if self.huefacet:
            map_args.append(self.huefacet)  
        
        if self.error_statistic[0]:
            error_stat = experiment.statistics[self.error_statistic]
            map_args.append(error_stat.name)
        else:
            error_stat = None
                        
        grid.map(_barplot, 
                 *map_args,
                 view = self,
                 stat_name = stat.name,
                 error_name = error_stat.name if error_stat is not None else None,
                 orientation = orientation,
                 grid = grid,
                 **kwargs)
        
        if orientation == 'horizontal':
            return dict(xscale = scale,
                        xlim = lim)
        else:
            return dict(yscale = scale,
                        ylim = lim)
            
def _barplot(*args, view, stat_name, error_name, orientation, grid, **kwargs):
    """ 
    A custom barchart function.  This is assembled from pieces cobbled
    together from seaborn v0.7.1.
    """
  
    data = pd.DataFrame({s.name: s for s in args}).sort_values(view.variable)
    categories = data[view.variable].unique()
 
    # plot the bars
    width = kwargs.pop('width', 0.8)
    ax = kwargs.pop('ax', None)

    if ax is None:
        ax = plt.gca()
    
    err_kws = {}
    errwidth = kwargs.pop('errwidth', None)
    if errwidth:
        err_kws['lw'] = errwidth
    else:
        err_kws['lw'] = mpl.rcParams["lines.linewidth"]
         
    errcolor = kwargs.pop('errcolor', '0.2')
    capsize = kwargs.pop('capsize', None)

    # Get the right matplotlib function depending on the orientation
    barfunc = ax.bar if orientation == "vertical" else ax.barh
    barpos = np.arange(len(categories))
    
    if view.huefacet:
        hue_names = grid.hue_names
        hue_level = data[view.huefacet].iloc[0]
        hue_idx = hue_names.index(hue_level)
        hue_offsets = np.linspace(0, width - (width / len(hue_names)), len(hue_names))
        hue_offsets -= hue_offsets.mean()
        nested_width = width / len(hue_names) * 0.98
        
        offpos = barpos + hue_offsets[hue_idx]
        barfunc(offpos,
                data[stat_name], 
                nested_width,
                align="center",
                **kwargs)
                
        if error_name:
            confint = data[error_name]
            errcolors = [errcolor] * len(offpos)
            _draw_confints(ax,
                           offpos,
                           data[data[view.huefacet] == hue_level][stat_name],
                           confint,
                           errcolors,
                           orientation,
                           errwidth = errwidth,
                           capsize = capsize)
                
    else:
        barfunc(barpos, data[stat_name], width, align="center", **kwargs)
         
        if error_name:
            confint = data[error_name]
            errcolors = [errcolor] * len(barpos)
            _draw_confints(ax,
                           barpos,
                           data[stat_name],
                           confint,
                           errcolors,
                           orientation,
                           errwidth = errwidth,
                           capsize = capsize)

    if orientation == "vertical":
        ax.set_xticks(np.arange(len(categories)))
        ax.set_xticklabels(categories)
    else:
        ax.set_yticks(np.arange(len(categories)))
        ax.set_yticklabels(categories)
 
    if orientation == "vertical":
        ax.xaxis.grid(False)
        ax.set_xlim(-.5, len(categories) - .5)
    else:
        ax.yaxis.grid(False)
        ax.set_ylim(-.5, len(categories) - .5)  
            
    return ax 


def _draw_confints(ax, at_group, stat, confints, colors, 
                   orient, errwidth=None, capsize=None, **kws):
 
    if errwidth is not None:
        kws.setdefault("lw", errwidth)
    else:
        kws.setdefault("lw", mpl.rcParams["lines.linewidth"] * 1.8)
         
    if isinstance(confints.iloc[0], tuple):
        ci_lo = [x[0] for x in confints]
        ci_hi = [x[1] for x in confints]
    else:
        ci_lo = [stat.iloc[i] - x for i, x in confints.reset_index(drop = True).items()]
        ci_hi = [stat.iloc[i] + x for i, x in confints.reset_index(drop = True).items()]
 
    for at, lo, hi, color in zip(at_group,
                                 ci_lo,
                                 ci_hi,
                                 colors):
        if orient == "vertical":
            if capsize is not None:
                kws['marker'] = '_'
                kws['markersize'] = capsize * 2
                kws['markeredgewidth'] = kws['lw']
            ax.plot([at, at], [lo, hi], color=color, **kws)
        else:
            if capsize is not None:
                kws['marker'] = '|'
                kws['markersize'] = capsize * 2
                kws['markeredgewidth'] = kws['lw']
            ax.plot([lo, hi], [at, at], color=color, **kws)


util.expand_class_attributes(BarChartView)
util.expand_method_parameters(BarChartView, BarChartView.plot)