"""
Boxplot Simulation.
"""
from Expe7132 import *
sns.set_theme(style='darkgrid', palette='pastel')
import pickle
import os

if not os.path.exists("Expe6585/"):
    os.makedirs("Expe6585/")


class Expe6585(Expe7132):
    """Provide Boxplots.
    """
    expe_no='6585'
    num_iter=30
    showfliers=False

    def dump_df(self):
        """Generate and dump a pd.DataFrame onto the disk.
        """
        self.clear()
        with rich.progress.Progress() as progress:
            task = progress.add_task(f'Performing Experiment {self.expe_no}: generating data...', total = self.num_iter * len(self.alpha_list))
            for i, alpha in enumerate(self.alpha_list):
                self.alpha = alpha
                self.simulate(num_iter = self.num_iter, progress = progress, task = task)
        pickle.dump(pd.DataFrame(self), open(f"Expe{self.expe_no}/Expe{self.expe_no}.d", "wb"))
        self.clear()

    def load_df(self):
        """Load the generated pd.DataFrame from the disk.
        """
        return pickle.load(open(f"Expe{self.expe_no}/Expe{self.expe_no}.d", "rb"))

    def plot(self, to_generate = False, dict_columns = None):
        """Plot the boxplots based on the pd.DataFrame.
        """

        # Prepare Data
        if to_generate:
            self.dump_df()
        self_df = self.load_df()

        self_df = self_df.rename(columns = self.rename())

        # Plot and save the figures.
        for ext_str in dict_columns:
            columns = dict_columns[ext_str]
            res_df = pd.DataFrame(columns = ['y', r'$\lambda$', 'hue'])
            for key in columns:
                key = self.rename(key)
                if key == r'$\lambda$':
                    continue
                temp_df = pd.DataFrame()
                temp_df[r'$\lambda$'] = self_df[r'$\lambda$']
                temp_df['hue'] = key
                temp_df['y'] = self_df[key]
                res_df = res_df.append(temp_df, ignore_index=True)
            fig, ax = matplotlib.pyplot.subplots(figsize = (12, 9))
            sns.boxplot(
                x = r'$\lambda$',
                y = 'y',
                hue = 'hue',
                data = res_df,
                ax = ax,
                showfliers = self.showfliers,
                )
            ax.legend(loc = 'upper right')
            ax.set(title=r'', xlabel = r'$\lambda$', ylabel = r'')
            new_xticks = [x for x in ax.get_xticks() if x % 3 == 0]
            new_xticklabels = [float(f'{float(x.get_text()):.2g}') for x in ax.get_xticklabels() if x.get_position()[0] in new_xticks]
            ax.set_xticks(new_xticks)
            ax.set_xticklabels(new_xticklabels)
            ax.tick_params(labelsize = 20)
            ax.set_xlabel(r'$\lambda$', fontsize=20)
            plt.legend(fontsize=20)
            fig.savefig(f'Expe{self.expe_no}/Expe{self.expe_no}_{ext_str}.pdf', dpi = 300)
            plt.close(fig)

    def do_more(self, to_generate):
        """Plot multiple boxplots.

        Args:
            to_generate: bool. Generate data if True. Load data if False.
        """
        dict_columns = dict()
        dict_columns.update({
            "df": (
                r'$\df/n$',
                r'$\hat{p}/n$',
                r'$\hat{n}/n$',
                r'$\trace[V]/n$',
                r'$|\df-\trace[\Sigma A]\trace[V]|/n$',
                r'$\lambda$',
                ),
            "A": (
                r'$\trace[\Sigma A]$',
                r'$\df/\trace[V]$',
                r'$|\trace[\Sigma A]-\df/\trace[V]|$',
                r'$\lambda$',
            ),
        self.plot(to_generate = to_generate, dict_columns = dict_columns)

if __name__ == "__main__":

    Expe6585().do_more(to_generate=False)
