#!/usr/bin/env -S ipython --no-banner -i --
"""
Quick performance plots used as a script to visualize results.

This is wonderfully useful for analyzing csv files of models!

  - Expect a csv header with "epoch" as a column (customizable via --index-col)
  - Requires a directory structure:
      ./results/
                run_id/perf.csv  # for mode 1 and 4
                      /log/perf*.csv  # for other modes
      where:
          - ./results/ can be changed via  --data-dir
          - "run_id" is any directory name
          - "perf.csv" and "perf*.csv" can be changed  via --data-fp-regex
          - the "log" directory is hardcoded...

Usage:
    $ python -m simplepytorch.plot_perf  -h

    # or if simplepytorch is installed:
    $ simplepytorch_plot -h

    # or alias it in a bash shell:
    $ alias sdp='python -m simplepytorch.perf_plot -- '
"""
import argparse as ap
import datetime as dt
import IPython
import re
import glob
import os
from os.path import join, basename, dirname
import pandas as pd
import numpy as np
import seaborn as sns
from textwrap import dedent
from matplotlib import pyplot as plt
import warnings
#  import mpld3
#  import mpld3.plugins


def get_run_ids(ns):
    dirs = glob.glob(join(ns.data_dir, '*'))
    for dir in dirs:
        run_id = basename(dir)
        if not re.search(ns.runid_regex, basename(dir)):
            continue
        yield run_id


def load_df_from_fp(fp, ns):
    print('found data', fp)
    try:
        if fp.endswith('.csv') or fp.endswith('.csv.gz'):
            df = pd.read_csv(fp).set_index(ns.index_col)
        elif fp.endswith('.h5'):
            df = pd.read_hdf(fp, ns.hdf_table)
        else:
            print('bug: file does not exist:', fp)
    except Exception as e:
        print(f"   WARNING: failed to load {fp}.  {e}")
        df = pd.DataFrame()
    return df


def _mode1_get_frames(ns):
    for run_id in get_run_ids(ns):
        #  if not exists(join(dir, 'lock.finished')):
        #      continue
        dirp = f'{ns.data_dir}/{run_id}'
        gen = (
            ((run_id, fname), load_df_from_fp(join(dirp, fname), ns))
            for fname in os.listdir(dirp)
            if re.match(ns.data_fp_regex, fname))
        yield from (x for x in gen if not x[-1].empty)


def _get_all_data(ns):
    cdfs_data = {}
    # date the plot was created.  nothing to do with timestamp column.
    #  timestamp = dt.datetime.utcnow().strftime('%Y%m%dT%H%M%S')
    for run_id in get_run_ids(ns):
        dirp = f'{ns.data_dir}/{run_id}/log'
        if not os.path.exists(dirp):
            print('skip', run_id, 'contains no log data')
            continue

        cdfs = pd.concat({
            (run_id, fname): load_df_from_fp(join(dirp, fname), ns)
            for fname in os.listdir(dirp)
            if re.search(f'{ns.data_fp_regex}', fname)},
            sort=False, names=['run_id', 'filename']
        )
        cdfs_data[run_id] = cdfs
    cdfs = pd.concat(cdfs_data.values())
    return cdfs


def _get_subplots(N:int, sharey:bool):
    if N < 5:
        M = 1
    else:
        M = int(np.sqrt(N))
        N = int(np.ceil(N/M))
    fig, axs = plt.subplots(N, M, figsize=(12,10), squeeze=False, sharex=True,
                            sharey=sharey)
    return fig, axs


def _subplots_add_legend(fig, handles_labels):
    # print(handles_labels)
    fig.tight_layout()
    fig.subplots_adjust(bottom=0.1)
    # choose just one axis as the legend, since they are all identical
    fig.legend(*handles_labels,
               loc='lower center', bbox_to_anchor=(0.5, -0.0),
               fancybox=False, shadow=False, ncol=3)


def make_plots(ns, cdfs, subplots:bool):
    cols = [col for col in cdfs.columns if re.search(ns.col_regex, col)]
    if ns.avg_cols:
        colname = 'Mean(%s)' % ns.col_regex
        cdfs[colname] = cdfs[cols].mean(axis=1)
        cols = [colname]

    plot_cols = [prettify(col) for col in cols]
    cdfs = cdfs.copy()
    cdfs.columns = [prettify(x) for x in cdfs.columns]
    cdfs.index.names = [prettify(x) for x in cdfs.index.names]
    if subplots:
        fig, axs = _get_subplots(len(plot_cols), sharey=bool(ns.ylim))
    for n, col in enumerate(plot_cols):
        if subplots:
            ax = axs.ravel()[n]
        else:
            fig, ax = plt.subplots(1,1)
        df = cdfs[col].unstack(level=0)
        if ns.rolling_mean:
            df = df.rolling(ns.rolling_mean).mean()
        if (df.dtypes.values == np.dtype('O')).all():
            print('WARNING: skip col with no numeric data:', col)
            ax.set_title(f'{col}, SKIPPED because no numeric data')
            continue
        _plot_lines = df.plot(
            ax=ax, title=col, legend=None if subplots else ns.legend, alpha=.9)
        #  _legend = mpld3.plugins.InteractiveLegendPlugin(
        #      *_plot_lines.get_legend_handles_labels())
        #  mpld3.plugins.connect(fig, _legend)
        if ns.ylim:
            ax.set_ylim(*ns.ylim)
        if not subplots:
            yield (fig, col)
    if subplots:
        if ns.legend:
            _subplots_add_legend(fig, ax.get_legend_handles_labels())
        yield (fig, ns.col_regex)


def savefig_with_symlink(fig, fp, symlink_fp):
    os.makedirs(dirname(fp), exist_ok=True)
    fig.savefig(fp, bbox_inches='tight')
    if os.path.islink(symlink_fp):
        os.remove(symlink_fp)
    prefix = os.path.dirname(os.path.commonprefix([fp, symlink_fp]))
    os.symlink(fp[len(prefix)+1:], symlink_fp)
    print('save fig', symlink_fp)


def main(ns):
    print(ns)
    # mode 1: compare each column across all files
    if ns.mode == 0:
        cdfs = mode_0(ns)
    elif ns.mode == 1:
        cdfs = mode_1_plots(ns)
    elif ns.mode == 2:
        cdfs = mode_2_plots(ns)
    elif ns.mode == 3:
        cdfs = mode_3_plots(ns)
    elif ns.mode == 4:
        cdfs = mode_4_plots(ns)
    else:
        raise Exception(f'not implemented mode: {ns.mode}')
    if ns.no_plot:
        print("type 'plt.show()' in IPython terminal to see result")
        ipython_start_cmd = []
    else:
        ipython_start_cmd = ['-c', 'plt.ion() ; plt.show(block=False)', '-i']
        #  plt.show(block=False)
    if ns.shell:
        dct = dict(locals())
        dct.update(globals())
        print("\n==== Global Namespace: ====\n{1}\n\n==== Local Namespace: ====\n{0}\n".format(
            '\n'.join([x for x in locals().keys() if not x.startswith('__')]),
            ', '.join(globals().keys())))
        IPython.start_ipython(['--no-banner'] + ipython_start_cmd, user_ns=dct)
    return cdfs


def mode_0(ns):
    cdfs = _get_all_data(ns)
    return cdfs


def _mode_1_and_4_get_perf_data_as_df(ns):
    dfs = dict(_mode1_get_frames(ns))
    cdfs = pd.concat(dfs, sort=False, names=['run_id', 'filename'])
    return cdfs


def mode_1_plots(ns):
    """Compare experiments.
    One plot for each metric, comparing the most recent result of each experiment
    """
    cdfs = _mode_1_and_4_get_perf_data_as_df(ns).reset_index('filename', drop=True)

    timestamp = dt.datetime.utcnow().strftime('%Y%m%dT%H%M%S.%f')  # date plot was created.  nothing to do with timestamp column.
    for fig, col in make_plots(ns, cdfs, ns.make_subplots):
        if ns.savefig and ns.make_subplots:
            # would need, I think, to make each column a subfigure instead of
            # axes plot.
            warnings.warn(
                "Sorry, but --mode1-subplots and --savefig haven't been made"
                " to work nicely together yet.")
        if not ns.savefig: continue
        savefig_with_symlink(
            fig,
            f'{ns.savefig_dir}/{ns.runid_regex}/mode1_{timestamp}__{col}.png',
            f'{ns.savefig_dir}/{ns.runid_regex}/mode1_latest__{col}.png')
    return cdfs


def mode_2_plots(ns):
    """Analyze an experiment's historical performance across many runs.
    One plot for each experiment and metric, to analyze history of runs for
    that experiment"""
    # mode 2: compare train to val performance
    cdfs_mode2 = {}
    timestamp = dt.datetime.utcnow().strftime('%Y%m%dT%H%M%S.%f')  # date plot was created.  nothing to do with timestamp column.
    for run_id in get_run_ids(ns):
        dirp = f'{ns.data_dir}/{run_id}/log'
        if not os.path.exists(dirp):
            print('skip', run_id, 'contains no log data')
            continue

        cdfs = pd.concat({
            (run_id, fname): load_df_from_fp(join(dirp, fname), ns)
            for fname in os.listdir(dirp)
            if re.search(f'{ns.data_fp_regex}', fname)},
            sort=False, names=['run_id', 'filename']
        )
        cdfs_mode2[run_id] = cdfs

        for fig, col in make_plots(ns, cdfs.reset_index('run_id', drop=True), ns.make_subplots):
            fig.suptitle(run_id)
            if not ns.savefig: continue
            # save to file
            savefig_with_symlink(
                fig,
                f'{ns.savefig_dir}/{run_id}/mode2_{timestamp}__{col}.png'.format(run_id=run_id),
                f'{ns.savefig_dir}/{run_id}/mode2_latest__{col}.png'.format(run_id=run_id))
    return cdfs_mode2


def prettify(x:str):
    if x is None:
        return x
    x = x.replace('_', ' ').split(' ')
    for i,xx in enumerate(x):
        if not any(y.isupper() for y in xx):
            x[i] = xx.title()
    return ' '.join(x)


def mode_3_plots(ns):
    """Compare across experiments, considering their history of runs.
    Basically combines mode 1 and mode 2. """
    cdfs = _get_all_data(ns)
    cols = [col for col in cdfs.columns if re.search(ns.col_regex, col)]
    if ns.avg_cols:
        colname = 'Mean(%s)' % ns.col_regex
        cdfs[colname] = cdfs[cols].mean(axis=1)
        cols = [colname]

    #  cdfs.groupby(['run_id', 'epoch']).agg(ns.mode3_agg_method)
    cdfs_unmodified = cdfs.copy()
    if ns.rolling_mean:
        cdfs = cdfs.rolling(ns.rolling_mean).mean()
    cdfs.columns = [prettify(col) for col in cdfs.columns]
    cdfs.index.names = [prettify(x) for x in cdfs.index.names]

    timestamp = dt.datetime.utcnow().strftime('%Y%m%dT%H%M%S.%f')  # date plot was created.  nothing to do with timestamp column.
    plot_cols = [prettify(col) for col in cols]

    if ns.make_subplots:
        fig, axs = _get_subplots(len(plot_cols), sharey=bool(ns.ylim))
    for n, col in enumerate(plot_cols):
        if ns.make_subplots:
            ax = axs.ravel()[n]
        else:
            fig, ax = plt.subplots(1,1)
        kws = dict(x=prettify(ns.index_col), y=col, hue='Run Id', ax=ax, data=cdfs)
        z = sns.lineplot(**kws, errorbar=('ci', ns.mode3_ci) if ns.mode3_ci else None, lw=2 if ns.mode3_ci == 0 else 1, legend=ns.legend)
        if ns.mode3_ci == 0:
            sns.lineplot(**kws, units="Filename", estimator=None, lw=1, alpha=.5)#, legend=False)
        ax.set_title(col)
        if ns.ylim:
            ax.set_ylim(*ns.ylim)
        if ns.legend and ns.make_subplots:
            _handles, _labels = ax.get_legend_handles_labels()
            ax.get_legend().remove()
            # note: the legend labels may be incorrect if --mode4-groupby is not
            # default value.
            _subplots_add_legend(fig, ax.get_legend_handles_labels())

        if ns.savefig and not ns.make_subplots:
            savefig_with_symlink(
                fig,
                f'{ns.savefig_dir}/{ns.runid_regex}/mode3_{timestamp}__{col}.png',
                f'{ns.savefig_dir}/{ns.runid_regex}/mode3_latest__{col}.png')
    if ns.savefig and ns.make_subplots:
            savefig_with_symlink(
                fig,
                f'{ns.savefig_dir}/{ns.runid_regex}/mode3_{timestamp}__{ns.col_regex}.png',
                f'{ns.savefig_dir}/{ns.runid_regex}/mode3_latest__{ns.col_regex}.png')
    return cdfs_unmodified


def _mode4_make_subplots(grp:pd.DataFrame, ns):
    regex, grp_num = ns.mode4_groupby.rsplit(':', 1)
    run_id = grp.index.to_frame()['Run Id'].unique()
    assert len(run_id) == 1, 'code bug'
    run_id = run_id[0]
    if ns.mode4_groupby is None:
        colgrps = [run_id] * len(grp.columns)
    else:
        colgrps = []
        match_error = False
        for col,filename in grp.columns:
            match = re.search(regex, col)
            if match is None:
                colgrps.append(None)
                print(f'--mode4-groupby: No match found for "{col}" given regex: "{regex}"')
                match_error = True
                continue
            try:  # some validation to give a helpful error message
                match.group(int(grp_num))
            except IndexError:
                msg = (
                    f'    match object: {match} '
                    f'\n    regex: "{regex}" '
                    f'\n    matched groups: {match.groups()} '
                    f'\n    col: {col}')
                raise Exception(
                    '--mode4-groupby had wrong group number :N. '
                    ' Helpful detail: \n' + msg)
            colgrps.append(match.group(int(grp_num)))
        if match_error:
            raise Exception(
                '--mode4-groupby regex did not match one or'
                ' more columns (see previously printed msgs)')
    assert len(colgrps) > 0, 'code bug'
    # choose nrows, ncols for the figure
    fig, axs = _get_subplots(N=len(set(colgrps)), sharey=bool(ns.ylim))
    colgrps = np.array(colgrps)
    for ax, subplot_title in zip(axs.ravel(), sorted(set(colgrps))):
        cols = grp.columns[colgrps == subplot_title]
        z = grp[cols].reset_index('Run Id', drop=True)
        assert z.columns.names[1] == 'Filename', 'code bug'
        z.columns.set_names([subplot_title, 'Filename'], inplace=True)
        z = z.copy()  # necessary to fix legend
        if ns.mode4_allruns == -1:
            # actually don't use filename, but keep it here in case want to extend
            # mode 4 to support a mode 2 or 3 style of lookback.
            z.columns = z.columns.droplevel('Filename')
        elif ns.mode4_allruns is not None:  # consider a subset of previous runs
            filenames = set(sorted(z.columns.levels[1])[-1*ns.mode4_allruns:])
            cols = [x for x in z.columns if x[1] in filenames]
            z = z[cols]
        if (z.dtypes.values == np.dtype('O')).all():
            print('WARNING: skip col with no numeric data:', col)
            ax.set_title(f'{col}, SKIPPED because no numeric data')
            continue
        z.plot(ax=ax, title=subplot_title)
        # print(ns)
        if ns.legend:
            _handles, _labels = ax.get_legend_handles_labels()
        ax.get_legend().remove()
        if ns.ylim:
            ax.set_ylim(*ns.ylim)
        #  if not ns.legend:
            #  ax.get_legend().remove()
    if ns.legend:
        # note: the legend labels may be incorrect if --mode4-groupby is not
        # default value.
        _subplots_add_legend(
            fig, (_handles, [re.search(regex, x).group(1) for x in _labels]))

    fig.suptitle(f'Experiment: {run_id}')
    fig.tight_layout()
    if ns.savefig:
        timestamp = dt.datetime.utcnow().strftime('%Y%m%dT%H%M%S')  # date plot was created.  nothing to do with timestamp column.
        savefig_with_symlink(
            fig,
            f'{ns.savefig_dir}/{run_id}/mode4_{timestamp}__{ns.col_regex}.png',
            f'{ns.savefig_dir}/{run_id}/mode4_latest__{ns.col_regex}.png')


def mode_4_plots(ns):
    """For each experiment, compare across columns. One plot per experiment"""
    print(ns.mode4_allruns)
    if ns.mode4_allruns == -1:
        cdfs = _mode_1_and_4_get_perf_data_as_df(ns)
    else:
        cdfs = _get_all_data(ns)
    cols = [col for col in cdfs.columns if re.search(ns.col_regex, col)]
    if ns.avg_cols:
        colname = 'Mean(%s)' % ns.col_regex
        cdfs[colname] = cdfs[cols].mean(axis=1)
        cols = [colname]
    cdfs_unmodified = cdfs.copy()
    cdfs = cdfs[cols]
    if ns.rolling_mean:
        cdfs = cdfs.rolling(ns.rolling_mean).mean()
    cdfs.columns = [prettify(col) for col in cdfs.columns]
    cdfs.index.names = [prettify(x) for x in cdfs.index.names]

    cdfs.unstack('Filename').groupby('Run Id').apply(
        lambda grp: _mode4_make_subplots(grp, ns))
    return cdfs_unmodified


def bap():
    class F(ap.ArgumentDefaultsHelpFormatter, ap.RawTextHelpFormatter): pass
    d = ('''
Plot results of experiments. Expect a default directory structure like:

        ./results/               # this default location is adjustable via --data-dir
            experiment_1a/
                log/perf*.csv    # used by --mode 2 and --mode 3
                                 # to compare multiple csv files within this experiment

                perf.csv         # used by --mode 1 and --mode 4
                                 # to compare one csv file representing each experiment
                                 # default filename is adjustable via --data-fp-regex.
            experiment_1b/
                perf.csv
            another_experiment/
                perf.csv
            ...

    where the CSV files should share the same header, and the header must
    include a column "epoch" (adjustable via --index-col).

    For example:
      - One has a history of runs of the form "results/experiment_1a/log/perf_{timestamp}.csv"
      - And a symlink 'results/experiment_1a/perf.csv -> log/perf_{latest_timestamp}.csv'
    For example:
      - One does K-fold cross validation and gets K csv files:  "results/experiment_1a/log/perf_{kfold_split_idx}.csv"
      - And a single perf.csv file summarizing a model evaluated on an out-of-distribution test dataset "results/experiment_1a/perf.csv"

Example command-line usage:

    # Plot "epoch vs {XYZ}" for data from experiment_1a and experiment_1b,
    # where {XYZ} is any column matching "Loss".

    python -m simplepytorch.plot_perf  'experiment_1' -c 'Loss'

    # note: the --mode option is very useful.

    ''')
    par = ap.ArgumentParser(description=d, formatter_class=F)
    A = par.add_argument
    A('runid_regex', help='find the run_ids matching the regular expression. run_id is the identifier for an experiment containing data')
    A('--data-dir', default='./results', help=' ')
    A('--data-fp-regex', default=r'perf.*\.csv(\.gz)?$', help=dedent('''\
    only consider data files matching the specified criteria:
    ./{data_dir}/{runid_regex}/{data_fp_regex}.'''))
    #  A('--', nargs='?', default='one', const='two')
    A('--rolling-mean', '--rm', type=int, help='smoothing')
    A('--ylim', nargs=2, type=float, help='Optional. Two numbers denoting the the y limit shared by all plots. Example: --ylim 0 1')
    A('--hdf-table-name', help='required if searching .h5 files')
    A('-c', '--col-regex', default='^(?!epoch|batch_idx|timestamp)', help='plot only columns matching regex.  By default, plot all except epoch and batch_idx.')
    A('--index-col', default='epoch',
      help='The x-axis of the plots.  All data files must contain this column.')
    A('--mode', default=1, type=int, choices=[0,1,2,3,4], help=dedent('''\
        `--mode 0` Don't plot anything, just collect data and drop into a shell.
        `--mode 1` Compare across experiments, with one plot per column (i.e. performance metric)
        `--mode 2` Within one experiment and column, visualize the history of all runs.
        `--mode 3` Combine 1 and 2.  compare across run_ids, with
                   confidence interval to consider history of runs for each run_id.
        `--mode 4` Compare across columns, one plot per experiment.
                   '''))
    A('--savefig-dir', default='./results/plots/',
      help='Default location to save figures if --savefig is defined')
    A('--mode3-ci', type=float,
      help='confidence interval for error bars.  If 0, also shows all lines individually')
    A('--make-figures', dest='make_subplots', action='store_false', help=dedent('''\
        If defined, generate individual figures instead of a figure with subplots.'''))
    A('--mode4-groupby', default="(Train|Val|Test) (.*):2", help=dedent('''\
        Optionally, make one sub-plot for each group of columns. To use, pass a
        regular expression extracting a string that we can group related
        columns followed by a ":N" where N is the one-indexed N-th capturing
        group to use.
        Example:  `--mode4-groupby "(Train|Val|Test) (.*):2"
          - this would place "train_loss" and "val_loss" into a subplot, and
            "train_acc" and "val_acc" in a different subplot.
        Note: This regex applies to the prettified text. The error messages
        will help you figure out what that means.
    '''))
    A('--mode4-allruns', nargs='?', type=int, default=-1, help=dedent('''\
        If supplied, give data on N 'most recent' runs for the given experiment,
        where 'most recent' means sort the csv filenames A-Z and take the last N.
        Example:  `--mode4-allruns 2` just shows the last two runs.'''))
    A('--best-effort', action='store_true',
      help=" Try to load a csv file, but don't raise error if cannot read a file.")
    A('--savefig', action='store_true', help=dedent('''\
        Save plots to file.  Can be useful with --no-plot --no-shell.
        Also can be useful with --make-figures.'''))
    A('--no-plot', action='store_true', help="if supplied, don't show the plots")
    A('--no-shell', action='store_false', dest='shell',
      help='do not drop into an IPython shell when finished')
    A('--no-legend', action='store_false', dest='legend', help='do not show legend')
    A('--avg-cols', action='store_true', help=(
        'if supplied, average the columns of the input csv file,'
        ' where columns are defined by -c parameter'))
    return par


if __name__ == "__main__":
    main(bap().parse_args())
