import os
import luigi
from matplotlib import pyplot as plt
import seaborn as sns
from figures import style as style
from figures.analysispipeline.data import SumStatsPlt
from figures.analysispipeline.datatypes import DatasetTask

__author__ = 'emmanuel'


class AllMedianPlots(DatasetTask):
    def requires(self):
        return [MedianPlot(self.dataset, measure) for measure in style.features_for_classification]

    def output(self):
        return [MedianPlot(self.dataset, measure).output() for measure in style.features_for_classification]


class MedianPlot(DatasetTask):
    measure = luigi.Parameter()

    def requires(self):
        return SumStatsPlt(self.dataset)

    def output(self):
        return luigi.LocalTarget(os.path.join(os.path.dirname(self.input().fn), "median", str(self.measure) + ".eps"))

    def run(self):
        self.output().makedirs()
        sum_stats_plt = self.input().load()
        names = sum_stats_plt.name.unique()
        names = [style.internal_name_to_print_name(name) for name in names]
        median_lines = (sum_stats_plt[sum_stats_plt.subsampling==1]
                           .groupby(['name', 'noise']).median().reset_index())

        fig, ax = plt.subplots()
        for name in names:
            color_ = style.color[name]
            named_lines = median_lines[median_lines.name == name]
            ax.plot(named_lines.noise, named_lines[self.measure],
                    '-' ,color=color_, label=name, linewidth=4)
        ax.legend(bbox_to_anchor=(1.7, 1.1))
        ax.set_title("Median " + style.feature_short_name_latex[self.measure], y=1.1)
        ax.set_xlabel("Noise")
        ax.xaxis.set_tick_params(labelsize=style.font_size)
        sns.despine(ax=ax, offset=10)
        fig.savefig(self.output().fn, bbox_inches="tight")