# -*- coding: utf-8 -*-
""" Toolbox for XAS experiments.

    Based on the LCLS LO59 experiment libraries.

    Time-resolved XAS and XMCD with uncertainties.

    Copyright (2019-) SCS Team
    Copyright (2017-2019) Loïc Le Guyader <loic.le.guyader@xfel.eu>
"""

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

def absorption(T, Io):
    """ Compute the absorption A = -ln(T/Io)
        Inputs:
            T: 1-D transmission value array of length N
            Io: 1-D Io monitor value array of length N

        Output:
            a structured array with:
                muA: absorption mean
                sigmaA: absorption standard deviation
                weights: sum of Io values
                muT: transmission mean
                sigmaT: transmission standard deviation
                muIo: Io mean
                sigmaIo: Io standard deviation
                p: correlation coefficient between T and Io
                counts: length of T
    """

    counts = len(T)
    assert counts == len(Io), "T and Io must have the same length"

    # return type of the structured array
    fdtype = [('muA', 'f8'), ('sigmaA', 'f8'), ('weights', 'f8'),
        ('muT', 'f8'), ('sigmaT', 'f8'), ('muIo', 'f8'), ('sigmaIo', 'f8'),
        ('p', 'f8'), ('counts', 'i8')]
    if counts == 0:
        return np.array([(np.NaN, np.NaN, np.NaN, np.NaN, np.NaN, np.NaN,
            np.NaN, np.NaN, 0)], dtype=fdtype)

    muT = np.mean(T)
    sigmaT = np.std(T)

    muIo = np.mean(Io)
    sigmaIo = np.std(Io)
    weights = np.sum(Io)

    p = np.corrcoef(T, Io)[0,1]

    muA = -np.log(muT/muIo)
    # from error propagation for correlated data
    sigmaA = (np.sqrt((sigmaT/muT)**2 + (sigmaIo/muIo)**2
        - 2*p*sigmaIo*sigmaT/(muIo*muT)))
    
    # direct calculation
    #mask = (Io != 0)
    #sigmaA = np.nanstd(-np.log(T[mask]/Io[mask]))

    return np.array([(muA, sigmaA, weights, muT, sigmaT, muIo, sigmaIo,
        p, counts)], dtype=fdtype)

def binning(x, data, func, bins=100, bin_length=None):
    """ General purpose 1-dimension data binning

        Inputs:
            x: input vector of len N
            data: structured array of len N
            func: a function handle that takes data from a bin an return
                a structured array of statistics computed in that bin
            bins: array of bin-edges or number of desired bins
            bin_length: if not None, the bin width covering the whole range

        Outputs:
            bins: the bins edges
            res: a structured array of binned data
    """

    if bin_length is not None:
        bin_start = np.amin(x)
        bin_end = np.amax(x)
        bins = np.arange(bin_start, bin_end+bin_length, bin_length)
    elif np.size(bins) == 1:
        bin_start = np.amin(x)
        bin_end = np.amax(x)
        bins = np.linspace(bin_start, bin_end, bins)
    bin_centers = (bins[1:]+bins[:-1])/2
    nb_bins = len(bin_centers)

    bin_idx = np.digitize(x, bins)
    dummy = func([])
    res = np.empty((nb_bins), dtype=dummy.dtype)
    for k in range(nb_bins):
        res[k] = func(data[bin_idx == k])

    return bins, res

def xas(nrun, bins=None, Iokey='SCS_SA3', Itkey='MCP3apd', nrjkey='nrj', Iooffset=0, plot=False):
    """ Compute the XAS spectra from a xarray nrun.

        Inputs:
            nrun: xarray of SCS data
            bins: an array of bin-edges or an integer number of
                desired bins or a float for the desired bin width.
            Iokey: string for the Io fields, typically 'SCS_XGM'
            Itkey: string for the It fields, typically 'MCP3apd'
            NRJkey: string for the nrj fields, typically 'nrj'
            Iooffset: offset to apply on Io
            plot: boolean, displays a XAS spectrum if True

        Outputs:
            a dictionnary containing:
                nrj: the bins center
                muA: the absorption
                sigmaA: standard deviation on the absorption
                sterrA: standard error on the absorption
                muIo: the mean of the Io
                counts: the number of events in each bin
    """
    
    stacked = nrun.stack(dummy_=['trainId','sa3_pId'])

    Io = stacked[Iokey].values + Iooffset
    It = stacked[Itkey].values
    nrj = stacked[nrjkey].values
    
    names_list = ['nrj', 'Io', 'It']
    rundata = np.vstack((nrj, Io, It))
    rundata = np.rec.fromarrays(rundata, names=names_list)
    
    def whichIo(data):
        """ Select which fields to use as I0 and which to use as I1
        """
        
        if 'mcp' in Iokey.lower():
            Io_sign = -1
        else:
            Io_sign = 1

        if 'mcp' in Itkey.lower():
            It_sign = -1
        else:
            It_sign = 1

        
        if len(data) == 0:
            return absorption([], [])
        else:
            return absorption(It_sign*data['It'], Io_sign*data['Io'])

    if bins is None:
        num_bins = 80
        energy_limits = [np.min(nrj), np.max(nrj)]
        bins = np.linspace(energy_limits[0], energy_limits[1], num_bins+1)
    elif type(bins) == int:
        energy_limits = [np.min(nrj), np.max(nrj)]
        bins = np.linspace(energy_limits[0], energy_limits[1], bins+1)
    elif type(bins) == float:
        energy_limits = [np.min(nrj), np.max(nrj)]
        bins = np.arange(energy_limits[0], energy_limits[1], bins)
        
    dummy, nosample = binning(rundata['nrj'], rundata, whichIo, bins)
    muA = nosample['muA']
    sterrA = nosample['sigmaA']/np.sqrt(nosample['counts'])

    bins_c = 0.5*(bins[1:] + bins[:-1])
    
    if plot:
        f = plt.figure(figsize=(6.5,6))
        gs = gridspec.GridSpec(2,1,height_ratios=[4,1])
        ax1 = plt.subplot(gs[0])
        ax1.plot(bins_c, muA, color='C1', label=r'$\sigma$')
        ax1.set_ylabel('XAS')
        ax1.set_xlabel('Energy (eV)')
        ax1.legend()
        ax1_twin = ax1.twinx()
        ax1_twin.bar(bins_c, nosample['muIo'], width=0.80*(bins_c[1]-bins_c[0]),
                color='C1', alpha=0.2)
        ax1_twin.set_ylabel('Io')
        try:
            proposalNB=int(nrun.attrs['runFolder'].split('/')[-4][1:])
            runNB=int(nrun.attrs['runFolder'].split('/')[-2][1:])
            ax1.set_title('run {:d} p{:}'.format(runNB, proposalNB))
        except:
            f.suptitle(nrun.attrs['plot_title'])
        
        ax2 = plt.subplot(gs[1])
        ax2.bar(bins_c, nosample['counts'], width=0.80*(bins_c[1]-bins_c[0]),
                color='C0', alpha=0.2)
        ax2.set_xlabel('Energy (eV)')
        ax2.set_ylabel('counts')
        plt.tight_layout()
    
    return {'nrj':bins_c, 'muA':muA, 'sterrA':sterrA, 'sigmaA':nosample['sigmaA'],
        'muIo':nosample['muIo'], 'counts':nosample['counts']}

def xasxmcd(dataP, dataN):
    """ Compute XAS and XMCD from data with both magnetic field direction
        Inputs:
            dataP: structured array for positive field
            dataN: structured array for negative field

        Outputs:
            xas: structured array for the sum
            xmcd: structured array for the difference
    """

    assert len(dataP) == len(dataN), "binned datasets must be of same lengths"
    assert not np.any(dataP['nrj'] - dataN['nrj']), "Energy points for dataP and dataN should be the same"

    muXAS = dataP['muA'] + dataN['muA']
    muXMCD = dataP['muA'] - dataN['muA']

    # standard error is the same for XAS and XMCD
    sigma = np.sqrt(dataP['sterrA']**2 + dataN['sterrA']**2)

    res = np.empty(len(muXAS), dtype=[('nrj', 'f8'), ('muXAS', 'f8'), ('sigmaXAS', 'f8'),
        ('muXMCD', 'f8'), ('sigmaXMCD', 'f8')])
    res['nrj'] = dataP['nrj']
    res['muXAS'] = muXAS
    res['muXMCD'] = muXMCD
    res['sigmaXAS'] = sigma
    res['sigmaXMCD'] = sigma

    return res