import lmfit as lm
import numpy as np
import scipy.special as sf
from scipy.optimize import curve_fit


def cosfunc(t, A, w, p, c, e):
    """
    definition of the sin function
    """

    return A * np.cos(2 * np.pi * w * t + p) + e * t + c


def fit_cos_eval(x, par):
    """
    
    par: dict
        fitting results
    
    """

    ys = cosfunc(x, par['amp'], par['freq'], par['phase'], par['offset'],
                 par['slope'])

    return ys


def regression_sine_fft():
    """
    fast fourier transformation for sine data
    """

    return []


def fit_cos_simple(x, y, freq=10.0):
    """
    simple sine regression
    
    x: vector or list
    y: vector or list
    
    freq: float
        initial frequency of regression analysis
        
    
    
    """

    tt = np.array(x)
    yy = np.array(y)

    guess_offset = np.mean(yy)
    offset_b = 0.4 * abs(guess_offset)

    guess_amp = abs(np.max(yy) - np.min(yy)) / 2.0

    param_bounds = ([
        0.3 * guess_amp, 0.99 * freq, -np.inf, guess_offset - offset_b, -np.inf
    ], [1.3 * guess_amp, 1.01 * freq, np.inf, guess_offset + offset_b, np.inf])

    popt, pcov = curve_fit(cosfunc, tt, yy, bounds=param_bounds)

    A, w, p, c, e = popt

    return {
        "amp": A,
        "freq": w,
        "phase": p,
        "offset": c,
        "slope": e,
    }


def fit_cos(x, y, freq=10.0):
    """
    sine regression
    
    x: vector or list
    y: vector or list
    
    freq: float
        initial frequency of regression analysis
        
    
    
    """

    # step 1

    res_step1 = fit_cos_simple(x, y, freq=freq)

    # step 2: lmfit
    res = {}

    mod = lm.models.Model(cosfunc)

    mod.set_param_hint(
        'A',
        value=res_step1['amp'],
        #min=res_step1['amp'] - 0.5 * abs(res_step1['amp']),
        #max=res_step1['amp'] + 0.5 * abs(res_step1['amp'])
    )

    mod.set_param_hint(
        'w',
        value=freq,
        vary=True,
        #min=res_step1['freq'] - 0.3 * abs(res_step1['freq']),
        #max=res_step1['freq'] + 0.3 * abs(res_step1['freq'])
    )

    mod.set_param_hint('p', value=res_step1['phase'], vary=True)

    mod.set_param_hint('c', value=res_step1['offset'],
                       vary=True)  #, min = -0.5, max = 0.5)

    mod.set_param_hint('e', value=res_step1['slope'], vary=True)

    parms_fit = [
        mod.param_hints['A']['value'], mod.param_hints['w']['value'],
        mod.param_hints['p']['value'], mod.param_hints['c']['value'],
        mod.param_hints['e']['value']
    ]

    abweichung = []
    chis = []
    chis_red = []
    results = []
    r2 = []

    methods = ['leastsq', 'powell']
    dof = len(y) - len(parms_fit)

    for method in methods:
        result = mod.fit(y, t=x, method=method, verbose=False)
        r2temp = 1 - result.residual.var() / np.var(y)
        #    r2temp = result.redchi / np.var(yfit, ddof=2)
        if r2temp < 0.:
            r2temp = 0
        r2.append(r2temp)

        chi = result.chisqr
        chis_red.append(result.redchi)
        abweichung.append(sf.gammaincc(dof / 2., chi / 2))
        chis.append(chi)
        results.append(result)

    ret = {}
    best = np.nanargmax(r2)

    res[f'amp'] = results[best].best_values['A']
    res[f'freq'] = results[best].best_values['w']
    res[f'phase'] = results[best].best_values['p']
    res[f'offset'] = results[best].best_values['c']
    res[f'slope'] = results[best].best_values['e']

    res[f'r2'] = r2[best]

    return res