#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""lakeshore
author    Benoit Dubois
copyright FEMTO ENGINEERING, 2019
license   GPL v3.0+
brief     Generate "340" calibration file for Lakeshore device.
details   Process Chebyshev polynomial fit on calibration data of cryogenic
          sensors.
"""

import argparse
import numpy as np
import matplotlib.pyplot as plt


HEADER = "Sensor Model:   {}\n\
Serial Number:  {}\n\
Data Format:    4      (Log Ohms/Kelvin)\n\
SetPoint Limit: {:+.3f}      (Kelvin)\n\
Temperature coefficient:  1 (Negative)\n\
Number of Breakpoints:   {:d}\n\
\n\
No.   Units  Temperature (K)\n\
"

FIT_RANGES = (14.2, 80)  # temperature defining limit between ranges
POLY_DEGREE = (9, 6, 7)  # degree of polynom used to fit data in each range
BREAK_NB = (49, 49, 50)  # number of data point in each range


# =============================================================================
def parse_cli():
    """Parse CLI parameters.
    :returns: populated namespace (parser)
    """
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description='Generate "340" calibration file for Lakeshore device.',
        epilog='Example:\n' \
        ' # gen340file file.dat -m CX-1050-BO CX-1050-BO -sn X123456 X234567' \
        ' --range 12.5 80 --degree 9 6 7 --break 49 49 50\n' \
        'Generate calibration files \'X123456.340\' and \'X234567.340\' ' \
        'with fits of degree 9, 6 and 9 and a number of data points ' \
        'of 49, 49 and 50 for, respectively, ' \
        'temperature < 12.5K, temperature between 12.5K and 80K ' \
        'and temperature over 80K, using data from \'file.dat\'.\n\n' \
        'Note: Columns of input file must be formated as: ' \
        '\'Temperatures  Sensor1_values  Sensor2_values...\'')

    parser.add_argument('ifile', type=str,
                        help='Input file with (raw) calibrated data')
    parser.add_argument('-m', '--model', type=str, nargs='+', required=True,
                        help='Model of sensor (ex: CX-1050-BO)')
    parser.add_argument('-sn', '--sn', type=str, nargs='+', required=True,
                        help='Serial number of sensor (ex: X123456)')
    parser.add_argument('--trange', type=float, nargs='+',
                        default=FIT_RANGES,
                        help='Temperature range limit value of fiting process')
    parser.add_argument('--degree', type=int, nargs='+',
                        default=POLY_DEGREE,
                        help='Degree of fit for each range')
    parser.add_argument('--breaks', type=int, nargs='+',
                        default=BREAK_NB,
                        help='Break values (number of data point in each range)')
    parser.add_argument('-g', action='store_true', dest='gdisplay', default=False,
                        help='Enable graphical data display (default=disable)')
    args_ = parser.parse_args()
    return args_


# =============================================================================
def fit_data(data, degree=9, is_log=True):
    """Specific data fit for Lakeshore temperature controller:
    - use 'x' or 'log10(x)' data.
    - normalize 'x' data.
    - fit data with Chebyshev polynomial.
    :param data: data to fit (array)
    :param degree: degree of fit (int)
    :param is_log: use log data or not (bool)
    :returns: Chebyshev polynomial coef and normalized data (array, array)
    """
    x = data[:, 1]  # sensor unit data (voltage, resistance)
    y = data[:, 0]  # temperature data
    if is_log is True:
        Z = np.log10(x)
    else:
        Z = x
    Zl = min(Z)
    Zu = max(Z)
    X = ((Z-Zl)-(Zu-Z))/(Zu-Zl)
    coef = np.polynomial.chebyshev.Chebyshev.fit(X, y, degree)
    if len(np.polynomial.chebyshev.chebroots(coef)) != 0:
        print("Fit on range [{}:{}] is not monotonic".format(min(x), max(x)))
    return coef, Z


# =============================================================================
def check_array_monocity(data):
    """Check monocity of a 1d array.
    :param data: data to check (array)
    :returns: 0 if monotonic else return indexes of monocity breaking (list)
    """
    ddiff = np.diff(data)
    if ddiff[0] > 0:
        ddiff = ddiff[np.where(ddiff < 0)]
        if ddiff is not None:
            return ddiff
    else:
        ddiff = ddiff[np.where(ddiff > 0)]
        if ddiff is not None:
            return ddiff
    return 0


# =============================================================================
def get_data_range(data, min_=None, max_=None):
    """Filter data that are not in the given range [min_:max_].
    'data' is a 2 column array and the filter is applied on the first column.
    :param data: data to process (array)
    :param min_: minimum range value (float)
    :param max_: maximum range value (float)
    :returns: data filtered (array)
    """
    if min_ is None and max_ is None:
        return data
    elif min_ is None:
        return data[np.where(data[:,0] < max_)]
    elif max_ is None:
        return data[np.where(min_ < data[:,0])]
    else:
        data = data[np.where(min_ <= data[:,0])]
        data = data[np.where(data[:,0] <= max_)]
        return data


# =============================================================================
def sort_data(data):
    """Take array of 'x' and 'y' of data (shape: (2,-1)) and sort data
    with respect to the first column.
    :param data: data to sort (array)
    :returns: data sortered (array)
    """
    xys = data[data[:,0].argsort()]
    return xys


# =============================================================================
def main():
    """Script main entry.
    """
    args = parse_cli()
    ifile = args.ifile
    model = args.model
    sn = args.sn
    plot = args.gdisplay
    trange = args.trange
    degree = args.degree
    breaks = args.breaks

    data = np.genfromtxt(ifile).T ##, delimiter='').T
    sensors_nb = np.shape(data)[0]

    for sensor in range(1, sensors_nb):
        xy = np.stack((data[0,:], data[sensor,:]), axis=-1)
        xys = sort_data(xy)
        # Process Chebychev coefficients
        cheb = list()
        Z = list()
        for i in range(len(trange)+1):
            if i == 0:
                xyr = get_data_range(xys, max_=trange[i])
            elif i == len(trange):
                xyr = get_data_range(xys, min_=trange[i-1])
            else:
                xyr = get_data_range(xys,
                                     min_=trange[i-1],
                                     max_=trange[i])
            coef, z = fit_data(xyr, degree[i])
            cheb.insert(i, coef)
            Z.insert(i, z)
        # Generate calibrated data
        for i in range(len(trange)+1):
            x = np.linspace(-1, 1, breaks[i])
            z = np.linspace(min(Z[i]), max(Z[i]), breaks[i])
            if i == 0:
                unit = z
                temp = cheb[i](x)
            else:
                unit = np.concatenate((z, unit))
                temp = np.concatenate((cheb[i](x), temp))
        id_ = np.arange(1, len(unit)+1)
        data_cal = np.stack((id_, unit, temp), axis=-1)
        #
        plt.figure()
        plt.plot(unit, temp, '-')
        #
        np.savetxt("{}.340".format(sn[sensor-1]),
                   data_cal,
                   fmt='%3u %8.5f %13.3f',
                   header=HEADER.format(model[sensor-1],
                                        sn[sensor-1],
                                        int(max(data_cal[:,2])),
                                        len(data_cal)),
                   comments='')

    if plot is True:
        plt.show()


# =============================================================================
main()
