import numpy as     np
from   lmfit import Parameters, minimize

def fit_vpsd_coefficients(self) -> None:
    """Fit velocity power spectral density (VPSD) coefficients.

    :return: None
    :rtype: None
    """

    # read VPSD
    freq, vpsd, freq_avg, vpsd_avg = [self.vpsd[key] for key in ["freq", "vpsd", "freq_avg", "vpsd_avg"]]

    # LMFIT parameters
    params = Parameters()

    # loop components
    for comp in self.vpsd_components.keys():

        # component dictionary
        comp_dict = self.vpsd_components[comp]

        # coefficients and vary
        coef_val = comp_dict["coef_val"]
        vary     = comp_dict["vary"]

        # loop coefficients
        for i in range(len(coef_val)):

            # add parameters
            params.add(comp + "_" + str(i), value=coef_val[i], min=coef_val[i]/10, max=coef_val[i]*10, vary=vary[i])

    # fit coefficients
    c = minimize(_func_res, params, args=(self, freq_avg, vpsd_avg))

    # loop components
    for comp in self.vpsd_components.keys():

        # component dictionary
        comp_dict = self.vpsd_components[comp]

        # coefficients
        coef_val = comp_dict["coef_val"]
        coef_err = comp_dict["coef_err"]

        # loop coefficients
        for i in range(len(coef_val)):

            # update coefficients with fitted values
            coef_val[i] = c.params[comp + "_" + str(i)].value
            coef_err[i] = c.params[comp + "_" + str(i)].stderr
    
    return None

def _func_res(params, self, freq_avg, vpsd_avg):

    # empty array for sum of components
    vpsd_tot = np.zeros(len(freq_avg))

    # loop components
    for comp in self.vpsd_components.keys():

        # component dictionary
        comp_dict = self.vpsd_components[comp]

        # component type
        type = comp_dict["type"]

        # type Constant
        if type == "Constant":
            
            # unpack coefficients
            c0 = params[comp + "_0"]

            # compute component
            vpsd_comp = c0

        # type Lorentz
        if type == "Lorentz":

            # unpack coefficients
            c0 = params[comp + "_0"]
            c1 = params[comp + "_1"]
            c2 = params[comp + "_2"]

            # compute component
            vpsd_comp = c0*c1**2/(c1**2+(freq_avg-c2)**2)

        # type Harvey
        if type == "Harvey":

            # unpack coefficients
            c0 = params[comp + "_0"]
            c1 = params[comp + "_1"]
            c2 = params[comp + "_2"]

            # compute component
            vpsd_comp = c0/(1+(c1*freq_avg)**c2)
    
        # add component to sum
        vpsd_tot += vpsd_comp
    
    # logarithmic residual
    logres = np.log10(vpsd_avg) - np.log10(vpsd_tot)

    # return residual
    return logres