import numpy as np

from ..utils.constants import gamma


class FCEval(object):
    name = 'FC Relaxation'
    params = None
    equation = ''

    @staticmethod
    def func(p, x, *args):
        xmode, ymode, func_args = args
        _x = x
        if xmode == 'f (Hz)':
            _x = x*2*np.pi
        if ymode == 'Susceptibility (no BPP)':
            return _x*FCEval.func_j(p, _x, *func_args)
        rate = FCEval.func_r1(p, _x, *func_args)
        if ymode == 'Rate':
            return rate
        elif ymode == 'Susceptibility':
            return _x*rate
        else:
            return 1/rate

    @staticmethod
    def func_r1(p, x, *args):
        counter = 0
        y = np.zeros_like(x)
        y_singles = []
        for arg in args:
            func_sd, func_cp, nuclei = arg
            try:
                nop_sd = len(func_sd.parameter)
            except TypeError:
                nop_sd = 0
            sd_params = p[counter:counter+nop_sd]
            counter += nop_sd
            tau = p[counter]
            counter += 1
            cp_params = [pp for pp in p[counter:counter+len(func_cp.parameter)]]
            counter += len(func_cp.parameter)

            cp_params.extend([gamma[n] for n in nuclei])
            prefactor = func_cp.func(*cp_params)

            if len(nuclei) == 2:
                t1 = func_sd.calc_t1_dipolar(x, tau, sd_params, prefactor=prefactor,
                                             gamma1=nuclei[0], gamma2=nuclei[1],
                                             inverse=False)
            else:
                t1 = func_sd.calc_t1_bpp(x, tau, sd_params, prefactor=prefactor,
                                         inverse=False)
            y_singles.append(t1)
            y += t1

        return np.array(y)

    @staticmethod
    def func_j(p, x, *args):
        counter = 0
        y = np.zeros_like(x)
        y_singles = []
        for arg in args:
            func_sd, func_cp, nuclei = arg
            try:
                nop_sd = len(func_sd.parameter)
            except TypeError:
                nop_sd = 0
            sd_params = p[counter:counter + nop_sd]
            counter += nop_sd
            tau = p[counter]
            counter += 1
            cp_params = [pp for pp in p[counter:counter + len(func_cp.parameter)]]
            counter += len(func_cp.parameter)

            cp_params.extend([gamma[n] for n in nuclei])
            prefactor = func_cp.func(*cp_params)
            j = prefactor*func_sd.specdens(x, tau, *sd_params)
            y_singles.append(j)
            y += j

        return np.array(y)
