import scipy
import scipy.odr
import numpy as np


# Fitting
def fit_linear(S, V, **kwargs):
    """Fit a linear function to the data
    
    :param S: data x values
    :param V: data y values
    :param xerr: data x uncertainties
    :param yerr: data y uncertainties
    :param p0: initial values of parameters
    :param odr: Fit method to use: True means ortogonal distance regression (ODR), False means ordinary least square fit, None (default) uses ODR only if either xerr or yerr is not None
    :param extend: fraction by which the returned fit trace extends beyond the first/last data point
    """
    result = fit_function(lambda s, v0, m: v0+m*s, S, V, p0=(np.mean(V), np.mean(np.diff(V)/np.diff(S))), **kwargs)
    param, param_error, function, [X, _] = result
    return param, param_error, function, [X, function(X, *param)]

def fit_gaussian(S, V, **kwargs):
    """Fit a gaussian function to the data
    
    :param S: data x values
    :param V: data y values
    :param xerr: data x uncertainties
    :param yerr: data y uncertainties
    :param p0: initial values of parameters
    :param odr: Fit method to use: True means ortogonal distance regression (ODR), False means ordinary least square fit, None (default) uses ODR only if either xerr or yerr is not None
    :param extend: fraction by which the returned fit trace extends beyond the first/last data point
    """
    gauss = lambda s, v0, vp, s0, sigma: v0+vp*np.exp(-0.5*((s-s0)/sigma)**2)
    v0 = np.min(V)
    vp = np.max(V)-v0
    s0 = S[np.argmax(V)]
    sigma = np.sqrt(np.abs(np.sum((S - s0) ** 2 * V) / np.sum(V)))
    result = fit_function(gauss, S, V, p0=(v0, vp, s0, sigma), **kwargs)
    param, param_error, function, [X, _] = result
    param[3] = np.abs(param[3]) # return the positive sigma (could use bounds instead, but that's more likely to fail)
    return param, param_error, function, [X, function(X, *param)]

def fit_lorenzian(S, V, **kwargs):
    """Fit a lorenzian function to the data
    
    :param S: data x values
    :param V: data y values
    :param xerr: data x uncertainties
    :param yerr: data y uncertainties
    :param p0: initial values of parameters
    :param odr: Fit method to use: True means ortogonal distance regression (ODR), False means ordinary least square fit, None (default) uses ODR only if either xerr or yerr is not None
    :param extend: fraction by which the returned fit trace extends beyond the first/last data point
    """
    lorenzian = lambda s, v0, vp, s0, gamma: v0 + vp/(1+((s-s0)/gamma)**2)
    v0 = np.min(V)
    vp = np.max(V)-v0
    s0 = S[np.argmax(V)]
    sigma = np.sqrt(np.abs(np.sum((S - s0) ** 2 * (V-v0)) / np.sum(V-v0)))
    result = fit_function(lorenzian, S, V, p0=(v0, vp, s0, sigma), **kwargs)
    param, param_error, function, [X, _] = result
    param[3] = np.abs(param[3]) # return the positive sigma (could use bounds instead, but that's more likely to fail)
    return param, param_error, function, [X, function(X, *param)]


def fit_double_lorenzian(S, V, *, s0=None, dry_run=False, **kwargs):
    """Fit a lorenzian function to the data
    
    :param S: data x values
    :param V: data y values
    :param xerr: data x uncertainties
    :param yerr: data y uncertainties
    :param s0: initial peak positions
    :param p0: initial values of parameters
    :param odr: Fit method to use: True means ortogonal distance regression (ODR), False means ordinary least square fit, None (default) uses ODR only if either xerr or yerr is not None
    :param extend: fraction by which the returned fit trace extends beyond the first/last data point
    """
    lorenzian = lambda s, v0, vp, s0, gamma, vp2, s02, gamma2: v0 + vp/(1+((s-s0)/gamma)**2) + vp2/(1+((s-s02)/gamma2)**2)
    
    # ! for double peak fit, initial conditions are very important
    v0 = np.min(V)
    vp = np.max(V)-v0
    if s0 is None:
        s0 = S[np.argmax(V)]
        sigma = np.sqrt(np.abs(np.sum((S - s0) ** 2 * (V-v0)) / np.sum(V-v0)))/10
        s0 = [s0-3*sigma, s0+3*sigma]    
    sigma0 = np.sqrt(np.abs(np.sum((S - s0[0]) ** 2 * (V-v0)) / np.sum(V-v0)))/10
    sigma1 = np.sqrt(np.abs(np.sum((S - s0[1]) ** 2 * (V-v0)) / np.sum(V-v0)))/10
    p0=(v0, vp, s0[0], sigma0, vp, s0[1], sigma1)
    
    if dry_run:
        mi, ma = min(S), max(S)
        X = np.linspace(mi, ma, 1000)
        return p0, (0,0,0,0), 0, [X, lorenzian(X, *p0)]
        
    result = fit_function(lorenzian, S, V, p0=p0, **kwargs)
    param, param_error, function, [X, _] = result
    param[3] = np.abs(param[3]) # return the positive sigma (could use bounds instead, but that's more likely to fail)
    return param, param_error, function, [X, function(X, *param)]
    

def fit_function(function, x, y, *, xerr=None, yerr=None, p0=None, odr=None, extend=0):
    """Fit a function to the data
    
    :param function: fit function(x, *param)
    :param x: data x values
    :param y: data y values
    :param xerr: data x uncertainties
    :param yerr: data y uncertainties
    :param p0: initial values of parameters
    :param odr: Fit method to use: True means ortogonal distance regression (ODR), False means ordinary least square fit, None (default) uses ODR only if either xerr or yerr is not None
    :param extend: fraction by which the returned fit trace extends beyond the first/last data point
    """
    if odr is None:
        odr = xerr is not None or yerr is not None
    
    if odr:
        # orthogonal distance regression
        data = scipy.odr.RealData(x, y, sx=xerr, sy=yerr)
        model = scipy.odr.Model(lambda beta, x: function(x, *beta))
        odr = scipy.odr.ODR(data, model, p0)
        # if not odr: odr.set_job(fit_type=2) # ordinary least-squares
        output = odr.run()
        param, param_error = output.beta, output.sd_beta
    else:
        # non-linear least squares
        param, cov = scipy.optimize.curve_fit(function, x, y, p0)
        param_error = np.sqrt(np.abs(cov.diagonal()))
    
    mi, ma = min(x), max(x)
    X = np.linspace(mi - (ma-mi)*extend, ma + (ma-mi)*extend, 1000)
    return param, param_error, function, [X, function(X, *param)]

def fit_exponential(S, V, **kwargs):
    """Fit an exponential function to the data
    
    :param S: data x values
    :param V: data y values
    :param xerr: data x uncertainties
    :param yerr: data y uncertainties
    :param p0: initial values of parameters
    :param odr: Fit method to use: True means ortogonal distance regression (ODR), False means ordinary least square fit, None (default) uses ODR only if either xerr or yerr is not None
    :param extend: fraction by which the returned fit trace extends beyond the first/last data point
    """
    exponential = lambda s, v0, vp, s0: v0 + vp*np.exp(s/s0)
    v0 = np.min(V)
    vp = np.max(V)-v0
    s0 = 1
    return fit_function(exponential, S, V, p0=(v0, vp, s0), **kwargs)