import json
import os
import pickle
import luigi
from matplotlib import pyplot as plt
import scipy as sp
import seaborn as sns
from figures import style as style
from figures.analysispipeline.data import DecisionTree, Table, SumStatsPlt
from figures.analysispipeline.datatypes import DatasetTask
from misc.plot_decision_tree import TreePlotter

__author__ = 'emmanuel'


class TreePlot(DatasetTask):
    def requires(self):
        return {"tree": DecisionTree(self.dataset),
                "run_sample_only": Table(self.dataset, table_name="run_sample_only"),
                "sum_stats_plt": SumStatsPlt(self.dataset)}

    def output(self):
        return {"tree_plot": luigi.LocalTarget(os.path.join(os.path.dirname(self.input()["tree"].fn), "decision_tree.pdf")),
                "measure_plots": luigi.LocalTarget(os.path.join(os.path.dirname(self.input()["tree"].fn), "measures"))}

    def run(self):
        run_sample_only = self.input()['run_sample_only'].load()
        sum_stats_plt = self.input()['sum_stats_plt'].load()
        sum_stats_plt.name = sum_stats_plt.name.apply(style.internal_name_to_print_name)

        with open(self.input()['tree'].fn, "rb") as f:
            decision_tree = pickle.load(f)
        tree = decision_tree.tree_
        my_palette = [style.color[name] for name in style.print_names]
        colors = [style.rgb_to_hex(col) for col in my_palette]

        feature_short_names_tree_labels = {-2: 'leaf', 0: 'r<SUB>ee</SUB>',
                       1: 'r<SUB>i/o</SUB>', 2: 'r<SUB>ei</SUB>', 3: 'r<SUP>(5)</SUP>'}
        tree_plotter = TreePlotter(feature_short_names_tree_labels, colors, style.print_names, label_mode="html")

        try:
            noise_type_ = json.loads(run_sample_only.json_parameters[0].replace("'", '"'))["noise_type"]
        except Exception as e:
            noise_type_ = "ShuffleEdgesEndsNoise"
        from figures.names import noise_type_rename
        noise_type = noise_type_rename[noise_type_]

        def plot_networks_over_noise(meas, plot_data, exclude_nets=(), title="Measure", ylabel="",
                            additional_hlines=(), counter="No counter"):
            fig, ax = plt.subplots()
            exclude = exclude_nets
            if "AP" in exclude_nets:
                exclude.append("API")
            nr_nets = len(style.print_names)
            noise_width = .1/2
            bar_width = noise_width/nr_nets
            bar_offsets = sp.arange(0, nr_nets) * bar_width - noise_width / 2
            plot_data = plot_data[['name', 'noise', meas]]
            for noise_level in sorted(plot_data.noise.unique()):
                noise_data = plot_data[plot_data.noise == noise_level]
                for k in range(nr_nets):
                    if style.print_names[k] not in exclude:
                        net_data = noise_data[noise_data.name == style.print_names[k]]
                        data = sp.asarray(net_data[meas])
                        ax.scatter([noise_level+bar_offsets[k]]*len(data), data,
                                    color=style.color[style.print_names[k]], alpha=1, label=style.print_names[k])
            ax.legend([p for p in style.print_names if p not in exclude], bbox_to_anchor=(1.6,1))
            sns.despine(ax=ax)
            noise_levels = plot_data.noise.unique()
            noise_levels.sort()
            if noise_levels.size > 10:
                noise_levels = [n for k, n in enumerate(noise_levels) if k % 2 == 0]
            ax.set_xticks(noise_levels)
            ax.set_xlabel(noise_type)
            ax.set_ylabel(ylabel)
            ax.set_title(title)
            for hline in additional_hlines:
                ax.axhline(hline, color='black', linewidth=2, zorder=-200, linestyle="--")
            plt.setp(ax.xaxis.get_majorticklabels(), rotation=-90)
            try:
                os.makedirs(self.output()["measure_plots"].fn)
            except FileExistsError as e:
                pass
            def save(app="", fig=None):
                base_path = os.path.join(self.output()["measure_plots"].fn, app)
                try:
                    os.makedirs(os.path.join(self.output()["measure_plots"].fn, app))
                except FileExistsError:
                    pass
                path = os.path.join(base_path, "{counter}.eps".format(counter=counter))
                fig.savefig(path, bbox_inches='tight')

            save(fig=fig, app="zoom_0")

            threshold = additional_hlines[0]

            ax.set_ylim([threshold-.02, threshold+.02])
            save(app="zoom_1", fig=fig)

            ax.set_ylim([threshold-.1, threshold+.1])
            save(app="zoom_2", fig=fig)

            ax.set_ylim([threshold-.2, threshold+.2])
            save(app="zoom_3", fig=fig)


        class NodeAction:
            def __init__(self):
                self.counter = 1

            def __call__(self, node, plotter):
                tree = plotter.tree
                is_leaf = tree.feature[node] == -2
                if not is_leaf:
                    measure = style.features_for_classification[tree.feature[node]]
                    measure_short_name = style.feature_short_name_latex[tree.feature[node]]
                    zero_classes = sp.argwhere(tree.value[node][0] == 0).T[0]
                    zero_classes_names = [name for nr, name in enumerate(style.print_names) if nr in zero_classes]
                    threshold = tree.threshold[node]
                    plot_networks_over_noise(measure, sum_stats_plt[sum_stats_plt.subsampling==1], exclude_nets=zero_classes_names,
                                             title=measure_short_name[:-1] + ": " + str(round(threshold,2)) + "$",
                                             additional_hlines = [threshold], ylabel=measure_short_name, counter=self.counter)
                    self.counter += 1


        tree_plotter.node_action = NodeAction()
        self.output()["tree_plot"].makedirs()
        tree_plotter.plot_tree(tree, self.output()["tree_plot"].fn)


class FeatureImportancePlot(DatasetTask):
    def requires(self):
        return DecisionTree(self.dataset)

    def output(self):
        return luigi.LocalTarget(os.path.join(os.path.dirname(self.input().fn), "feature_importance.eps"))

    def run(self):
        with open(self.input().fn, "rb") as f:
            decision_tree = pickle.load(f)

        fig, ax = plt.subplots()
        ax.bar(range(len(decision_tree.feature_importances_)),decision_tree.feature_importances_, color="black")
        xticks = sp.asarray(range(len(decision_tree.feature_importances_))) + .5
        xtick_labels =  [style.feature_short_name_latex[f]for f in range(len(decision_tree.feature_importances_))]
        ax.set_xticks(xticks)
        ax.set_xticklabels(xtick_labels)
        ax.set_title("Feature Importance")
        ax.tick_params(axis='x', length=0)
        style.despine(ax=ax)
        self.output().makedirs()
        fig.savefig(self.output().fn, bbox_inches="tight")