"""Metrics for use when analyzing the results of an EA run.

Basically, these classes allow us to think of an analysis as a
graph of nodes, where there are certain actions we can perform on
each type of data to produce new, processed types of data.

Each class is a thin wrapper around a DataFrame, and we assume that
the raw data comes in a format as generated by LEAP's
FitessStatesCSVProbe operator.
"""
import csv
import os
import shutil
import sys
import tempfile

import click
from matplotlib import patches
from matplotlib import pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import toolz

from leap_ec.probe import FitnessStatsCSVProbe


##############################
# Function best()
##############################
@toolz.curry
def max_value(df, time_col, metric_col):
    """Compute the best (highest) fitness value found across a group's 'metric_col' attribute."""
    return pd.Series({'max_value': df[metric_col].max()})


##############################
# Function auc()
##############################
@toolz.curry
def auc(df, time_col, metric_col):
    """A fast area-under-curve function."""
    df = df[[time_col, metric_col]].sort_values(time_col)
    df['previous_time'] = [df[time_col].min()] + df[time_col][0:len(df[time_col]) - 1].to_list()
    df['xdiff'] = df[time_col] - df.previous_time
    auc = sum(df.xdiff * df[metric_col])
    return pd.Series({'auc': auc})


##############################
# Class CurveAnalyzer
##############################
class CurveAnalyzer():
    """
    There are four fundamental ways that we often analyze performance curves:

     1. plotting the raw curves for all runs, grouped by each experimental configuration,
     2. averaging each curve across runs to get one mean curve for each experiment,
     3. averaging each curve across time to get a scalar performance metric for each run,
     4. averaging each curve across both runs and time to get an average scalar metrics for each experiment.

    This class takes the data for (1) as input, and provides functions for creating
    dataframes for (2), (3), and (4).

    It accepts a dataframe following the conventions used by `FitnessStatsCSVProbe`.
    That is, we have four kinds of columns:

     1. a `time_col` to serve as the "x axis" (ex. steps, generations, evals, or births),
     2. a list of `metric_cols` that measure performance at each time step (ex. best-so-far fitness),
     3. a `run_col` that uniquely identifies each multivariate time series (run), and
     3. a list of `experiment_cols` indicating the experimental configuration of the runs (such 
        as the name of the algorithm used, or the value of various hyperparameters).
    """

    def __init__(self, df,
                 time_col: str = FitnessStatsCSVProbe.time_col,
                 metric_cols: list = FitnessStatsCSVProbe.default_metric_cols,
                 experiment_cols: list = (),
                 run_col: str = 'job'):
        assert(df is not None)
        assert(time_col is not None)
        assert(metric_cols is not None)
        assert(len(metric_cols) > 0)
        assert(experiment_cols is not None)
        assert(len(experiment_cols) > 0)
        assert(run_col is not None)
        self.df = df
        self.time_col = time_col
        self.metric_cols = list(metric_cols)
        self.experiment_cols = list(experiment_cols)
        self.run_col = run_col
        assert(time_col in df.columns), f"Expected '{time_col}' column to be in {df.columns}."
        assert(run_col in df.columns), f"Expected '{run_col}' column to be in {df.columns}."
        assert(set(metric_cols).issubset(df.columns)), f"Expected '{metric_cols}' to be a subset of {df.columns}."
        assert(set(experiment_cols).issubset(df.columns)), f"Expected '{experiment_cols}' to be a subset of {df.columns}."

    def avg_curves(self):
        """Return a dataframe that reports the mean and std of each metric for each experimental group,
        averaged over all of the runs.

        This is analogous to scalar_metrics_per_run, except we're aggregating over runs, instead time.
        """
        # We'll be grouping over time and the experiment columns
        group_cols = [self.time_col] + self.experiment_cols

        # Drop the run_col column, since we're aggregating over it
        df_minus_jobs = self.df[group_cols + self.metric_cols]

        # Compute mean and std within each group
        avg_df = df_minus_jobs.groupby(group_cols).agg(['mean', 'std'])

        # Convert the names of the new columns from (metric, mean) to 'metric_mean', etc.
        avg_df.columns = ['_'.join(col).strip() for col in avg_df.columns.values]
        avg_df.reset_index(inplace=True)
        return avg_df

    def scalar_metrics_per_run(self, metric_col: str, scalar_measure=auc):
        """Compute one or more scalar peformance measure of each metric for each run.

        This is analogous to avg_curves, except we're aggregating over time, instead of across runs.
        
        For example, we might use this to take a set of best-so-far curves and reduce
        them to a series convergence times, or area-under-curve score (one value per run)."""

        # We'll be grouping over the run and experiment columns
        group_cols = [self.run_col] + self.experiment_cols

        # Drop the time_col column, since we're aggregating over it
        #df_minus_time = self.df[group_cols + self.metric_cols]

        # Tell the measure which columns we're using
        parameterized_measure = scalar_measure(time_col=self.time_col, metric_col=metric_col)

        df_scalars = self.df.groupby(group_cols).apply(parameterized_measure)
        # df_scalars = df[group + ['eval', 'bsf']].groupby(group).apply(best)
        # df_auc = df[group + ['eval', 'bsf']].groupby(group).apply(auc)
        # df_scalars['auc'] = df_auc.auc
        df_scalars.reset_index(inplace=True)

        # These are the new columns that we've created via aggregation
        scalar_metric_col = scalar_measure.__name__
        assert(set(group_cols).issubset(df_scalars.columns)), f"Expected {group_cols} to be a subset of {df_scalars.columns}."
        return df_scalars

    def avg_scalar_metrics(self, metric_col :str, scalar_measure=auc):
        """
        Aggregate over both runs and time, computing the mean and std of several scalar
        metrics across runs.

        For example, we can use this to compute the average area-under-best-so-far-curve
        for a collection of runs.
        """
        df_scalars = self.scalar_metrics_per_run(metric_col=metric_col, scalar_measure=scalar_measure)
        scalar_metric_col = scalar_measure.__name__
        assert(scalar_metric_col in df_scalars.columns), f"Expected '{scalar_metric_col}' to appear in {df_scalars.columns}."


        # We'll be grouping over only the experiment columns
        group_cols = self.experiment_cols

        # Drop the run_col column, since we're aggregating over it
        df_minus_jobs = df_scalars[group_cols + [scalar_metric_col]]

        # Compute mean and std of the scalar metrics within each group
        avg_df = df_minus_jobs.groupby(group_cols).agg(['mean', 'std'])

        # Convert the names of the new columns from (metric, mean) to 'metric_mean', etc.
        avg_df.columns = ['_'.join(col).strip() for col in avg_df.columns.values]
        avg_df.reset_index(inplace=True)
        return avg_df


##############################
# Class CurvePlotter
##############################
class CurvePlotter():
    def __init__(self, analyzer):
        assert(analyzer is not None)
        self.analyzer = analyzer
        plt.style.use('ggplot')

    def plot_curves(self, metric_col, title: str, xlabel: str = None, ylabel: str = None, ylog : bool = False, ax=None):
        """Plot all of the curves in a single image, using the specified metric on the y axis."""
        assert(metric_col in self.analyzer.metric_cols), f"Expected '{metric_col}' column to be in {self.analyzer.metric_cols}."

        if ax is None:
            _, ax = plt.subplots(figsize=(8,6))
        
        cmap = plt.get_cmap("tab10")
        legend_patches = []

        for i, (group_key, group_df) in enumerate(self.analyzer.df.groupby(self.analyzer.experiment_cols)):
            patch = patches.Patch(color=cmap(i), label=group_key)
            legend_patches.append(patch)
            group_df.groupby([self.analyzer.run_col]).plot(x=self.analyzer.time_col, y=metric_col,
                                                           color=cmap(i), ax=ax)

        plt.legend(handles=legend_patches)

        if xlabel:
            plt.xlabel(xlabel)
        if ylabel:
            plt.ylabel(ylabel)
        plt.title(title)
        if ylog:
            plt.yscale('log')

    def plot_avg_curves(self, metric_col, title: str, error_bars: bool, ylim, xlabel: str = None, ylabel: str = None, ylog: bool = False, ax=None):
        """Plot the mean curves in a single image, using the specified metric on the y axis."""
        assert(metric_col in self.analyzer.metric_cols), f"Expected '{metric_col}' column to be in {self.analyzer.metric_cols}."

        if ax is None:
            _, ax = plt.subplots(figsize=(8,6))

        df = self.analyzer.avg_curves()
        time_col = self.analyzer.time_col
        # avg_curves() produces _mean and _std columsn for each metric_col
        mean_col = f"{metric_col}_mean"
        std_col = f"{metric_col}_std"
        assert(mean_col in df.columns), f"Expected '{mean_col}' column to be in {df.columns}."
        assert(std_col in df.columns), f"Expected '{std_col}' column to be in {df.columns}."

        # We want to group the data for each experiment into a separate curve
        group_cols = self.analyzer.experiment_cols

        cmap = plt.get_cmap("tab10")
        legend_patches = []

        # Iterate through each group and plot it
        for i, (name, group_df) in enumerate(df.groupby(group_cols, as_index=False)):
            # Optionally population the legend using the specified column
            #label = group_df[legend_column].iloc[0] if legend_column else name
            patch = patches.Patch(color=cmap(i), label=name)
            legend_patches.append(patch)
            if error_bars:
                group_df.plot(x=time_col, y=mean_col, yerr=std_col, capsize=4, fmt='o-', ax=ax, grid='on', ms=10, color=cmap(i))
            else:
                group_df.plot(x=time_col, y=mean_col, style='o-', ax=ax, grid='on', color=cmap(i))

        plt.legend(handles=legend_patches)

        # If the user set the legend to None, remove it
        # if not legend_column:
        #     plt.gca().get_legend().remove()
            
        if ylim:
            plt.ylim(*ylim)

        if xlabel:
            plt.xlabel(xlabel)
        if ylabel:
            plt.ylabel(ylabel)
        plt.title(title)
        if ylog:
            plt.yscale('log')

    def plot_scalars_bar(self, metric_col: str, scalar_measure=auc, title='Performance by Experimental Group', ax=None):
        assert(metric_col is not None)
        assert(scalar_measure is not None)

        if ax is None:
            _, ax = plt.subplots(figsize=(8,6))

        df = self.analyzer.scalar_metrics_per_run(metric_col=metric_col, scalar_measure=scalar_measure)
        columns = self.analyzer.experiment_cols + [ scalar_measure.__name__ ]

        df[columns].boxplot(by=self.analyzer.experiment_cols, ax=ax)
        plt.title(title)
        plt.suptitle('')  # Git rid of the auto-generated boxplot() sup-title
        plt.xticks(rotation = 45)
        plt.tight_layout()

    def plot_avg_scalars(self, metric, independent_vars: list, line=True):
        assert(independent_vars is not None)

        df = self.analyzer.avg_scalar_metrics(metrics=(metric))

        metric_col = metric.__name__
        assert(metric_col in self.analyzer.metric_cols), f"Expected '{metric_col}' column to be in {self.analyzer.metric_cols}."

        for v in independent_vars:
            cmap = sns.color_palette("coolwarm", as_cmap=True)
            plt.figure()
            x, y = np.array(df[v]), np.array(df[metric])
            plt.scatter(df[v], df[metric], c=df[metric], cmap=cmap)

            # Linear fit
            if line:
                m,b = np.polyfit(x, y, 1) 
                end_points = np.array([min(x), max(x)])
                plt.plot(end_points, m*end_points+b, '--k')

            plt.title(f"{v}")


##############################
# CLI groups
##############################
@click.group()
def cli():
    """Experiments in evolving controllers for simple 2-D robots."""
    pass


@cli.group()
def plot():
    """Plotting commands."""
    pass


##############################
# Command analyze
##############################
@cli.command()
@click.argument('files', nargs=-1)
#@click.option('--time-col', default='step', type=str, help="Name of column that represent time (ex. generation, step, eval).")
#@click.option('--experiment-cols', default='step', type=str, nargs=-1, help="Name of columns that indicate experiment groups (ex. algorithm name, hyperparameter values).")
def analyze(files):
    """
    Analyze a set of performance-curve files by computing various metrics and averages.

    Outputs four new files in the current directory.
    """

    curves_file = './curves_combined.csv'

    # Combine files into one big CSV
    df = cat_csv_files(files)
    df.to_csv(curves_file)

    all_runs_df = pd.read_csv(curves_file, skipinitialspace=True, comment='#')

    analyzer = CurveAnalyzer(df=all_runs_df, experiment_cols=['experiment'])
    analyzer.avg_curves().to_csv('./avg_curves_by_experiment.csv', index=False)
    analyzer.scalar_metrics_per_run(metric_col='bsf', scalar_measure=auc).to_csv('./scalars_by_run.csv', index=False)
    analyzer.avg_scalar_metrics(metric_col='bsf', scalar_measure=auc).to_csv('./avg_scalars_by_experiment.csv', index=False)


##############################
# Command plot all
##############################
@plot.command()
@click.argument('curves-file')
@click.option('--metric-col', default='bsf', type=str, help="Name of column to plot on the y axis (ex. bsf, max_fitness).")
@click.option('--error/--no-error', type=bool, default=True)
@click.option('--ylim', type=(float, float), default=(None, None))
@click.option('--modulo', type=int, default='1')
@click.option('--time-col', default='step', type=str, help="Name of column that represent time (ex. generation, step, eval).")
@click.option('--ylog/--no-ylog', type=bool, default=False)
def all(curves_file, metric_col, error, ylim, modulo, time_col,ylog):
    """Plot a single best-of-generation fitness curve from a CSV file."""
    assert(os.path.exists(curves_file))
    df = pd.read_csv(curves_file, skipinitialspace=True, comment='#')
    df = df[df[time_col] % modulo == 0]  # Select every nth step
    analyzer = CurveAnalyzer(df, time_col=time_col, experiment_cols=['experiment'])
    plotter = CurvePlotter(analyzer)

    plt.figure(figsize=(20, 8))
    plt.subplot(131)
    plotter.plot_curves(metric_col, f"Indiviudal {metric_col} Curves", ylog=ylog, ax=plt.gca())
    plt.subplot(132)
    plotter.plot_avg_curves(metric_col, f"Average {metric_col} Curves", error, ylim, ylog=ylog, ax=plt.gca())
    plt.subplot(133)
    plotter.plot_scalars_bar(metric_col, title='Area Under Curve', ax=plt.gca())
    plt.show()


##############################
# Command plot curves
##############################
@plot.command()
@click.argument('curves-file')
@click.option('--metric-col', default='bsf', type=str, help="Name of column to plot on the y axis (ex. bsf, max_fitness).")
@click.option('--title', type=str, default='Performance Curves')
@click.option('--time-col', default='step', type=str, help="Name of column that represent time (ex. generation, step, eval).")
@click.option('--ylog/--no-ylog', type=bool, default=False)
def curves(curves_file, metric_col, title, time_col, ylog):
    """Plot the individual performance curves from a CSV file."""
    assert(os.path.exists(curves_file))
    df = pd.read_csv(curves_file, skipinitialspace=True, comment='#')
    analyzer = CurveAnalyzer(df, time_col=time_col, experiment_cols=['experiment'])
    plotter = CurvePlotter(analyzer)
    plotter.plot_curves(metric_col, title, ylog=ylog)
    plt.show()


##############################
# Command plot avg-curves
##############################
@plot.command('avg-curves')
@click.argument('average-bsf-file')
@click.option('--metric-col', default='bsf', type=str, help="Name of column to plot on the y axis (ex. bsf, max_fitness).")
@click.option('--title', type=str, default='Average Performance Curves')
@click.option('--error/--no-error', type=bool, default=True)
@click.option('--ylim', type=(float, float), default=(None, None))
@click.option('--time-col', default='step', type=str, help="Name of column that represent time (ex. generation, step, eval).")
@click.option('--ylog/--no-ylog', type=bool, default=False)
def avg_curves(average_bsf_file, metric_col, title, error, ylim, time_col, ylog):
    """Plot average fitness curves from an average-fitness CSV file."""
    assert(os.path.exists(average_bsf_file))
    df = pd.read_csv(average_bsf_file, skipinitialspace=True, comment='#')
    analyzer = CurveAnalyzer(df, time_col=time_col, experiment_cols=['experiment'])
    plotter = CurvePlotter(analyzer)
    plotter.plot_avg_curves(metric_col, title, error, ylim, ylog=ylog)
    plt.show()


##############################
# Command plot auc
##############################
@plot.command('auc')
@click.argument('curves-file')
@click.option('--metric-col', default='bsf', type=str, help="Name of column to plot on the y axis (ex. bsf, max_fitness).")
@click.option('--title', type=str, default='Area-under-curve by experiment')
@click.option('--time-col', default='step', type=str, help="Name of column that represent time (ex. generation, step, eval).")
def auc_command(curves_file, metric_col, title, time_col):
    """Plot a single best-of-generation fitness curve from a CSV file."""
    assert(os.path.exists(curves_file))
    df = pd.read_csv(curves_file, skipinitialspace=True, comment='#')
    analyzer = CurveAnalyzer(df, time_col=time_col, experiment_cols=['experiment'])
    plotter = CurvePlotter(analyzer)
    plotter.plot_scalars_bar(metric_col, title=title)
    plt.show()


##############################
# Function cat_csvs()
##############################
def cat_csv_files(files: list):
    """Combine a bunch of CSVs with the same columns into one CSV with a single header.
    
    The result is written to the given stream."""
    # TODO Use pandas instead of csv to read input, because we like to use CSV files that have comments at the top
    # header = None
    # writer = csv.writer(stream)

    dataframes = []
    for fname in files:
        f_df = pd.read_csv(fname, skipinitialspace=True, comment='#')
        dataframes.append(f_df)

    return pd.concat(dataframes)
        # with open(fname, 'r') as f:
        #     reader = csv.reader(f)

        #     if header is None:
        #         header = next(reader)
        #         writer.writerow(header)
        #     else:
        #         this_header = next(reader)
        #         if this_header != header:
        #             raise ValueError(f"Inconsistent headers.  First file was:\n{header}\n\nbut found a file with\n{this_header}.")

        #     for row in reader:
        #         writer.writerow(row)


##############################
# Entry point
##############################
if __name__ == '__main__':
    cli()