from __future__ import annotations

__version__ = '0.1.2'

import os
from argparse import ArgumentParser
from pathlib import Path
import sys
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import simps

from nmreval.io.dsc import DSCReader, Cyclohexane, ReferenceValue

parser = ArgumentParser(description='Calibrate DSC data')
parser.add_argument('sample', type=str, help='filename of DSC sample')
parser.add_argument('empty', type=str, help='filename of empty pan')
parser.add_argument('reference', help='filename of reference', type=str)
parser.add_argument('--cooling', help='Plot found cooling rates', action='store_true')


def evaluate(sample: str|Path, empty: str|Path, reference: str|Path,
             ref_points: ReferenceValue = Cyclohexane, show_cooling: bool = False):

    sample = DSCReader(sample)
    empty = DSCReader(empty)
    reference = DSCReader(reference)
    print(sample)

    if show_cooling:
        fig, ax = plt.subplots()
        print('\n')
        for k, v in sample.cooling.items():
            print('Plot run {} with cooling rate {} K/min'.format(k, v))
            c = sample.flow_data(v, mode='c')
            ax.plot(c[0], c[1], label=str(v)+' K/min')
        ax.set_xlabel('T / K')
        plt.legend()
        plt.show()

        return

    run_list = []
    if len(sample.heating) > 1:
        run = None
        print('\nMultiple heat rates found:')
        for k, v in sample.heating.items():
            print('  run {}: {} K/min'.format(k, v))
        while run not in sample.heating:
            # choose your own adventure!!!
            value = input('\nPlease select a run (press Enter for all heat rates): ')
            if value == '':
                run_list = list(sample.heating.keys())
                break
            else:
                run = int(value)
                run_list = [run]
    else:
        run_list = list(sample.heating.keys())

    for run in run_list:
        rate = sample.heating[run]

        print('\nProcessing heat rate {} K/min'.format(rate))

        print('Load data of heating data')
        len_sample = sample.length(run)

        # sanity checks
        try:
            reference_data = reference.flow_data(rate)
        except IndexError:
            print('ERROR: Reference measurement has no heat rate {} K/min'.format(rate))
            print('Stop evaluation')
            sys.exit()

        try:
            run_baseline = empty.get_run(rate)
        except ValueError:
            print('ERROR: Empty measurement has no heat rate {} K/min'.format(rate))
            print('Stop evaluation')
            sys.exit()

        len_baseline = empty.length(run_baseline)
        max_length = None
        if len_baseline != len_sample:
            print('WARNING: measurements differ by {} points'.format(abs(len_baseline - len_sample)))
            max_length = min(len_baseline, len_sample)

        sample_data = sample.flow_data(rate, length=max_length)
        empty_data = empty.flow_data(rate, length=max_length)

        # plot input data
        fig1, ax1 = plt.subplots(2, 3, **{'figsize': (10, 6)})
        ax1[0, 0].set_title('raw data')
        ax1[0, 0].set_xlabel('T / K')

        ax1[0, 0].plot(sample_data[0], sample_data[1], 'k-', label='Sample')
        ax1[0, 0].plot(empty_data[0], empty_data[1], 'b-', label='Empty')
        ax1[0, 0].plot(reference_data[0], reference_data[1], 'r-', label='Reference')
        ax1[0, 0].legend()

        print('Substract empty data')
        sample_baseline = sample_data.copy()
        sample_baseline[1] = sample_data[1] - empty_data[1]

        # plot baseline correction
        ax1[0, 1].set_title('baseline correction')
        ax1[0, 1].set_xlabel('T / K')

        ax1[0, 1].plot(sample_data[0], sample_data[1], 'k--', label='Raw')
        ax1[0, 1].plot(sample_baseline[0], sample_baseline[1], 'k-', label='Baseline corrected')
        ax1[0, 1].plot(empty_data[0], empty_data[1], 'b-', label='Empty')
        ax1[0, 1].legend()

        print('Load isothermal data around heat rate')
        mean_isotherms = []
        for offset, where, ls in [(-1, 'low', '-'), (1, 'high', '--')]:
            # read isotherms and baseline correct
            len_baseline = empty.length(run_baseline+offset)
            len_sample = sample.length(run+offset)
            if len_baseline != len_sample:
                print('WARNING: {} T isotherms differ by {} points'.format(where, abs(len_baseline-len_sample)))

            max_length = min(len_baseline, len_sample)
            isotherm_sample = sample.isotherm_data(run_baseline+offset, length=max_length)
            isotherm_empty = empty.isotherm_data(run+offset, length=max_length)

            isotherm_sample[1] -= isotherm_empty[1]

            # get mean isotherm value
            m = np.polyfit(isotherm_sample[0, 200:-200], isotherm_sample[1, 200:-200], 0)[0]
            mean_isotherms.append(m)
            print('Calculated {} heat flow: {} mW'.format(where, m))

            ax1[0, 2].plot(isotherm_sample[0], isotherm_sample[1], 'k--')

        # calculate slope from difference between isotherms
        slope = (mean_isotherms[1]-mean_isotherms[0]) / (sample_data[2, -1] - empty_data[2, 0])
        print('Heat flow slope from isotherms: {} per minute'.format(slope*60))

        # calculate mean slope of heat flow at points in the beginning
        slope_baseline = np.gradient(sample_baseline[1, int(4000/rate):int(9000/rate)],
                                     sample_baseline[2, 300]-sample_baseline[2, 299]).mean()
        print('Heat flow slope from initial heating: {} per minute'.format(slope_baseline*60))

        drift_corrected = sample_baseline[1] - mean_isotherms[0] - (sample_baseline[2]-empty_data[2, 0])*slope
        drift_from_slope = sample_baseline[1] - mean_isotherms[0] - (sample_baseline[2]-empty_data[2, 0])*slope_baseline

        # plot
        ax1[0, 2].axhline(mean_isotherms[0], linestyle=':')
        ax1[0, 2].axhline(mean_isotherms[1], linestyle=':')

        ax1[0, 2].plot(sample_baseline[2], sample_baseline[1], 'k-', label='Baseline corrected')
        ax1[0, 2].plot(sample_baseline[2], drift_corrected, 'g-', label='Corrected (isotherm)')
        ax1[0, 2].plot(sample_baseline[2], drift_from_slope, 'b-', label='Corrected (heating)')

        ax1[0, 2].plot(sample_baseline[2], mean_isotherms[0] + (sample_baseline[2]-empty_data[2, 0])*slope, 'g--')
        ax1[0, 2].plot(sample_baseline[2], mean_isotherms[0] + slope_baseline*(sample_baseline[2]-empty_data[2, 0]),
                       'b--')

        ax1[0, 2].plot(sample_baseline[2, int(4000/rate):int(9000/rate)],
                       sample_baseline[1, int(4000/rate):int(9000/rate)], 'r--')

        ax1[0, 2].set_title('time dependence')
        ax1[0, 2].set_xlabel('t / s')
        ax1[0, 2].legend()

        melts = []
        for i, (trans_temp, enthalpy) in enumerate(ref_points.transitions):
            print(trans_temp, enthalpy)

            # region around reference peaks
            # NOTE: limits are hard coded for cyclohexane, other references need other limits
            low_border = np.argmin(abs(reference_data[0]-(trans_temp-15)))
            high_border = np.argmin(abs(reference_data[0]-(trans_temp+15)))
            ref_zoom = reference_data[:, low_border:high_border]
            x_val = np.array([[ref_zoom[0, 0], 1],
                              [ref_zoom[0, -1], 1]])
            y_val = np.array([ref_zoom[1, 0],
                              ref_zoom[1, -1]])
            print('Baseline correct reference of {} transition'.format(trans_temp))
            sol = np.linalg.solve(x_val, y_val)
            ref_zoom[1] -= (ref_zoom[0] * sol[0] + sol[1])
            peak_max = ref_zoom[:, np.argmax(ref_zoom[1])]
            integration_limit = np.argmin(abs(ref_zoom[0]-peak_max[0]+3)), np.argmin(abs(ref_zoom[0]-peak_max[0]-3))

            # substract baseline around reference peaks
            x_val = np.array([[ref_zoom[0, integration_limit[0]], 1],
                              [ref_zoom[0, integration_limit[1]], 1]])
            y_val = np.array([ref_zoom[1, integration_limit[0]],
                              ref_zoom[1, integration_limit[1]]])

            print('Baseline correct reference of {} transition'.format(trans_temp))
            sol = np.linalg.solve(x_val, y_val)
            ref_zoom[1] -= (ref_zoom[0] * sol[0] + sol[1])

            # calculate onset slope (use points at position of maximum gradient +/- 100/rate)
            ref_grad = np.gradient(ref_zoom[1])
            max_grad = np.argmax(ref_grad)

            x_val = np.array([[ref_zoom[0, max_grad-int(100/rate)], 1],
                              [ref_zoom[0, max_grad+int(100/rate)+1], 1]])
            y_val = np.array([ref_zoom[1, max_grad-int(100/rate)],
                              ref_zoom[1, max_grad+int(100/rate)+1]])
            sol = np.linalg.solve(x_val, y_val)
            onset = sol[0]*ref_zoom[0] + sol[1]

            melts.append(-sol[1]/sol[0])

            # plot
            ax1[1, i].set_title('reference: {:.2f} K'.format(trans_temp))
            ax1[1, i].set_xlabel('T / K')

            ax1[1, i].plot(reference_data[0], reference_data[1], 'r-')
            ax1[1, i].plot(ref_zoom[0, max_grad], ref_zoom[0, max_grad], 'kx')
            ax1[1, i].plot(ref_zoom[0], onset, 'k--')
            ax1[1, i].axhline(0, color='k', linestyle='--')

            ax1[1, i].set_xlim(ref_zoom[0, integration_limit[0]], ref_zoom[0, integration_limit[1]])
            ax1[1, i].set_ylim(-max(ref_zoom[1])/10, max(ref_zoom[1])*1.1)

            print('Onset of transition: {:.2f} K, should be at {:.2f}'.format(melts[-1], trans_temp))
            if enthalpy is not None:
                # integrate over low temperature peak to calibrate y axis
                # NOTE: again, this is only valid for cyclohexane
                # delta H in J/g: Integrate Peak over time and divide by weight
                area = 1e-3 * simps(ref_zoom[1, integration_limit[0]:integration_limit[1]],
                                    ref_zoom[2, integration_limit[0]:integration_limit[1]],
                                    even='avg')
                calib_y_axis = enthalpy / (area / reference.weight)
                print("Calibration factor of peak: {}".format(calib_y_axis))

                sample_baseline[1] *= calib_y_axis

        fig1.delaxes(ax1[1, 2])
        fig1.tight_layout()

        plt.show()

        # give a choice how to compensate for long-time drift
        mode = None
        while mode not in ['i', 'h']:
            mode = input('\nUse [i]sotherms or initial [h]eating for long-time correction? (Default: i) ')
            if mode == '':
                mode = 'i'

        if mode == 'h':
            print('\nCorrect slope from initial heating')
            sample_baseline[1] = drift_from_slope
        else:
            print('\nCorrect slope from isotherm')
            sample_baseline[1] = drift_corrected

        # calibrate T axis
        print('\nCalibrate temperature')
        real_trans = np.array([ref_points.transition1, ref_points.transition2])
        t_vals = np.array([[melts[0], 1],
                           [melts[1], 1]])
        calibration_temp = np.linalg.solve(t_vals, real_trans)
        print('T_real = {:.4f} * T_meas {:+.4f}'.format(*calibration_temp))

        sample_baseline[0] = calibration_temp[0] * sample_baseline[0] + calibration_temp[1]

        print('Convert to capacity')
        cp = sample_baseline[1] * 60. / rate / sample.weight / 1000.
        if sample.weight is None:
            raise ValueError('No sample weight given')

        # plot final results in separate figure
        fig2, ax2 = plt.subplots()
        ax2.set_title('{} K/min: Heat flow vs. heat capacity (close to cont.)'.format(rate))
        ax2.set_xlabel('Temperature / K')

        ax2.plot(sample_baseline[0], sample_baseline[1], label='heat flow')
        ax2.plot(sample_baseline[0], cp, label='heat capacity')

        plt.legend()
        plt.show()

        outname = os.path.splitext(sample.name)[0] + '_' + str(rate) + 'K-min.dat'
        header = 'Made with version: {}\n'.format(__version__)
        header += 'T/K\tCp/J/(gK)\theat flow/mW'

        print()
        print('Save to', outname)
        np.savetxt(outname, np.c_[sample_baseline[0], cp, sample_baseline[1]], header=header)


if __name__ == '__main__':
    args = parser.parse_args()
    evaluate(args.sample, args.empty, args.reference, show_cooling=args.cooling)
