# JMG. Glorian: V1, 5 avril 2021

# The CASSIS team realises that a screenshot of the display of the
# Rotational Diagram module is not the best solution for publishable results. 
# This python script will produce a nice pdf figure that you can use in your articles.
# The font, annotations and titles can be changed to your taste.


import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import ticker, rcParams
import csv
from math import floor, log10

my_path = '/Applications/cassis'
input_dir = os.path.join(my_path, 'delivery/script/data')
output_dir = os.path.join(my_path, 'delivery/script/outputs')

# Input: files saved in the Rotational Diagram display window of CASSIS, through the Save tab
file_saved = os.path.join(input_dir,'Saved_values.csv')  # Saved displayed values
file_fit = os.path.join(input_dir,'Saved_fit')  # Saved fit infos

# Output: location of the pdf file
output_diagram = os.path.join(output_dir,'python_rotational_diagram.pdf')


def fexp(f):
    return int(floor(log10(abs(f)))) if f != 0 else 0


def fman(f):
    return f/10**fexp(f)


def fman2(f, exp_val):
    return f/10**exp_val


def get_rounded_min_max(data, step, minor_step):
    data_min = min(data)
    data_max = max(data)
    rounded_min = np.floor(data_min / step) * step
    rounded_max = np.ceil(data_max / step) * step
    if abs(rounded_min - data_min) < minor_step:
        rounded_min -= step
    if abs(rounded_max - data_max) < minor_step:
        rounded_max += step
    return rounded_min, rounded_max


def get_nticks(axis_min, axis_max, step_major, step_minor):
    nticks_major = (axis_max - axis_min) // step_major + 1
    nticks_minor = np.floor(step_major / step_minor)
    return nticks_major, nticks_minor


def func(x, a, b):
    return a * x + b


def read_rot_data():
    with open(file_saved, newline='') as csvfile:
        reader = csv.reader(csvfile, delimiter='\t')
        xdata = []
        ydata = []
        yerrorData = []
        num_line = 0
        for line in reader:
            if num_line != 0:
                xdata = np.append(xdata, np.double(line[0]))
                ydata = np.append(ydata, np.double(line[1]))
                yerrorData = np.append(yerrorData, np.double(line[2]))
            num_line = num_line + 1

    return xdata, ydata, yerrorData


def read_rot_fit():
    fit_parameters = {}
    with open(file_fit) as file_object:
        for line in file_object:
            if line.strip() != '':
                (key, val) = line.rsplit('=', maxsplit=1)
                fit_parameters[key] = val.strip()
    return fit_parameters


def plot_diagram(xdata, ydata, yerrorData, fit_parameters, Trot, Trot_dev, N_col, error_N_col):
    # Choose font properties
    font = {'family': 'serif',
            'name': 'Times',
            'color': 'black',
            'weight': 'normal',
            'style': 'normal',
            'size': 14,
            }
    # The following updates the font properties for the entire plot
    for k, v in font.items():
        try:
            rcParams.update({'font.'+k: v})
        except KeyError:
            pass  # ignore if error
    rcParams.update({'text.usetex': True})

    fig, ax = plt.subplots()

    # plot the linear fit ; if fit is wanted on top of data points, use a 'large' zorder, e.g. 10
    popt = [float(fit_parameters['Coeff_A']), float(fit_parameters['Coeff_B'])]
    ax.plot(xdata, func(xdata, *popt), 'r-', label='fit: a=%5.3f, b=%5.3f' % tuple(popt), zorder=0)
    # plt.show(block=False)

    # Titles for the axes
    ax.set_ylabel(r'$\ln(N_{\rm u}/g_{\rm u})$', fontsize=16)  # , labelpad=0)
    ax.set_xlabel(r'$E_{\rm u}/k_B$ (K)', fontsize=16)

    # Plot points and error bars
    line, caps, bars = ax.errorbar(
        xdata,  # X
        ydata,  # Y
        yerr=yerrorData,  # Y-errors
        fmt="ko",  # define color, marker and linestyle as in plot ; some possible values :
        # - filled markers : 'o', 'v', '^', '<', '>', '8', 's', 'p', '*', 'h', 'H', 'D', 'd', 'P', 'X'
        # - linestyle : '-', '--', '-.', ':'
        mec='k',  # marker edge color (overrides fmt)
        mfc='k',  # marker face color (overrides fmt)
        ms=5,  # marker size
        linewidth=3,  # width of plot line if linestyle specified in fmt
        elinewidth=0.5,  # width of error bar line
        ecolor='k',  # color of error bar
        capsize=5,  # cap length for error bar
        capthick=0.5  # cap thickness for error bar
    )

    xticks_method = 'user'  # choose between 'auto' or 'user'
    yticks_method = 'user'  # choose between 'auto' or 'user'
    if xticks_method == 'user':
        xstep_major = 50
        xstep_minor = 10
        # choose between one of the following two lines for either
        # (i) limits computer by the program or (ii) your own plot limits
        # xplot_min, xplot_max = get_rounded_min_max(xdata, xstep_major, xstep_minor)
        xplot_min, xplot_max = 50, 455

        ax.set_xlim([xplot_min, xplot_max])
        xticks_major, xticks_minor = get_nticks(xplot_min, xplot_max, xstep_major, xstep_minor)
        ax.xaxis.set_major_locator(plt.MaxNLocator(xticks_major))
        ax.xaxis.set_minor_locator(ticker.AutoMinorLocator(xticks_minor))
    else:  # 'auto' or any other value
        ax.xaxis.set_minor_locator(ticker.AutoMinorLocator())

    if yticks_method == 'user':  # choose step for major and minor yticks
        ystep_major = 1.0
        ystep_minor = 0.2
        # choose between one of the following two lines for either
        # (i) limits computer by the program or (ii) your own plot limits
        # yplot_min, yplot_max = get_rounded_min_max(ydata, ystep_major, ystep_minor)
        yplot_min, yplot_max = 23.5, 28

        ax.set_ylim([yplot_min, yplot_max])
        yticks_major, yticks_minor = get_nticks(yplot_min, yplot_max, ystep_major, ystep_minor)
        ax.yaxis.set_major_locator(plt.MaxNLocator(yticks_major))
        ax.yaxis.set_minor_locator(ticker.AutoMinorLocator(yticks_minor))
    else:  # 'auto' or any other value
        ax.yaxis.set_minor_locator(ticker.AutoMinorLocator())

    # Ticks inside and on all sides
    ax.tick_params(axis='x', which='both', direction='in', top=True)
    ax.tick_params(axis='y', which='both', direction='in', right=True)

    # Annotations
    exp_val = fexp(N_col)
    n_string = r'$N = ({cden} \pm {err})\times 10^{{ {power} }}$ cm$^{{ -2 }}$'.format(
        cden=round(fman(N_col), 1),
        err=round(fman2(error_N_col, exp_val), 1),
        power=exp_val
    )
    t_string = r'$T = ({} \pm {})$ K'.format(round(Trot, 1), round(Trot_dev, 1))
    xmin, xmax = ax.get_xlim()
    ymin, ymax = ax.get_ylim()
    xdiff = xmax - xmin
    ydiff = ymax - ymin
    n_x, n_y = 0.45 * xdiff + xmin, 0.925 * ydiff + ymin
    t_x, t_y = 0.45 * xdiff + xmin, 0.85 * ydiff + ymin
    # can replace automatic placement by un-commenting the following two lines
    # n_x, n_y = 245, 27.5
    # t_x, t_y = 245, 27
    ax.text(n_x, n_y, n_string, fontdict=font)
    ax.text(t_x, t_y, t_string, fontdict=font)

    # Some adjustment in the margins and save pdf file
    plt.subplots_adjust(left=None, bottom=0.12, right=None, top=None, wspace=None, hspace=None)
    plt.show()
    fig.savefig(output_diagram)


if __name__ == '__main__':

    xdata, ydata, yerrorData = read_rot_data()
    fit_parameters = read_rot_fit()

    Trot = float(fit_parameters['Tex'])
    Trot_dev = float(fit_parameters['Tex_error'])
    N_col = float(fit_parameters['Ntot'])
    error_N_col = float(fit_parameters['Ntot_error'])
    print("N = {:.2e} cm-2 +/- {:.2e}".format(N_col, error_N_col))
    print("T = {:.2f} K +/- {:.2f}".format(Trot, Trot_dev))

    plot_diagram(xdata, ydata, yerrorData, fit_parameters, Trot, Trot_dev, N_col, error_N_col)
