#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
WARNING : Do not execute this script in the Jython console of CASSIS.

Python script to create a corner plot from the output of the MCMC algorithm in CASSIS.

@author: Emmanuel Caux & Sandrine Bottinelli, IRAP, Toulouse, France
"""
import sys
import numpy as np
from scipy.special import erf
import matplotlib.pyplot as pl
import os
from os.path import dirname
from os.path import abspath
import corner as corner

pl.rcParams["font.family"] = "DejaVu Sans"
pl.rcParams["xtick.direction"] = "in"
pl.rcParams["ytick.direction"] = "in"
pl.rcParams["font.weight"] = "normal"
# pl.rcParams["font.stretch"] = "extra-condensed"  # matplotlib doc : property currently not implemented
pl.rcParams["axes.formatter.limits"] = [-2, 3]  # numbers <0.01 or >1000. will be displayed in scientific format


if __name__ == '__main__':

    # USER INPUTS
    n_sig = 1  # vertical lines to plot at +/- nsigma for nsigma in [1, ..., n_sig]

    fraction_walkers_rejected = 0.1  # default fraction of walkers to be rejected

    input_dir = '/Applications/cassis/delivery/script/data'
    output_dir = '/Applications/cassis/delivery/script/outputs'
    if not os.path.isdir(output_dir):
        os.makedirs(output_dir)

    # filename = "W3OH_AE-CH3OH_2C_cont_lte"
    # filename = "OVI"
    # filename = "Full_GRID-Statistics2"
    filename = "cornerData"
    inputFilename = os.path.join(input_dir, filename + '.dat')
    outputFilename = os.path.join(output_dir, filename + '_corner')
    # END USER INPUTS

    # print("myScript = ", sys.argv[0])
    if len(sys.argv) > 2:  # call from jython script
        myPath = sys.argv[1]
        # print("myPath = ", myPath)
        myName = sys.argv[2]
        # print("myName = ", myName)
        if len(sys.argv) == 4:
            fraction_walkers_rejected = float(sys.argv[3])
        inputFilename = os.path.join(myPath, myName + '.dat')
        outputFilename = os.path.join(myPath, myName + '_corner')  # component id will be appended

    # file structure :
    # l.1 : component type (lte or radex)
    # l.2 : column names (e.g. #  nmol_1  tex_1  nmol_2  tex_2) ; always starts with # (the walker number)
    # l.3+: parameter values for each walker

    # read the data and make sure new lines are specified with \n, not \r ; TODO : is this necessary?
    file_contents = open(inputFilename).read().replace("\r", "\n").splitlines()

    # read the data and remove leading underscore from names
    myData = np.genfromtxt(inputFilename, comments='!', skip_header=1, delimiter='\t', names=True)
    new_dtype = [(name.strip(), myData.dtype[name]) for name in myData.dtype.names]
    new_data = myData.astype(new_dtype)
    # myData.dtype.names = (name.lstrip('_') for name in myData.dtype.names)
    # get the number of data lines
    nb_elements = len(myData)

    # remove unwanted data
    linesSkipped = int(fraction_walkers_rejected * nb_elements)
    myData = myData[linesSkipped:]

    #  walker, chi2, rate are not parameters we are interested in, so remove the corresponding columns:
    unwanted = ('f0', 'Chi2', 'chi2', 'rate', 'Nb', 'Walker', 'nb', 'walker', 'StepNb')
    select = [x for x in myData.dtype.names if x not in unwanted]
    myData = myData[select]

    # determine the number of components and their "names" (1, 2, etc)
    par_list = myData.dtype.names  # list of parameters

    par_comp_dict = {}
    for par in par_list:
        comp = par.split('_')[-1]
        if comp.isnumeric():
            if comp in par_comp_dict:
                par_comp_dict[comp].append(par)
            else:
                par_comp_dict[comp] = [par]

    pc = [erf(sig / np.sqrt(2)) for sig in range(1, n_sig + 1, 1)]  # 0.6827, 0.9545, 0.9973 for sigma = 1, 2, 3 resp.
    quantiles_m = [0.5 - p / 2 for p in pc]
    quantiles_p = [0.5 + p / 2 for p in pc]
    quantiles_plot = quantiles_m + quantiles_p

    powerlimits = pl.rcParams["axes.formatter.limits"]

    fig = []

    for cpt, par_list_comp in par_comp_dict.items():
        physicalPar = [v for v in par_list_comp if len(set(myData[v])) > 1]  # list of varying parameters
        if not physicalPar:
            print('No plot for component {} : it contains only fixed parameters.'.format(cpt))
            continue
        data_comp_plot = myData[physicalPar]
        m_true = np.array([np.median(data_comp_plot[v]) for v in physicalPar])  # array of median values
        s_true = np.array([np.std(data_comp_plot[v]) for v in physicalPar])  # array of standard dev values

        # define rangePar = list of length 2 tuples providing plotting range (+/- 4 sigma)
        rangePar = [(np.percentile(data_comp_plot[v], (0.5 - erf(4. / np.sqrt(2)) / 2) * 100.),
                     np.percentile(data_comp_plot[v], (0.5 + erf(4. / np.sqrt(2)) / 2) * 100.)) for v in physicalPar]

        samples = np.array(data_comp_plot.tolist())
        ndim = len(physicalPar)

        # call corner without titles
        corner_fig = corner.corner(samples, labels=physicalPar, truths=m_true,
                                   quantiles=quantiles_plot,
                                   # contour levels at 1, 2, 3 sigmas :
                                   levels=1.0 - np.exp(-0.5 * np.arange(1, 3.1, 1) ** 2),
                                   label_kwargs={"fontsize": 11}, range=rangePar)

        # add titles formatted as desired depending on whether log(value) is within power limits.
        axes = np.array(corner_fig.axes).reshape((ndim, ndim))
        for i in range(ndim):
            ax = axes[i, i]
            par = i
            q_50 = m_true[par]
            q_m = [q_50 - np.percentile(data_comp_plot[physicalPar[par]], q * 100.) for q in quantiles_m]
            q_p = [np.percentile(data_comp_plot[physicalPar[par]], q * 100.) - q_50 for q in quantiles_p]
            power = int(np.log10(np.abs(q_50)))

            # Format the quantile display.
            if (power >= powerlimits[0]) and (power < powerlimits[1]):
                title_fmt = ".2f"
            else:
                title_fmt = ".2e"

            if title_fmt == ".2e":
                fmt = "{{0:{0}}}".format(".2f").format
                title = r"$({{{0}}}_{{-{1}}}^{{+{2}}})\times10^{{{3}}}$"
                title = title.format(fmt(q_50 / np.power(10., power)),
                                     ', '.join([fmt(q / np.power(10., power)) for q in q_m]),
                                     ', '.join([fmt(q / np.power(10., power)) for q in q_p]),
                                     power)
            else:
                fmt = "{{0:{0}}}".format(title_fmt).format
                title = r"${{{0}}}_{{-{1}}}^{{+{2}}}$"
                title = title.format(fmt(q_50),
                                     ', '.join([fmt(q) for q in q_m]),
                                     ', '.join([fmt(q) for q in q_p])
                                     )

            title = "{0} = {1}".format(physicalPar[par], title)

            ax.set_title(title)

        corner_fig.savefig(outputFilename+'_'+str(cpt)+'.png')
        fig.append(corner_fig)

    pl.show(block=True)
