"""
Written by
  Alessandro Motta <alessandro.motta@brain.mpg.de>
"""
import argparse
import logging
import matplotlib.pyplot as plt
import os
import seaborn as sns

from abcsmc.loader import ABCLoader, SQLDataStore

logging.getLogger('matplotlib').setLevel(logging.INFO)
# HACKHACKHACK(amotta): To make text editable
plt.rcParams['pdf.fonttype'] = 42
plt.rcParams['ps.fonttype'] = 42

name = {
    "ER": "ER-ESN",
    "EXP": "EXP-LSM",
    "LL": "LAYERED",
    "SYN": "SYNFIRE",
    "SYH": "SYNFIRE",
    "ERFEVER": "FEVER",
    "API": "API",
    "SORN": "STDP-SORN",
    "adapted": "Adapted",
    # NOTE(amotta): Not sure what these are
    "CDF": "CDF",
    "PDF": "PDF",
    "noise": "Noise",
    "joint": "Shared",
    "model": "Model",
}


def style_cbar_0_1(cbar, label=None):
    if label is not None:
        cbar.set_label(label)
    cbar.set_ticks([0, .2, .4, .6, .8, 1])
    cbar.set_ticklabels([0, "", "", "", "", 1])
    cbar.outline.set_linewidth(0)


def plot_confusion_matrix(output_dir, loader, label_config={}):
    conf_mat_dict = loader.confusion_matrix_dict()
    for index, confusion_matrix in enumerate(conf_mat_dict.values()):
        output_file = os.path.join(output_dir, f"confusion_{index}")

        confmat = confusion_matrix.rename(columns=name, index=name)

        reorder = []
        for (_, n) in name.items():
            # NOTE(amotta): Avoid duplicates
            if not n in reorder: reorder.append(n)
        reorder = [n for n in reorder if n in confmat.columns.to_list()]

        # NOTE(amotta): Use same model order as in all other figures
        confmat = confmat[reorder].reindex(reorder)

        fig, ax = plt.subplots()
        hm = sns.heatmap(confmat,
                         annot=False, annot_kws={'size': 5}, vmin=0, vmax=1,
                         cmap='Greys', ax=ax, square=True)
        hm.set_facecolor('xkcd:salmon')

        cbar = ax.collections[0].colorbar
        style_cbar_0_1(cbar, "$p$")

        plt.setp(ax.xaxis.get_majorticklabels(), **label_config)
        plt.setp(ax.yaxis.get_majorticklabels(), rotation=0)
        ax.set_xlabel("Posterior estimate")
        ax.set_ylabel("Original model")

        with open(output_file + '.txt', "w") as f: f.write(str(confmat))
        plt.savefig(output_file + '.pdf', bbox_inches="tight")


def plot_particle_counts(output_dir, loader):
    results = loader.results_including_initial()
    particles = loader.particle_counts()

    for abc_id, gt_model_name in zip(results.abc_smc_id, results.gt_model_name):
        fig = plt.figure(tight_layout=True)
        sns.heatmap(particles.loc[abc_id, :],
                    vmin=0, vmax=2000, cmap="Greys", square=True,
                    annot=True, fmt='.0f', annot_kws={'size': 10})

        output_file = 'particle-counts_abc-{}_{}.pdf'.format(abc_id, gt_model_name.lower())
        output_file = os.path.join(output_dir, output_file)
        plt.savefig(output_file, bbox_inches='tight')
        plt.close(fig)


def plot_model_probs(output_dir, loader):
    results = loader.results_including_initial()
    model_probs = loader.model_probs()

    for abc_id, gt_model_name in zip(results.abc_smc_id, results.gt_model_name):
        fig = plt.figure(tight_layout=True)
        sns.heatmap(model_probs.loc[abc_id, :],
                    vmin=0, vmax=1, cmap="Greys", square=True,
                    annot=True, fmt='.2f', annot_kws={'size': 10})

        output_file = 'model-probs_abc-{}_{}.pdf'.format(abc_id, gt_model_name.lower())
        output_file = os.path.join(output_dir, output_file)
        plt.savefig(output_file, bbox_inches='tight')
        plt.close(fig)


def make_plots(output_dir, db_file):
    db_name = os.path.split(db_file)[-1]
    db_name = os.path.splitext(db_name)[0]
    output_dir = os.path.join(output_dir, db_name)

    try:
        os.mkdir(output_dir)
    except FileExistsError:
        logging.warning(f"Output directory \"{output_dir}\" already exists")

    loader = ABCLoader(SQLDataStore(db_file))

    conf_mat_config = {'rotation': 75, 'rotation_mode': "anchor", 'horizontalalignment': "right"}
    plot_confusion_matrix(output_dir, loader, conf_mat_config)
    plot_particle_counts(output_dir, loader)
    plot_model_probs(output_dir, loader)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-o', '--output-dir', type=str, default=os.getcwd(),
                        help='Path to output directory in which the plots will be saved. Default: working directory')
    parser.add_argument('input_files', type=str, nargs=argparse.REMAINDER,
                        help='Path to database files')
    args = parser.parse_args()

    for db_file in args.input_files:
        make_plots(args.output_dir, db_file)
