import os
import luigi
from matplotlib import pyplot as plt
import matplotlib as mpl
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 BoxPlot(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), "boxplot", str(self. measure) + ".eps"))

    def run(self):
        fig, ax = plt.subplots()
        self.output().makedirs()
        old_lw = mpl.rcParams['lines.linewidth']
        mpl.rcParams['lines.linewidth'] = 2
        sum_stats_plt = self.input().load()
        zero_noise_data = sum_stats_plt[(sum_stats_plt.noise == 0) & (sum_stats_plt.subsampling == 1)]
        sns.boxplot(x="name", y=self.measure, data=zero_noise_data, ax=ax)
        sns.despine()
        plt.setp(ax.xaxis.get_majorticklabels(), rotation=-45)
        ax.set_title(style.feature_short_name_latex[self.measure], y=1.05)
        ax.set_ylabel("")
        colors = [style.color[n] for n in style.internal_names]
        for (artist_, color_) in zip(ax.artists, colors):
            artist_.set_facecolor(color_)
            artist_.set_edgecolor("black")
            artist_.set_linewidth(1)
        fig.savefig(self.output().path, bbox_inches="tight")
        mpl.rcParams['lines.linewidth'] = old_lw


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

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