"""
Rank summarization results.
"""
import os
import sys
import time
import argparse
from datetime import datetime
from itertools import product

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from scipy.stats import sem
from scipy.stats import gmean
from tqdm import tqdm

here = os.path.abspath(os.path.dirname(__file__))
sys.path.insert(0, here + '/../')
from postprocess import util as pp_util
from experiments import util as exp_util
from config import rank_args
from remove import get_mean_df


def process(args, exp_hash, out_dir, logger):
    begin = time.time()

    color, line, label = pp_util.get_plot_dicts()

    df_list = []
    df_li_list = []
    df_time_list = []
    df_mem_list = []

    for tree_type in args.tree_type:

        in_dir = os.path.join(args.in_dir,
                              tree_type,
                              f'exp_{exp_hash}',
                              'summary')

        # get resource usage
        ckpt_dir = os.path.join(in_dir, f'ckpt_{args.ckpt[0]}')

        fp_time = os.path.join(ckpt_dir, 'runtime.csv')
        fp_mem = os.path.join(ckpt_dir, 'mem.csv')
        assert os.path.exists(fp_time), f'{fp_time} does not exist!'
        assert os.path.exists(fp_mem), f'{fp_mem} does not exist!'

        df_time_list.append(pd.read_csv(fp_time))
        df_mem_list.append(pd.read_csv(fp_mem))

        # get loss
        for ckpt in args.ckpt:
            ckpt_dir = os.path.join(in_dir, f'ckpt_{ckpt}')

            fp = os.path.join(ckpt_dir, 'loss_rank.csv')
            fp_li = os.path.join(ckpt_dir, 'loss_rank_li.csv')
            assert os.path.exists(fp), f'{fp} does not exist!'
            assert os.path.exists(fp_li), f'{fp_li} does not exist!'

            df_list.append(pd.read_csv(fp))
            df_li_list.append(pd.read_csv(fp_li))

    df_all = pd.concat(df_list)
    df_li_all = pd.concat(df_li_list)
    df_time_all = pd.concat(df_time_list)
    df_mem_all = pd.concat(df_mem_list)

    # average ranks among different checkpoints and/or tree types
    group_cols = ['dataset']

    df_all = df_all.groupby(group_cols).mean().reset_index()
    df_li_all = df_li_all.groupby(group_cols).mean().reset_index()
    df_time_all = df_time_all.groupby(group_cols).mean().reset_index()
    df_mem_all = df_mem_all.groupby(group_cols).mean().reset_index()

    # compute average ranks
    skip_cols = ['dataset', 'tree_type', 'poison_frac']

    df = get_mean_df(df_all, skip_cols=skip_cols, sort='ascending')
    df_li = get_mean_df(df_li_all, skip_cols=skip_cols, sort='ascending')

    logger.info(f'\nLoss:\n{df}')
    logger.info(f'\nLoss (li):\n{df_li}')

    # # combine dataframes
    # index = df_li.index
    # df = df_li.reset_index().merge(df.reset_index(), on='index', how='left')
    # means_df = df[['index', 'mean_x', 'mean_y']].copy()
    # sems_df = df[['index', 'sem_x', 'sem_y']].copy()

    # # rename and clean up
    # means_df.index = means_df['index']
    # sems_df.index = means_df['index']
    # del means_df['index']
    # del sems_df['index']
    # means_df.columns = ['Subgroup A', 'All datasets']
    # sems_df.columns = ['Subgroup A', 'All datasets']

    # print(means_df)
    # print(sems_df)

    # # plot
    # fig, ax = plt.subplots(figsize=(4, 4))
    # means_df.plot.bar(yerr=sems_df, ax=ax, rot=45,
    #                   title=f'Loss ({len(means_df)} datasets)', capsize=3,
    #                   ylabel='Avg. rank', xlabel='Method')

    # plot
    n_datasets = len(df_all['dataset'].unique())
    n_li_datasets = len(df_li_all['dataset'].unique())

    label_dict = {'LeafInfluence': 'LeafInf.', 'SubSample': 'SubS.', 'Target': 'RandomSL'}

    df = df.rename(columns={'mean': 'All datasets'}, index=label_dict)
    df_li = df_li.rename(columns={'mean': 'SDS'}, index=label_dict)

    # reorder methods
    order = ['BoostIn', 'LeafInfSP', 'TreeSim', 'TREX', 'SubS.', 'LOO', 'RandomSL', 'Random']
    order_li = ['LeafRefit', 'LeafInf.', 'BoostIn', 'LeafInfSP', 'TreeSim', 'TREX', 'SubS.', 'LOO',
                'RandomSL', 'Random']

    # order_li = ['BoostIn', 'LeafInfSP', 'TreeSim', 'TREX', 'SubS.', 'LOO', 'RandomSL', 'Random',
    #             'LeafRefit', 'LeafInf.']

    df = df.reindex(order)
    df_li = df_li.reindex(order_li)

    labels = [c if i % 2 != 0 else f'\n{c}' for i, c in enumerate(df.index)]
    labels_li = [c if i % 2 != 0 else f'\n{c}' for i, c in enumerate(df_li.index)]

    pp_util.plot_settings(fontsize=28)
    width = 22
    height = pp_util.get_height(width, subplots=(1, 2))

    fig, axs = plt.subplots(1, 2, figsize=(width, height), gridspec_kw={'width_ratios': [6, 8]})

    ax = axs[0]
    df.plot(kind='bar', y='All datasets', yerr='sem', ax=ax, title=None, capsize=3,
            ylabel='Average rank', xlabel=None, legend=True, color='#3e9ccf')
    ax.set_xticklabels(labels, rotation=0)

    ax = axs[1]
    df_li.plot(kind='bar', y='SDS', yerr='sem', ax=ax, title=None, capsize=3,
               ylabel=None, xlabel=None, legend=True, color='#ff7600')
    ax.set_xticklabels(labels_li, rotation=0)

    ax.axvline(1.5, color='gray', linestyle='--')
    # ax.axvline(7.5, color='gray', linestyle='--')

    logger.info(f'\nSaving results to {out_dir}/...')

    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, 'roar.pdf'), bbox_inches='tight')

    df.to_csv(os.path.join(out_dir, 'loss_rank.csv'))
    df_li.to_csv(os.path.join(out_dir, 'loss_rank_li.csv'))

    logger.info(f'\nTotal time: {time.time() - begin:.3f}s')


def main(args):

    exp_dict = {'n_test': args.n_test, 'poison_frac': args.poison_frac}
    exp_hash = exp_util.dict_to_hash(exp_dict)

    assert len(args.tree_type) > 0
    out_dir = os.path.join(args.in_dir,
                           'rank',
                           f'exp_{exp_hash}',
                           f'+'.join(args.tree_type))

    # create logger
    os.makedirs(out_dir, exist_ok=True)
    logger = exp_util.get_logger(os.path.join(out_dir, 'log.txt'))
    logger.info(args)
    logger.info(f'\ntimestamp: {datetime.now()}')

    process(args, exp_hash, out_dir, logger)


if __name__ == '__main__':
    main(rank_args.get_poison_args().parse_args())

"""
Rank summarization results.
"""
import os
import sys
import time
import argparse
from datetime import datetime
from itertools import product

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from scipy.stats import sem
from scipy.stats import gmean
from tqdm import tqdm

here = os.path.abspath(os.path.dirname(__file__))
sys.path.insert(0, here + '/../')
from postprocess import util as pp_util
from experiments import util as exp_util
from config import rank_args
from remove import get_mean_df


def process(args, exp_hash, out_dir, logger):
    begin = time.time()

    color, line, label = pp_util.get_plot_dicts()

    df_list = []
    df_li_list = []
    df_time_list = []
    df_mem_list = []

    for tree_type in args.tree_type:

        in_dir = os.path.join(args.in_dir,
                              tree_type,
                              f'exp_{exp_hash}',
                              'summary')

        # get resource usage
        ckpt_dir = os.path.join(in_dir, f'ckpt_{args.ckpt[0]}')

        fp_time = os.path.join(ckpt_dir, 'runtime.csv')
        fp_mem = os.path.join(ckpt_dir, 'mem.csv')
        assert os.path.exists(fp_time), f'{fp_time} does not exist!'
        assert os.path.exists(fp_mem), f'{fp_mem} does not exist!'

        df_time_list.append(pd.read_csv(fp_time))
        df_mem_list.append(pd.read_csv(fp_mem))

        # get loss
        for ckpt in args.ckpt:
            ckpt_dir = os.path.join(in_dir, f'ckpt_{ckpt}')

            fp = os.path.join(ckpt_dir, 'loss_rank.csv')
            fp_li = os.path.join(ckpt_dir, 'loss_rank_li.csv')
            assert os.path.exists(fp), f'{fp} does not exist!'
            assert os.path.exists(fp_li), f'{fp_li} does not exist!'

            df_list.append(pd.read_csv(fp))
            df_li_list.append(pd.read_csv(fp_li))

    df_all = pd.concat(df_list)
    df_li_all = pd.concat(df_li_list)
    df_time_all = pd.concat(df_time_list)
    df_mem_all = pd.concat(df_mem_list)

    # average ranks among different checkpoints and/or tree types
    group_cols = ['dataset']

    df_all = df_all.groupby(group_cols).mean().reset_index()
    df_li_all = df_li_all.groupby(group_cols).mean().reset_index()
    df_time_all = df_time_all.groupby(group_cols).mean().reset_index()
    df_mem_all = df_mem_all.groupby(group_cols).mean().reset_index()

    # compute average ranks
    skip_cols = ['dataset', 'tree_type', 'poison_frac']

    df = get_mean_df(df_all, skip_cols=skip_cols, sort='ascending')
    df_li = get_mean_df(df_li_all, skip_cols=skip_cols, sort='ascending')

    logger.info(f'\nLoss:\n{df}')
    logger.info(f'\nLoss (li):\n{df_li}')

    # # combine dataframes
    # index = df_li.index
    # df = df_li.reset_index().merge(df.reset_index(), on='index', how='left')
    # means_df = df[['index', 'mean_x', 'mean_y']].copy()
    # sems_df = df[['index', 'sem_x', 'sem_y']].copy()

    # # rename and clean up
    # means_df.index = means_df['index']
    # sems_df.index = means_df['index']
    # del means_df['index']
    # del sems_df['index']
    # means_df.columns = ['Subgroup A', 'All datasets']
    # sems_df.columns = ['Subgroup A', 'All datasets']

    # print(means_df)
    # print(sems_df)

    # # plot
    # fig, ax = plt.subplots(figsize=(4, 4))
    # means_df.plot.bar(yerr=sems_df, ax=ax, rot=45,
    #                   title=f'Loss ({len(means_df)} datasets)', capsize=3,
    #                   ylabel='Avg. rank', xlabel='Method')

    # plot
    n_datasets = len(df_all['dataset'].unique())
    n_li_datasets = len(df_li_all['dataset'].unique())

    label_dict = {'LeafInfluence': 'LeafInf.', 'SubSample': 'SubS.', 'Target': 'RandomSL'}

    df = df.rename(columns={'mean': 'All datasets'}, index=label_dict)
    df_li = df_li.rename(columns={'mean': 'SDS'}, index=label_dict)

    # reorder methods
    order = ['BoostIn', 'LeafInfSP', 'TreeSim', 'TREX', 'SubS.', 'LOO', 'RandomSL', 'Random']
    order_li = ['LeafRefit', 'LeafInf.', 'BoostIn', 'LeafInfSP', 'TreeSim', 'TREX', 'SubS.', 'LOO',
                'RandomSL', 'Random']

    # order_li = ['BoostIn', 'LeafInfSP', 'TreeSim', 'TREX', 'SubS.', 'LOO', 'RandomSL', 'Random',
    #             'LeafRefit', 'LeafInf.']

    df = df.reindex(order)
    df_li = df_li.reindex(order_li)

    labels = [c if i % 2 != 0 else f'\n{c}' for i, c in enumerate(df.index)]
    labels_li = [c if i % 2 != 0 else f'\n{c}' for i, c in enumerate(df_li.index)]

    pp_util.plot_settings(fontsize=28)
    width = 22
    height = pp_util.get_height(width, subplots=(1, 2))

    fig, axs = plt.subplots(1, 2, figsize=(width, height), gridspec_kw={'width_ratios': [6, 8]})

    ax = axs[0]
    df.plot(kind='bar', y='All datasets', yerr='sem', ax=ax, title=None, capsize=3,
            ylabel='Average rank', xlabel=None, legend=True, color='#3e9ccf')
    ax.set_xticklabels(labels, rotation=0)

    ax = axs[1]
    df_li.plot(kind='bar', y='SDS', yerr='sem', ax=ax, title=None, capsize=3,
               ylabel=None, xlabel=None, legend=True, color='#ff7600')
    ax.set_xticklabels(labels_li, rotation=0)

    ax.axvline(1.5, color='gray', linestyle='--')
    # ax.axvline(7.5, color='gray', linestyle='--')

    logger.info(f'\nSaving results to {out_dir}/...')

    plt.tight_layout()
    plt.savefig(os.path.join(out_dir, 'roar.pdf'), bbox_inches='tight')

    df.to_csv(os.path.join(out_dir, 'loss_rank.csv'))
    df_li.to_csv(os.path.join(out_dir, 'loss_rank_li.csv'))

    logger.info(f'\nTotal time: {time.time() - begin:.3f}s')


def main(args):

    exp_dict = {'n_test': args.n_test, 'poison_frac': args.poison_frac}
    exp_hash = exp_util.dict_to_hash(exp_dict)

    assert len(args.tree_type) > 0
    out_dir = os.path.join(args.in_dir,
                           'rank',
                           f'exp_{exp_hash}',
                           f'+'.join(args.tree_type))

    # create logger
    os.makedirs(out_dir, exist_ok=True)
    logger = exp_util.get_logger(os.path.join(out_dir, 'log.txt'))
    logger.info(args)
    logger.info(f'\ntimestamp: {datetime.now()}')

    process(args, exp_hash, out_dir, logger)


if __name__ == '__main__':
    main(rank_args.get_poison_args().parse_args())

Rank summarization results. # get resource usage # get loss # average ranks among different checkpoints and/or tree types # compute average ranks # # combine dataframes # index = df_li.index # df = df_li.reset_index().merge(df.reset_index(), on='index', how='left') # means_df = df[['index', 'mean_x', 'mean_y']].copy() # sems_df = df[['index', 'sem_x', 'sem_y']].copy() # # rename and clean up # means_df.index = means_df['index'] # sems_df.index = means_df['index'] # del means_df['index'] # del sems_df['index'] # means_df.columns = ['Subgroup A', 'All datasets'] # sems_df.columns = ['Subgroup A', 'All datasets'] # print(means_df) # print(sems_df) # # plot # fig, ax = plt.subplots(figsize=(4, 4)) # means_df.plot.bar(yerr=sems_df, ax=ax, rot=45, #                   title=f'Loss ({len(means_df)} datasets)', capsize=3, #                   ylabel='Avg. rank', xlabel='Method') # plot # reorder methods # order_li = ['BoostIn', 'LeafInfSP', 'TreeSim', 'TREX', 'SubS.', 'LOO', 'RandomSL', 'Random', #             'LeafRefit', 'LeafInf.'] # ax.axvline(7.5, color='gray', linestyle='--') # create logger