""" Mathematical functions and models """

import numpy as np
import scipy
import scipy.constants
import qtt.pgeometry
import matplotlib.pyplot as plt


def gaussian(x, mean, s, amplitude=1, offset=0):
    """ Model for Gaussuan function

       $$y = offset + amplitude * np.exp(-(1/2)*(x-mean)^2/s^2)$$

    Args:
        x (array): data points
        mean, std, amplitude, offset: parameters
    Returns:
        y (array)

    """
    y = offset + amplitude * np.exp(- (x - mean) * (x - mean) / (2 * s * s))
    return y


def double_gaussian(signal_amp, params):
    """ A model for the sum of two Gaussian distributions. 

    Args:
        signal_amp (array): x values of the data
        params (array): parameters of the two gaussians, [A_dn, A_up, sigma_dn, sigma_up, mean_dn, mean_up]
            amplitude of first (second) gaussian = A_dn (A_up) 
            standard deviation of first (second) gaussian = sigma_dn (sigma_up)
            average value of the first (second) gaussian = mean_dn (mean_up)

    Returns:
        double_gauss (np.array): model of a double gaussian
    """
    [A_dn, A_up, sigma_dn, sigma_up, mean_dn, mean_up] = params
    gauss_dn = gaussian(signal_amp, params[4], params[2], params[0])
    gauss_up = gaussian(signal_amp, params[5], params[3], params[1])
    double_gauss = gauss_dn + gauss_up
    return double_gauss


def _cost_double_gaussian(signal_amp, counts, params):
    """ Cost function for fitting of double Gaussian. 

    Args:
        signal_amp (array): x values of the data
        counts (array): y values of the data
        params (array): parameters of the two gaussians, [A_dn, A_up, sigma_dn, sigma_up, mean_dn, mean_up]
            amplitude of first (second) gaussian = A_dn (A_up) 
            standard deviation of first (second) gaussian = sigma_dn (sigma_up)
            average value of the first (second) gaussian = mean_dn (mean_up)

    Returns:
        cost (float): value which indicates the difference between the data and the fit
    """
    model = double_gaussian(signal_amp, params)
    cost = np.linalg.norm(counts - model)
    return cost


def fit_double_gaussian(signal_amp, counts, maxiter=None, maxfun=5000, verbose=1, par_guess=None):
    """ Fitting of double gaussian
    
    Fitting the Gaussians and finding the split between the up and the down state, 
    separation between the max of the two gaussians measured in the sum of the std.

    Args:
        signal_amp (array): x values of the data
        counts (array): y values of the data
        maxiter (int): maximum number of iterations to perform
        maxfun (int): maximum number of function evaluations to make
        verbose (int): set to >0 to print convergence messages
        par_guess (None or array): optional, initial guess for the fit parameters: 
            [A_dn, A_up, sigma_dn, sigma_up, mean_dn, mean_up]

    Returns:
        par_fit (array): fit parameters of the double gaussian: [A_dn, A_up, sigma_dn, sigma_up, mean_dn, mean_up]
        par_guess (array): initial guess for the fit parameters, either the ones give to the function, or generated by the function: [A_dn, A_up, sigma_dn, sigma_up, mean_dn, mean_up]
        separation (float): separation between the max of the two gaussians measured in the sum of the std
        split (float): value that seperates the up and the down level

    """
    def func(params): return _cost_double_gaussian(signal_amp, counts, params)
    maxsignal = np.percentile(signal_amp, 98)
    minsignal = np.percentile(signal_amp, 2)
    if par_guess is None:
        A_dn = np.max(counts[:int((len(counts) / 2))])
        A_up = np.max(counts[int((len(counts) / 2)):])
        sigma_dn = (maxsignal - minsignal) * 1 / 20
        sigma_up = (maxsignal - minsignal) * 1 / 20
        mean_dn = minsignal + 1 / 4 * (maxsignal - minsignal)
        mean_up = minsignal + 3 / 4 * (maxsignal - minsignal)
        par_guess = np.array([A_dn, A_up, sigma_dn, sigma_up, mean_dn, mean_up])
    par_fit = scipy.optimize.fmin(func, par_guess, maxiter=maxiter, maxfun=maxfun, disp=verbose >= 2)

    # separation is the difference between the max of the gaussians devided by the sum of the std of both gaussians
    separation = (par_fit[5] - par_fit[4]) / (abs(par_fit[2]) + abs(par_fit[3]))
    # split equal distant to both peaks measured in std from the peak
    split = par_fit[4] + separation * abs(par_fit[2])

    result_dict = {'parameters initial guess': par_guess, 'separation': separation, 'split': split}

    return par_fit, result_dict


def exp_function(x, a, b, c):
    """ Model for exponential function

       $$y = a + b * np.exp(-c * x)$$

    Args:
        x (array): x values of the data
        a = offset
        b = starting value
        c = 1/typical decay time
    Returns:
        y (array): model for exponantial decay

    """
    y = a + b * np.exp(-c * x)
    return y


def cost_exp_decay(x_data, y_data, params, threshold=None):
    """ Cost function for exponential decay.

    Args:
        x_data (array): the data for the input variable
        y_data (array): the data for the measured variable
        params (array): parameters of the exponential decay function, [A,B, gamma]
        threshold (float or None or 'auto'): if the difference between data and model is larger then the threshold, then the cost penalty is reduced.
            If None use normal cost function. If 'auto' use automatic detection (at 95th percentile)
    Returns:
        cost (float): value which indicates the difference between the data and the fit
    """
    model = exp_function(x_data, params[0], params[1], params[2])
    cost = qtt.pgeometry.robustCost(y_data - model, thr=threshold)
    cost = np.linalg.norm(cost)
    return cost


def test_cost_exp_decay():
    params=[0,1,1]
    x_data=np.arange(0, 20)
    y_data=exp_function(x_data, *params)
    y_data[-1]+=10
    c=cost_exp_decay(x_data, y_data, params)
    assert(c==10.0)
    c=cost_exp_decay(x_data, y_data, params, threshold=None)
    assert(c==10.0)
    c=cost_exp_decay(x_data, y_data, params, threshold='auto')
    assert(c<10.0)
    
def fit_exp_decay(x_data, y_data, maxiter=None, maxfun=5000, verbose=1, par_guess=None, threshold=None):
    """ Fit a exponential decay. 

    Args:
        x_data (array): the data for the input variable
        y_data (array): the data for the measured variable
        maxiter (int): maximum number of iterations to perform
        maxfun (int): maximum number of function evaluations to make
        verbose (int): set to >0 to print convergence messages
        par_guess (None or array): optional, initial guess for the fit parameters: [A,B, gamma]
        threshold (float or None): threshold for the cost function.
            If the difference between data and model is larger then the threshold, these data are not taken into account for the fit. 
            If None use automatic detection (at 95th percentile)
    Returns:
        par_fit (array): fit parameters of the exponential decay, [A,B, gamma]

    """

    func = lambda params: cost_exp_decay(x_data, y_data, params, threshold)
    if par_guess is None:
        maxsignal = np.percentile(y_data, 98)
        minsignal = np.percentile(y_data, 2)
        gamma = 1 / (x_data[int(len(x_data) / 2)])
        A = minsignal
        B = maxsignal
        par_guess = np.array([A, B, gamma])

    par_fit = scipy.optimize.fmin(func, par_guess, maxiter=maxiter, maxfun=maxfun, disp=verbose >= 2)
    return par_fit


def gauss_ramsey(x_data, params):
    """ Model for the measurement result of a pulse Ramsey sequence while varying the free evolution time, the phase of the second pulse 
    is made dependent on the free evolution time. This results in a gaussian decay multiplied by a sinus. 
    Function as used by T.F. Watson et all., example in qtt/docs/notebooks/example_fit_ramsey.ipynb
    
    $$ gauss_ramsey = A * exp(-(C*x_data)**2) * sin(2pi*ramseyfreq * x_data - angle) +B  $$
    
    Args:
        x_data (array): the data for the input variable
        params (array): parameters of the gauss_ramsey function, [A,C,ramseyfreq,angle,B]
        
    Result:
        gauss_ramsey (array): model for the gauss_ramsey
    """
    [A,C,ramseyfreq,angle,B] = params    
    gauss_ramsey = A * np.exp(-(C*x_data)**2) * np.sin(2*np.pi*ramseyfreq * x_data - angle) +B 
    return gauss_ramsey


def cost_gauss_ramsey(x_data, y_data, params, weight_power=0):
    """ Cost function for gauss_ramsey. 
    
    Args:
        x_data (array): the data for the input variable
        y_data (array): the data for the measured variable
        params (array): parameters of the gauss_ramsey function, [A,C,ramseyfreq,angle,B]
        weight_power (float)
        
    Returns:
        cost (float): value which indicates the difference between the data and the fit
    """
    model = gauss_ramsey(x_data, params)
    #cost = np.sum([(np.array(y_data)[1:] - np.array(model)[1:])**2*(np.array(x_data)[1:]-np.array(x_data)[:-1])**weight_power])
    cost = np.sum([(np.array(y_data)[1:] - np.array(model)[1:])**2*(np.diff(x_data))**weight_power])
    return cost

def fit_gauss_ramsey(x_data, y_data, weight_power=0, maxiter=None, maxfun=5000, verbose=1, par_guess=None):
    """ Fit a gauss_ramsey. The function gauss_ramsey gives a model for the measurement result of a pulse Ramsey sequence while varying 
    the free evolution time, the phase of the second pulse is made dependent on the free evolution time.
    This results in a gaussian decay multiplied by a sinus. Function as used by T.F. Watson et all., 
    see function 'gauss_ramsey' and example in qtt/docs/notebooks/example_fit_ramsey.ipynb
    
    Args:
        x_data (array): the data for the independant variable
        y_data (array): the data for the measured variable
        weight_power (float)
        maxiter (int): maximum number of iterations to perform
        maxfun (int): maximum number of function evaluations to make
        verbose (int): set to >0 to print convergence messages
        par_guess (None or array): optional, initial guess for the fit parameters: [A,C,ramseyfreq,angle,B]
        
    Returns:
        par_fit (array): array with the fit parameters: [A,C,ramseyfreq,angle,B]
        result_dict (dict): dictionary containing a description, the par_fit and par_guess
    
    """
    func = lambda params: cost_gauss_ramsey(x_data, y_data, params, weight_power=weight_power)
    if par_guess is None:
        A = (np.max(y_data) - np.min(y_data))/2
        C = 1/(1e-6)
        ramseyfreq = 1/(1e-6)
        angle = 0
        B = (np.min(y_data) + (np.max(y_data) - np.min(y_data))/2)
        par_guess = np.array([A,C,ramseyfreq,angle,B])
        
    par_fit = scipy.optimize.fmin(func, par_guess, maxiter=maxiter, maxfun=maxfun, disp=verbose >= 2)
    
    result_dict = {'description':'Function to analyse the results of a Ramsey experiment, fitted function: gauss_ramsey = A * exp(-(C*x_data)**2) * sin(2pi*ramseyfreq * x_data - angle) +B','parameters fit': par_fit,'parameters initial guess': par_guess}
    
    return par_fit, result_dict

def test_fit_gauss_ramsey(fig=None):
    """ Tests if the function fit_gauss_ramsey is working
    
    Args:
        None
        
    Returns:
        Nothing
    """
    y_data = np.array([0.6019, 0.5242, 0.3619, 0.1888, 0.1969, 0.3461, 0.5276, 0.5361,
       0.4261, 0.28  , 0.2323, 0.2992, 0.4373, 0.4803, 0.4438, 0.3392,
       0.3061, 0.3161, 0.3976, 0.4246, 0.398 , 0.3757, 0.3615, 0.3723,
       0.3803, 0.3873, 0.3873, 0.3561, 0.37  , 0.3819, 0.3834, 0.3838,
       0.37  , 0.383 , 0.3573, 0.3869, 0.3838, 0.3792, 0.3757, 0.3815])
    x_data = np.array([i*1.6/40 for i in range(40)])
    
    par_fit_test,_ = fit_gauss_ramsey(x_data*1e-6,y_data)
    
    assert(np.abs(np.abs(par_fit_test[0])-0.21)<0.1)
    assert(np.abs(par_fit_test[-2]-1.88)<0.1)
    assert(np.abs(par_fit_test[-1]-0.38)<0.1)
    
    test_x = np.linspace(0,x_data.max()*1e-6,200)
    test_y = gauss_ramsey(test_x,par_fit_test)
    
    if fig is not None:
        plt.figure(10);
        plt.clf()
        plt.plot(x_data, y_data, 'o', label='input data')
        plt.plot(test_x*1e6, test_y, label='fitted curve')
        plt.legend(numpoints=1)
        
if __name__ == '__main__':
    test_fit_gauss_ramsey(fig=10)    

def linear_function(x, a, b):
    """ Linear function with offset"""
    return a * x + b


def Fermi(x, cc, A, T, kb=1):
    """ Fermi distribution 

    Arguments:
        x (numpy array): independent variable
        cc (float): center of Fermi distribution
        A (float): amplitude of Fermi distribution
        T (float): temperature Fermi distribution
        kb (float, default: 1): temperature scaling factor

    Returns:
        y (numpy array): value of the function

    .. math::

        y =  A*(1/ (1+\exp( (x-cc)/(kb*T) ) ) )
    """
    y = A * 1. / (1 + np.exp((x - cc) / (kb * T)))
    return y


def FermiLinear(x, a, b, cc, A, T, l=1.16):
    """ Fermi distribution with linear function added 

    Arguments:
        x (numpy array): independent variable 
        a, b (float): coefficients of linear part
        cc (float): center of Fermi distribution
        A (float): amplitude of Fermi distribution
        T (float): temperature Fermi distribution in mili Kelvin
        l (float): leverarm divided by kb

    The default value of the leverarm is
        (100 ueV/mV)/kb = (100*1e-6*scipy.constants.eV )/kb = 1.16.

    For this value the input variable x should be in mV, the 
    temperature T in K. We input the leverarm divided by kb for numerical stability.

    Returns:
        y (numpy array): value of the function

    .. math::

        y = a*x + b + A*(1/ (1+\exp( l* (x-cc)/(T) ) ) )

    """
    y = a * x + b + A * 1. / (1 + np.exp(l * (x - cc) / (T)))
    return y


def logistic(x, x0=0, alpha=1):
    """ Logistic function

    Args:
        x (array): array with values
        x0 (float):  parameters of function
        alpha (float): parameters of function

    Example:
        >>> y=logistic(0, 1, alpha=1)
    """
    f = 1 / (1 + np.exp(-2 * alpha * (x - x0)))
    return f


def test_logistic_and_linear_function():
    x_data = np.arange(-10, 10, 0.1)

    _ = logistic(x_data, x0=0, alpha=1)
    assert logistic(0, x0=0, alpha=1) == 0.5

    _ = linear_function(x_data, 1, 2)
    assert linear_function(0, 1, 2) == 2
    assert linear_function(3, 1, 2) == 5


if __name__ == '__main__':
    test_logistic_and_linear_function()
