import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# TODO:
# Split into multiple classes and make it even more reusable!


class PlottingAdapter:

    def __init__(self):
        pass

#  Bar plots

    @classmethod
    def plot_bar(cls, x, y):
        fig = plt.figure()
        if len(x) > 10:
            plt.xticks(rotation=45, ha='right')  # 'vertical')
            fig.subplots_adjust(bottom=0.2)
        ax = sns.barplot(x, y, palette="Set3")
        fig = ax.get_figure()

        return fig

    @classmethod
    def plot_bar_positive_negative(cls, data):
        fig = plt.figure()
        data['param'] = data['PARAMETER'] + '\n' + data['PARAMETER.1']
        ax = sns.barplot(data.index, data['CORRELATION'], palette="Set3")
        plt.xticks(data.index, data['param'], rotation=45)  # 'vertical')
        plt.ylabel("Parameter interdependence based on PCC")
        fig = ax.get_figure()

        return fig

    @classmethod
    def plot_bar_stacked_scgrp(cls, data, groups):
        sns.set_palette('deep')
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ind = np.arange(len(data['PARAMETER NAME'].drop_duplicates().values.tolist()))
        width = 0.5

        plots = []
        prop_iter = iter(sns.color_palette().as_hex())
        for key, value in groups.iteritems():
            p = ax.bar(ind, data[data['group'] == value]['COMPOSITE SCALED SENSITIVITY'], width, color=next(prop_iter))
            plots.append(p)

        plt.ylabel('Grouped Parameter importance to observations, based on Composite Scaled Sensitivity')
        plt.xticks(ind + width/2., data['PARAMETER NAME'].drop_duplicates())
        plt.legend(plots, groups.values())

        return fig

    @classmethod
    def plot_bar_stacked2(cls, data):
        colnames = list(data)
        del colnames[:1]  # only vectors left in list
        fig = plt.figure()
        ax = fig.add_subplot(111)
        df = data[colnames]
        df.plot.bar(stacked=True, linewidth=0, ax=ax)
        plt.xticks(data.index, data['PARAMETER FOR EACH VECTOR ELEMENT'], rotation=45)  # 'vertical')
        plt.ylabel("Vector contribution")

        # Shrink current axis by 20%
        box = ax.get_position()
        ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
        ax.legend(loc='center left', bbox_to_anchor=(1, 0.5),
                  fancybox=True, shadow=True)
        plt.setp(plt.gca().get_legend().get_texts(), fontsize='6')

        return fig

    @classmethod
    def plot_bar_stacked(cls, data):
        colnames = list(data)
        del colnames[:3]  # only vectors left in list
        fig = plt.figure()
        ax = fig.add_subplot(111)
        df = data[colnames]

        df.plot.bar(stacked=True, linewidth=0, ax=ax)
        plt.xticks(data['PARAMETER NAME'].index, data['PARAMETER NAME'], rotation=45)  # 'vertical')
        plt.ylabel("Composite scaled sensitivity")

        # Shrink current axis by 20%
        box = ax.get_position()
        ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
        ax.legend(loc='center left', bbox_to_anchor=(1, 0.5), fancybox=True, shadow=True)
        plt.setp(plt.gca().get_legend().get_texts(), fontsize='6')

        return fig

#  Scatter plots

    @classmethod
    def plot_scatter_os(cls, data, annotate=True):
        data_names = ('SIMULATED EQUIVALENT', 'OBSERVED or PRIOR VALUE')
        axis_names = ('Simulated Equivalent', 'Observed Value')
        return cls.plot_scatter(data, data_names, axis_names, annotate)

    @classmethod
    def plot_scatter_ww(cls, data, annotate=True):
        data_names = ('WEIGHTED SIMULATED EQUIVALENT', 'WEIGHTED OBSERVED or PRIOR VALUE')
        axis_names = ('Weighted Simulated Equivalent', 'Weighted Observation')
        return cls.plot_scatter(data, data_names, axis_names, annotate)

    @classmethod
    def plot_scatter_ws(cls, data, annotate=True):
        data_names = ('SIMULATED EQUIVALENT', 'WEIGHTED RESIDUAL')
        axis_names = ('Unweighted Simulated Values', 'Weighted Residual')
        return cls.plot_scatter(data, data_names, axis_names, annotate, True)

    @classmethod
    def plot_scatter(cls, data, touple_data, touple_axis_names, annotate=True, zeroline=False):

        fig, ax = plt.subplots()
        ax.scatter(data[touple_data[0]], data[touple_data[1]], zorder=10)

        if annotate:
            for i, txt in enumerate(data['OBSERVATION or PRIOR NAME']):
                ax.annotate(txt, (data[touple_data[0]][i], data[touple_data[1]][i]))

        if zeroline:
            lims = ax.get_xlim()
            ax.plot(lims, np.zeros(len(lims)), ls="--", c=".3", zorder=0)
        else:
            lims = [np.min([ax.get_xlim(), ax.get_ylim()]), np.max([ax.get_xlim(), ax.get_ylim()])]
            ax.plot(lims, lims, ls="--", c=".3", zorder=0)
            ax.set_aspect('equal')
            ax.set_xlim(lims)
            ax.set_ylim(lims)

        plt.xlabel(touple_axis_names[0])
        plt.ylabel(touple_axis_names[1])

        return fig

# Line Plots

    @classmethod
    def plot_facet(cls, data, col_wrap=1):
        sns.set(style="ticks")

        grid = sns.FacetGrid(data, col="Parameter", hue="Parameter", col_wrap=col_wrap, size=3.5)

        grid.map(plt.plot, "ObservationIndex", "Data",  marker="o", ms=4).set_axis_labels('DSS', 'Observation')

        grid.set(xticks=np.arange(len(data.Observation.drop_duplicates().values.tolist())))
        grid.set_xticklabels(data.Observation.drop_duplicates().values.tolist(), rotation=90, fontsize=8)

        return grid

    @classmethod
    def plot_lines_itter(cls, data, relative):
        sns.set_style("darkgrid")
        parameters = data['Parameter'].drop_duplicates().values.tolist()

        fig = plt.figure()
        ax = fig.add_subplot(111)

        for param in parameters:
            df = data[data['Parameter'] == param]
            if relative:
                initial = df['ESTIMATE'].iloc[0]
                pd.options.mode.chained_assignment = None
                df.ESTIMATE = df.ESTIMATE / initial
            ax.plot(df['ITERATION'], df['ESTIMATE'], label=param)
        ax.legend()
        plt.xlabel("Iteration")
        plt.ylabel("Parameter Estimate")

        return fig

    @classmethod
    def plot_multiple_in_one(cls, data):
        sns.set_style("darkgrid")
        data = data.sort_values("ObservationIndex")
        parameters = data['Parameter'].drop_duplicates().values.tolist()
        fig = plt.figure()
        ax = fig.add_subplot(111)
        for param in parameters:
            df = data.loc[data['Parameter'] == param]
            ax.plot(df['ObservationIndex'], df['Data'], label=param)
        # df = data.loc[data['Parameter'] == parameters[0]]
        r = np.arange(df["Observation"].size)
        if len(r) > 10:
            plt.xticks(r, df["Observation"], rotation=45)  # 'vertical')
            fig.subplots_adjust(bottom=0.2)
        plt.xticks(r, df["Observation"])

        ax.legend()
        plt.ylabel("DSS")
        plt.xlabel("Observations")
        return fig

    @classmethod
    def plot_lines(cls, data):
        sns.set_style("darkgrid")
        fig = plt.figure()
        ax = fig.add_subplot(111)
        colnames = list(data)
        del colnames[:2]
        for param in colnames:
            ax.plot(data.index, data[[param]], label=param)
        ax.set_xticklabels(data["OBSERVATION NAME"], rotation=45, ha='right')
        r = np.arange(data["OBSERVATION NAME"].size)
        plt.xticks(r, data["OBSERVATION NAME"])

        ax.legend()
        plt.xlabel("Observation Name")
        plt.ylabel("One percent scaled sensitivity")

        return fig

    @classmethod
    def plot_multiple_lines_annotated(cls, data):
        sns.set_style("darkgrid")
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.plot(data['ITERATION'], data['SSWR-(OBSERVATIONS ONLY)'], label='OBSERVATIONS ONLY', marker='x')
        ax.plot(data['ITERATION'], data['SSWR-(PRIOR INFORMATION ONLY)'], label='PRIOR INFORMATION ONLY', marker=',')
        ax.plot(data['ITERATION'], data['SSWR-(TOTAL)'], label='Total', marker='.')
        for i, txt in enumerate(data['# OBSERVATIONS INCLUDED']):
            ax.annotate(txt, (data['ITERATION'][i], data['SSWR-(OBSERVATIONS ONLY)'][i]))
        ax.legend()
        plt.xlabel('Iteration')
        plt.ylabel('Sum of squared weighted Residuals')

        return fig
