import os
#import beeswarm
import luigi
import matplotlib
#matplotlib.use("agg", force=True)
from matplotlib import pyplot as plt
from figures import style as style
from figures.analysispipeline.data import SumStatsPlt
from figures.analysispipeline.datatypes import DatasetTask
import warnings


class BeePlot(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), "beeplot", str(self.measure) + ".eps"))

    def run(self):
        self.output().makedirs()
        sum_stats_plt = self.input().load()
        fig, ax = plt.subplots()
        colors= [style.color[name] for name in style.internal_names]
        zero_noise_data = sum_stats_plt[(sum_stats_plt.noise == 0) & (sum_stats_plt.subsampling == 1)]

        if len(zero_noise_data) == 0:
            zero_noise_data = sum_stats_plt
            warnings.warn("Bad hack here", DeprecationWarning)
        beeswarm_data = [zero_noise_data[zero_noise_data.name == name_][self.measure].as_matrix()
                         for name_ in style.internal_names]
        bs, ax = beeswarm.beeswarm(beeswarm_data, ax=ax, col=colors, labels=style.print_names, s=5,
                                   lw=0.1, edgecolor='white')
        ax.set_title(style.feature_short_name_latex[self.measure], y=1.05)
        plt.setp(ax.xaxis.get_majorticklabels(), rotation=-45)
        style.despine(ax=ax)
        fig.savefig(self.output().fn, bbox_inches="tight")


class AllBeePlots(DatasetTask):
    def requires(self):
        return [BeePlot(self.dataset, measure) for measure in style.features_for_display]

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