"""
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


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_data(output_file, data, label_config={}):
    fig, ax = plt.subplots()
    hm = sns.heatmap(data,
                     annot=True, annot_kws={'size': 5}, vmin=0, vmax=1,
                     cmap='Greys', ax=ax, square=True)
    cbar = ax.collections[0].colorbar
    style_cbar_0_1(cbar, "$p$")

    plt.setp(ax.xaxis.get_majorticklabels())
    plt.setp(ax.yaxis.get_majorticklabels(), rotation=0)
    ax.set_xlabel("Noise")
    ax.set_ylabel("Subsampling")

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


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))
    loader.group_parameters_from_pars = ["noise"]
    loader.group_parameters_from_noise = ["subsampling_fraction"]

    plot_config = {'rotation': 75, 'rotation_mode': "anchor", 'horizontalalignment': "right"}

    average_accuracy = loader.average_mass_at_tround_truth()
    plot_data(os.path.join(output_dir, 'average_accuracy'), average_accuracy, plot_config)

    map_accuracy = loader.maximum_a_posteriori()
    plot_data(os.path.join(output_dir, 'map_accuracy'), map_accuracy, plot_config)


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)
