"""
WARNING : Do not execute this script in the Jython console of CASSIS.

Python script to plot chi2 contours for all pairs of variables
in the output file produced by a CASSIS regular grid script.
A red cross indicates the position of the minimum chi2.
If at least 10 points within 3-sigma, contour values are expressed as 1-, 2-, 3-sigma,
otherwise they are given as chi2 - min(chi2).
N.B.: Delta chi2 at 1-, 2-, 3-sigma depends on the number of free parameters

Saves a png version of the plot :
- if called from a jython script, same location and name as the input file
- if called directly, define the variables below
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import ticker
from scipy import stats, special
import os
import sys

# User's settings

dirname_out = '/Applications/cassis/delivery/script/outputs/'
dirname_in = '/Applications/cassis/delivery/script/data/'
filename = 'cornerDataRG.dat'
filein = os.path.join(dirname_in, filename)
filename_out = None  # if None, will be taken to be the same as filename, replacing the extension by .png

plot_grid_points = True

plot_chi2_1d = False  # if True : plot chi2 as a function of a single parameter in the diagonal
# in this plot, the grey points are the chi2 values,
# the black points are the chi2 values when all other parameters are set to their best value.

# End user's settings


def read_data(filein):
    # Read the first line and extract reduced chi2min
    with open(filein, 'r') as f:
        elts = f.readline().split()
        chi2min_r = float(elts[1].split(sep='=')[1])
        nfree = int(elts[2].split(sep='=')[1])

    # Read the data in a pandas dataframe
    data = pd.read_csv(filein, sep='\t', header=1)

    # remove first column (number)
    data = data.iloc[:, 1:]

    # remove constant columns if any
    for col_name, col in data.items():
        if len(set(col)) == 1:
            data = data.drop(columns=col_name)

    # find out whether same isotopic ratio in components
    # number of free parameters from cassis less than number of columns minus 1 (chi2 column)
    if nfree < data.shape[1] - 1:
        for col_name, col in data.items():
            if 'iso' in col_name and col_name != 'iso_1':
                data = data.drop(columns=col_name)
        data = data.rename(columns={'iso_1': 'iso'})
    # if 'iso_1' in data.columns.tolist() and 'iso_2' in data.columns.tolist():
    #     if all([v1 == v2 for v1, v2 in zip(data['iso_1'].values, data['iso_2'].values)]):

    return data, chi2min_r, nfree


def get_comp_vars(data_cols):  # currently not used
    """
    Separate variables per component
    :param data_cols: A list of the column names
    :return: A dictionary with key = component number and value = list of variables
    """
    tmp = [vname.split(sep='_')[1][0] for vname in data_cols[:-1]]
    cpt = set(tmp)
    var_dict = {}
    for c in cpt:
        c_vars = [vname for vname in data_cols if '_' + c in vname]
        var_dict[c] = c_vars
    return var_dict


def get_limits(grid_vals):
    """
    Determine extended limits (1/50 of the full range) for a given parameter.
    :param grid_vals: A list of the (ordered) grid values.
    :return: Lower and upper extended limits.
    """
    vmin = grid_vals[0] - np.abs(grid_vals[-1] - grid_vals[0]) / 50
    vmax = grid_vals[-1] + np.abs(grid_vals[-1] - grid_vals[0]) / 50
    return vmin, vmax


def chi2_triangle(data, fileout, chi2min_red=None):
    col_names = data.columns.tolist()[:-1]
    # var_dict = get_comp_vars(col_names)

    chi2min = min(data['chi2'].values)
    print('chi2min = ', '{:.3f}'.format(chi2min))

    i_chi2min = data.index[data['chi2'] == chi2min]
    best_values = {}
    for col in list(data.columns):
        best_values[col] = data[col][i_chi2min].values[0]
    pprint_vals = ' ; '.join([f'{key} = {val:.3f}' for key, val in best_values.items()])
    print("Best values : ", pprint_vals)

    grid_values = {}
    for cname in col_names:
        uniq_vals = list(set(data[cname].values))
        uniq_vals.sort()
        grid_values[cname] = uniq_vals

    ncols_tot = data.shape[1]
    nfree = ncols_tot - 1
    conf_levs = [special.erf(float(n) / np.sqrt(2)) for n in range(1, 4)]  # [0.6827, 0.9545, 0.9973]
    sigmas = [stats.chi2.ppf(cl, nfree) for cl in conf_levs]
    print('Delta chi2 at 1, 2, 3 sigma for {} free parameters = '.format(nfree),
          ', '.join(['{:.3f}'.format(s) for s in sigmas]))

    levs = sigmas

    # sigma5 = stats.chi.ppf(special.erf(5. / np.sqrt(2)), nfree)  # 0.9999994
    # print(str(sigma5))
    # data_sub = data.copy()[data['chi2'] <= (chi2min + sigma5)]  # for histograms
    data_sub = data

    nx = nfree if plot_chi2_1d else nfree - 1  # number of rows and columns
    ny = nx if nx > 1 else 2
    factor = 2.0  # size of one side of one panel
    lbdim = 0.5 * factor  # size of left/bottom margin
    trdim = 0.4 * factor if plot_chi2_1d else 0.3 * factor  # size of top/right margin
    whspace = 0.05  # w/hspace size
    plotdimx = factor * ny + factor * (ny - 1.) * whspace
    plotdimy = factor * nx + factor * (nx - 1.) * whspace
    dimx = lbdim + plotdimx + trdim
    dimy = lbdim + plotdimy + trdim

    fig, axes = plt.subplots(nx, ny, figsize=(dimx, dimy))

    # Format the figure.
    fig.subplots_adjust(left=lbdim / dimx, bottom=lbdim / dimy,
                        right=(lbdim + plotdimx) / dimx, top=(lbdim + plotdimy) / dimy,
                        wspace=whspace, hspace=whspace)

    # formatter = ticker.ScalarFormatter(useMathText=True)
    # formatter.set_scientific(True)

    names_x = col_names if plot_chi2_1d else col_names[:-1]
    names_y = col_names if plot_chi2_1d else col_names[1:]

    for i, icol in enumerate(names_x):
        # plot contours in lower left of the diagonal
        for j, jcol in enumerate(names_y):
            if nx == 1:
                ax = axes[i]
            else:
                ax = axes[j, i]

            if j < i:
                ax.set_frame_on(False)
                ax.set_xticks([])
                ax.set_yticks([])
                continue

            col_names_tmp = col_names.copy()
            col_names_tmp.remove(icol)
            if icol != jcol:
                col_names_tmp.remove(jcol)
            data2plot = data.copy()
            for k in col_names_tmp:
                data2plot = data2plot.loc[data2plot[k] == best_values[k]]

            if data2plot.shape[0] < 3:
                print('Not enough points to draw the contours.')
                continue

            # Limits and labels
            xmin, xmax = get_limits(grid_values[icol])
            ymin, ymax = get_limits(grid_values[jcol])
            xlabel = icol #if 'Log' not in icol else icol[4:-1]
            ylabel = jcol #if 'Log' not in jcol else jcol[4:-1]

            if i == j:
                if plot_chi2_1d:  # plot chi2 "1D" in the diagonal
                    # ymax = 1.1 * get_delta_chi2(nfree, 3)
                    # data2plot = data2plot[(data2plot['Chi2'].values - chi2min) <= ymax]
                    ax.plot(data_sub[icol].values, data_sub['chi2'].values - chi2min, color='0.8', marker='+', ms=1.0,
                            linestyle='')

                    ax.plot(data2plot[icol].values, data2plot['chi2'].values - chi2min, 'ko', ms=1.5)

                    ax.axvline(best_values[icol], color='r', linestyle='--')

                    ax.set_xlim([xmin, xmax])
                    # ax.set_yticklabels([])
                    ax.yaxis.tick_right()
                    ax.yaxis.set_label_position("right")
                    ax.set_ylabel(r'$\Delta\chi^2$')
                    if i < nx - 1:
                        ax.set_xticklabels([])
                    else:
                        [l.set_rotation(45) for l in ax.get_xticklabels()]
                        ax.set_xlabel(xlabel)
                        ax.xaxis.set_label_coords(0.5, -0.3)

                    continue

            x = data2plot[icol].values
            y = data2plot[jcol].values
            z = data2plot['chi2'].values - chi2min

            if len([zi for zi in z if (levs[0] <= zi <= levs[-1])]) >= 5:  # at least 6 points within 1-3 sigma
                ctr = ax.tricontour(x, y, z, linewidths=1.5,
                                    levels=levs,
                                    colors='k'
                                    )
                strs = [r'$1\sigma$', r'$2\sigma$', r'$3\sigma$']
                fmt = {l: s for l, s in zip(ctr.levels, strs)}
                ax.clabel(ctr, ctr.levels, inline=True, fmt=fmt, fontsize=10)
            else:
                ctr_default = ax.tricontour(x, y, z, 3, linewidths=1.5, colors='k')
                ax.clabel(ctr_default, ctr_default.levels, inline=True, fontsize=10)

            if plot_grid_points:
                ax.plot(x, y, color='0.8', marker='o', ms=2.0, linestyle='', label='Grid locations')
                handles, labels = ax.get_legend_handles_labels()

            # Plot point for chi2min
            ax.plot(best_values[icol], best_values[jcol], 'r+')

            # Set limits and labels
            ax.set_xlim([xmin, xmax])
            ax.set_ylim([ymin, ymax])

            if j < ny - 1 and nx != 1:
                ax.set_xticklabels([])
            else:
                # if 'Log' in icol:
                #     ax.set_xticklabels(['{:.2e}'.format(10.**l) for l in ax.get_xticks()])
                    # ax.xaxis.set_major_formatter(formatter)
                [l.set_rotation(45) for l in ax.get_xticklabels()]
                ax.set_xlabel(xlabel)
                ax.xaxis.set_label_coords(0.5, -0.3)
            if i > 0:
                ax.set_yticklabels([])
            else:
                # if 'Log' in jcol:
                #     # ax.set_yticklabels(['{:.2e}'.format(10.**l) for l in ax.get_yticks()])
                #     ax.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, pos: str(10.**x)))
                ax.set_ylabel(ylabel)

    if nx == 1:
        ax = axes[1]
        ax.set_frame_on(False)
        ax.set_xticks([])
        ax.set_yticks([])

    # print some info in upper right corner
    chi2min_str = r'$\chi^2_{\rm min}$'
    text = r'{} = {:.3f}'.format(chi2min_str, chi2min) + '\n'
    if chi2min_red:
        text += r'{}, reduced = {:.3f}'.format(chi2min_str, chi2min_red) + '\n'
    text += r'Number of free parameters : {}'.format(nfree) + '\n'
    # text += r'$\Delta\chi^2$ at 1, 2, 3$\sigma$ : {}'.format(', '.join(['{:.3f}'.format(s) for s in sigmas]))
    for i, sig in enumerate(sigmas):
        text += r'$\Delta\chi^2$ at {}$\sigma$ : {:.3f}'.format(i + 1, sig) + '\n'

    ax = axes[1] if nx == 1 else axes[0, ny - 1]
    xpos = 0.3 if plot_chi2_1d else 0
    ax.annotate(text, xy=(xpos, 1), xycoords='axes fraction', va='top')

    if plot_grid_points:
        ax.legend(handles=handles, labels=labels, loc='lower center', frameon=False)

    fig.savefig(fileout)
    plt.show(block=True)


if __name__ == '__main__':
    if len(sys.argv) >= 2:  # if called from script
        filein = sys.argv[1]
        dirname_out = None
        filename_out = None

    if os.path.isfile(filein):
        print('Data file : ', filein)
        data, chi2min_r, nfree = read_data(filein)

        output_dir = os.path.dirname(filein) if dirname_out is None else dirname_out
        if not os.path.isdir(output_dir):
            os.makedirs(output_dir)
        if filename_out is None:
            filename_out = os.path.splitext(os.path.split(filein)[-1])[0] + '.png'
        fileout = os.path.join(output_dir, filename_out)
        chi2_triangle(data, fileout, chi2min_red=chi2min_r)
    else:
        print('File ' + filein + ' not found.')
