"""
Written by
  Alessandro Motta <alessandro.motta@brain.mpg.de>
"""
import argparse
import logging
import matplotlib.pyplot as plt
import numpy as np
import os
import pandas as pd
import re
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, 1])
    cbar.set_ticklabels([0, 1])
    cbar.outline.set_linewidth(0)


def plot_matrix(output_dir, data, label_config={}):
    data.sort(key=lambda datum: datum[0], reverse=True)

    side_lens, accuracies = zip(*data)
    accuracies = np.asarray(accuracies).reshape((1, -1))
    data = pd.DataFrame(data=accuracies, columns=side_lens)

    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, "Average\nprobability")

    plt.setp(ax.xaxis.get_majorticklabels(), **label_config)
    plt.setp(ax.yaxis.get_majorticklabels(), rotation=0)
    ax.set_xlabel("EM volume (µm)")

    output_file = os.path.join(output_dir, "barrelsubvolume_average_accuracy")
    plt.savefig(output_file + '.pdf', bbox_inches="tight", transparent=True)


def load_data(db_file):
    # HACKHACKHACK(amotta): Extract cube length from file name
    side_len = int(re.search(r'_(\d+)um_', db_file).group(1))

    loader = ABCLoader(SQLDataStore(db_file))
    accuracy = loader.means().mass_at_gt_model
    return (side_len, accuracy)


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()

    config = {'rotation': 75, 'rotation_mode': "anchor", 'horizontalalignment': "right"}
    data = list(map(load_data, args.input_files))
    plot_matrix(args.output_dir, data, config)
