import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import gaussian_kde


class _pyjoyplotter():

    def __init__( self, data=None, x=None, y=None, hue=None, kind=None, offset=0.75,
                 bottom_offset=0, cmap='Dark2', smooth=1, order=None, bins=None, weights=None,
                 figsize=None, ax=None, xlim=None, show_counts=False, weight_heights=False):
        ' initialise class, check args'
        assert kind in ['line', 'density', 'hist']
        assert type(data) == pd.core.frame.DataFrame 
        assert (x in data) and (hue in data), \
                'x and hue must be columns in dataframe'
        if kind == 'line': assert (y in data), \
                'y must be column in dataframe for line plot'
        
        self.data = data
        self.x = x
        self.y = y
        self.hue = hue
        self.categories = list(set(data[hue])) if order == None else order
        self.n = len(self.categories)
        self.offset = offset
        self.bottom_offset = bottom_offset
        self.colours = self._get_colours(cmap)
        self.n_c = len(self.colours)
        self.smooth = smooth
        self.kind = kind
        self.bins = bins
        self.weights = weights
        self.figsize = figsize
        self.ax = ax
        self.xlim = xlim
        self.show_counts = show_counts
        self.weight_heights = weight_heights

    
    def _get_colours(self, cmap):
    
        if type(cmap) == str:
            assert cmap in [
            m for m in plt.cm.datad if not m.endswith("_r")]
            cm = plt.get_cmap(cmap)
            
            return [cm(i) for i in range(cm.N)]

        if type(cmap) == list:
            assert len(cmap) == self.n
            return cmap

    def _get_ax(self):
        if not self.ax:
            fig = plt.figure(figsize=self.figsize)
            ax = plt.axes(frameon=False)
            self.ax = ax
        else:
            ax = self.ax
            ax.set_frame_on(False)
        return ax

    def _print_numbers(self):
        ax = self.ax
        xlim = self.xlim
        xpos = (xlim[0] + 5*xlim[1])/6.0
        for i, c in enumerate(self.categories):
            i = self.n - (i + 1)
            ypos = i * self.offset + self.bottom_offset
            val = len(self.data[self.data[self.hue] == c])
            ax.text(xpos, ypos, '{:,d}'.format(val), ha='center')

    def _line_plot(self):
        ax = self._get_ax()

        self.data = self.data.sort_values(by=self.x)

        for i, c in enumerate(self.categories):
        
            i = self.n - (i + 1)
            df = self.data[
                self.data[self.hue] == c
                ].rolling(self.smooth).mean()

            x_d = df.loc[df[self.hue] == c, self.x].values
            y_d = df.loc[df[self.hue] == c, self.y].values
            y_d = (y_d - np.nanmin(y_d)
                )/(np.nanmax(y_d) - np.nanmin(y_d))

            y_d += i * self.offset + self.bottom_offset
            y_min = np.tile(np.nanmin(y_d), (len(y_d)))
        
            col = self.colours[i % self.n_c]

            ax.plot(x_d, y_d, 
                    color=col, 
                    label=c,
                    alpha=0.8)
                    
            ax.fill_between(x_d, y_d, y_min,
                    alpha=0.6,
                    color=col)

        x_min = self.data[self.x].min()
        x_max = self.data[self.x].max()
        self.xlim = (x_min, x_max)

        ax.set_xlim(self.xlim)  
        ax.set_xlabel(self.x)
        ax.set_yticks([self.offset*i + self.bottom_offset for i in range(self.n)[::-1]])
        ax.set_yticklabels(self.categories)
    
        if self.show_counts:
            self._print_numbers()

        return ax


    def _density_plot(self):
        ax = self._get_ax()

        x_min = self.data[self.x].min()
        x_max = self.data[self.x].max()
        if not self.xlim:
            xlim = (x_min, x_max)
            self.xlim = xlim
        else:
            xlim = self.xlim

        
        height_weights = [len(self.data[self.data[self.hue] == c]) for c in self.categories]
        height_weights = [float(hw) / max(height_weights) for hw in height_weights]

        for idx, c in enumerate(self.categories):
        
            i = self.n - (idx + 1)

            values = self.data[self.data[self.hue] == c][self.x].values
            kernel = gaussian_kde(values)
            x_d = np.linspace(np.nanmin(values), np.nanmax(values), 200)
            y_d = kernel(x_d)

            y_d = (y_d - np.nanmin(y_d)
                )/(np.nanmax(y_d) - np.nanmin(y_d))
            if self.weight_heights:
                y_d *= height_weights[idx]

            y_d += i * self.offset + self.bottom_offset
            y_min = np.tile(np.nanmin(y_d), (len(y_d)))
        
            col = self.colours[i % self.n_c]

            ax.plot(x_d, y_d, 
                    color=col, 
                    label=c,
                    alpha=0.8)
                    
            ax.fill_between(x_d, y_d, y_min,
                    alpha=0.6,
                    color=col)

        ax.set_xlim(xlim)
        ax.set_xlabel(self.x)
        ax.set_yticks([self.offset*i + self.bottom_offset for i in range(self.n)[::-1]])
        ax.set_yticklabels(self.categories)
    
        if self.show_counts:
            self._print_numbers()

        return ax

    def _hist_plot(self):
        ax = self._get_ax()

        if type(self.bins) != list:
            self.bins = self.n * [self.bins]
        else:
            self.bins = list(reversed(self.bins))
        if type(self.weights) != list:
            self.weights = self.n * [self.weights]
        else:
            self.weights = list(reversed(self.weights))
        
        for i, c in enumerate(self.categories):
        
            i = self.n - (i + 1)
            df = self.data[
                self.data[self.hue] == c]

            x_d = df.loc[df[self.hue] == c, self.x].values
            x_d = x_d[~np.isnan(x_d.astype("float"))]

            col = self.colours[i % self.n_c]
            if self.weights[i]:
                hist = np.histogram(x_d, bins=self.bins[i], weights=x_d)
            else:
                hist = np.histogram(x_d, bins=self.bins[i])
            
            new_hist = (hist[0] / np.nanmax(hist[0]),
                    hist[1])

            width = (new_hist[1] - np.roll(new_hist[1], 1))[1:]  
            bot = i * self.offset + self.bottom_offset
            left = new_hist[1][:-1]
            height = new_hist[0]    
            
            ax.bar(left, height, 
                    width=width, bottom=bot, alpha=0.8, align='edge')

        ax.set_xlabel(self.x)
        ax.set_yticks([self.offset*i + self.bottom_offset for i in range(self.n)[::-1]])
        ax.set_yticklabels(self.categories)

        if self.show_counts:
            self._print_numbers()

        return ax
    
    def _plot(self):
        
        if self.kind == 'line':
            return self._line_plot()

        elif self.kind == 'density':
            return self._density_plot()

        else:
            assert self.kind == 'hist', 'Unrecognized kind: {}'.format(self.kind)
            return self._hist_plot()


def plot(data=None, x=None, y=None, hue=None, kind='line', offset=0.75, bottom_offset=0,
         cmap='Dark2', smooth=1, order=None, bins=10, weights=None, figsize=None, xlim=None,
         ax=None, show_counts=False, weight_heights=False):
    '''
    Create 'Joy Plot':
        data (pd.DataFrame): DataFrame holding all data
        x (str)  : DataFrame column to use as x value
        y (str)  : DataFrame column to use as y values 
        hue (str): DataFrame column to use to group data
        kind (str): specify plot type; line or hist
        offset (int/float): vertical seperation between plots
        cmap (str/list): name of matplotlib cmap, or list
                 of colors to be used for plots
        smooth (int): smoothing window, if smoothing to be applied
        order (list): order of categories - top to bottom
        bins (int/list): bins if using hist. int for all hists to have same bins
                 else list of bin no. for each hist
        weights (boolean/list): should the histogram be weighted?

    ''' 


    plotter = _pyjoyplotter(data=data, x=x, y=y, hue=hue, offset=offset,
                            bottom_offset=bottom_offset, cmap=cmap, smooth=smooth, kind=kind,
                            order=order, bins=bins, weights=weights, figsize=figsize, xlim=xlim,
                            ax=ax, show_counts=show_counts, weight_heights=weight_heights)
    return plotter._plot()  
    
