""" 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, 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.
            plot: bool: If True, plots the data and the result of the fit.
        Outputs:
            ndarray with beam radius at 1/e^2 in mm and standard error from the fit
            in mm.
    '''
    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]
    if intensities[0] > intensities[-1]:
        func = integPowerDown
    else:
        func = integPowerUp
        
    if p0 is None:
        p0 = [np.mean(positions), 0.1, np.max(intensities)/2]
    popt, pcov = curve_fit(func, positions, intensities, p0=p0)
    print('w0 = (%.1f +/- %.1f) um'%(popt[1]*1e3, pcov[1,1]**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()
    return np.array([popt[1], pcov[1,1]**0.5])