""" Toolbox for SCS.

    Various utilities function to quickly process data measured at the SCS instruments.

    Copyright (2019) SCS Team.
"""
import matplotlib.pyplot as plt
import numpy as np
from scipy.special import erfc
from scipy.optimize import curve_fit

def knife_edge(nrun, axisKey='scannerX', signalKey='FastADC4peaks', p0=None, full=False, plot=False):
    ''' Calculates the beam radius at 1/e^2 from a knife-edge scan by fitting with
        erfc function: f(a, u) = a*erfc(u) or f(a, u) = a*erfc(-u) where 
        u = sqrt(2)*(x-x0)/w0 with w0 the beam radius at 1/e^2 and x0 the beam center.
        Inputs:
            nrun: xarray Dataset containing the detector signal and the motor 
                  position.
            axisKey: string, key of the axis against which the knife-edge is 
                  performed.
            signalKey: string, key of the detector signal.
            p0: list, initial parameters used for the fit: x0, w0, a. If None, a beam
                radius of 100 um is assumed.
            full: bool: If False, returns the beam radius and standard error. If True,
                returns the popt, pcov list of parameters and covariance matrix from 
                curve_fit.
            plot: bool: If True, plots the data and the result of the fit.
        Outputs:
            If full is False, ndarray with beam radius at 1/e^2 in mm and standard 
            error from the fit in mm. If full is True, returns popt and pcov from 
            curve_fit function.
    '''
    def integPowerUp(x, x0, w0, a):
        return a*erfc(-np.sqrt(2)*(x-x0)/w0)

    def integPowerDown(x, x0, w0, a):
        return a*erfc(np.sqrt(2)*(x-x0)/w0)

    #get the number of pulses per train from the signal source:
    dim = nrun[signalKey].dims[1]
    #duplicate motor position values to match signal shape
    #this is much faster than using nrun.stack()
    positions = np.repeat(nrun[axisKey].values, 
                          len(nrun[dim])).astype(nrun[signalKey].dtype)
    #sort the data to decide which fitting function to use
    sortIdx = np.argsort(positions)
    positions = positions[sortIdx]
    intensities = nrun[signalKey].values.flatten()[sortIdx]

    # estimate a linear slope fitting the data to determine which function to fit
    slope = np.cov(positions, intensities)[0][1]/np.var(positions) 
    if slope < 0:
        func = integPowerDown
        funcStr = 'a*erfc(np.sqrt(2)*(x-x0)/w0)'
    else:
        func = integPowerUp
        funcStr = 'a*erfc(-np.sqrt(2)*(x-x0)/w0)'
    if p0 is None:
        p0 = [np.mean(positions), 0.1, np.max(intensities)/2]
    popt, pcov = curve_fit(func, positions, intensities, p0=p0)
    print('fitting function:', funcStr)
    print('w0 = (%.1f +/- %.1f) um'%(popt[1]*1e3, pcov[1,1]**0.5*1e3))
    print('x0 = (%.3f +/- %.3f) mm'%(popt[0], pcov[0,0]**0.5*1e3))
    print('a = %e +/- %e '%(popt[2], pcov[2,2]**0.5*1e3))
    
    if plot:
        xfit = np.linspace(positions.min(), positions.max(), 1000)
        yfit = func(xfit, *popt)
        plt.figure(figsize=(7,4))
        plt.scatter(positions, intensities, color='C1', label='exp', s=2, alpha=0.01)
        plt.plot(xfit, yfit, color='C4', 
                 label=r'fit $\rightarrow$ $w_0=$(%.1f $\pm$ %.1f) $\mu$m'%(popt[1]*1e3, pcov[1,1]**0.5*1e3))
        leg = plt.legend()
        for lh in leg.legendHandles: 
            lh.set_alpha(1)
        plt.ylabel(signalKey)
        plt.xlabel(axisKey + ' position [mm]')
        plt.tight_layout()
    if full:
        return popt, pcov
    else:
        return np.array([popt[1], pcov[1,1]**0.5])