import scipy
import numpy as np


# Fitting
def fit_linear(S, V, **kwargs):
    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):
    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, log=False, **kwargs):
    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_function(function, x, y, p0=None, extend=0, **kwargs):
    """
    :param log: if the y-data is log-scaled
    """
    param, cov = scipy.optimize.curve_fit(function, x, y, p0, **kwargs)
    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):
    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)