import numpy as np

from ..distributions import *
from ..distributions.energy import EnergyBarriers
from ..distributions.intermolecular import FFHS
from ..nmr.relaxation import Relaxation
from ..utils.constants import gamma


class _AbstractFC:
    name = 'Havriliak-Negami'
    type = 'Field Cycling'
    equation = ''
    params = ['C', r'\tau']
    bounds = [(0, None), (0, None)]
    choices = [('x axis', 'xaxis', {'Frequency': 'freq', 'Omega': 'omega'}),
               ('y axis', 'yaxis', {'omega / T_1': 'chi', '1/ T_1': 'rate', 'T_1': 'time'})]
    relax = Relaxation()

    @classmethod
    def func(cls, x, c, *args, xaxis='freq', yaxis='chi'):
        return cls._calc_relax(x, c, *args, xaxis=xaxis, yaxis=yaxis)

    @classmethod
    def _calc_relax(cls, x, c, *args, xaxis='freq', yaxis='chi', is_bpp=True, **kwargs):
        _x = x
        if xaxis == 'freq':
            _x = 2*np.pi * x

        if is_bpp:
            r1 = cls.relax.t1(_x, *args, prefactor=c, inverse=False)
        else:
            r1 = cls.relax.t1_dipolar(_x, *args, prefactor=c, inverse=False, **kwargs)

        if yaxis == 'rate':
            return r1
        elif yaxis == 'chi':
            return _x * r1
        elif yaxis == 'time':
            return 1. / r1
        else:
            raise ValueError(f'Unknown yaxis option {yaxis}, not `chi`, `rate`, `time`.')


class ColeColeFC(_AbstractFC):
    name = 'Cole-Cole'
    params = _AbstractFC.params + [r'\alpha']
    bounds = _AbstractFC.bounds + [(0, 1)]
    relax = Relaxation(distribution=ColeCole)


class ColeDavidsonFC(_AbstractFC):
    name = 'Cole-Davidson'
    params = _AbstractFC.params + [r'\gamma']
    bounds = _AbstractFC.bounds + [(0, 1)]
    relax = Relaxation(distribution=ColeDavidson)


class HavriliakNegamiFC(_AbstractFC):
    name = 'Havriliak-Negami'
    params = _AbstractFC.params + [r'\alpha', r'\gamma']
    bounds = _AbstractFC.bounds + [(0, 1), (0, 1)]
    relax = Relaxation(distribution=HavriliakNegami)


class KWWFC(_AbstractFC):
    name = 'KWW'
    params = _AbstractFC.params + [r'\beta']
    bounds = _AbstractFC.bounds + [(0, 1)]
    relax = Relaxation(distribution=KWW)


class LogGaussianFC(_AbstractFC):
    name = 'Log-Gaussian'
    params = _AbstractFC.params + [r'\sigma']
    bounds = _AbstractFC.bounds + [(0, None)]
    relax = Relaxation(distribution=LogGaussian)


class FFHSFC(_AbstractFC):
    name = 'FFHS'
    relax = Relaxation(distribution=FFHS)


class EnergyFC(_AbstractFC):
    name = 'Energy distribution'
    params = ['C', 'T'] + EnergyBarriers.parameter
    bounds = [(0, None), (0, None), (0, None), (0, None)]
    ralax = Relaxation(distribution=EnergyBarriers)


class _AbstractFCDipolar(_AbstractFC):
    name = 'AbstractFC (het. dip.)'
    choices = _AbstractFC.choices + [(r'\gamma (obs.)', 'gamma_obs', gamma), (r'\gamma (coup.)', 'gamma_coup', gamma)]

    @classmethod
    def func(cls, x, c, *args, gamma_obs=gamma['1H'], gamma_coup=gamma['2H'],
             xaxis='freq', yaxis='chi'):
        return cls._calc_relax(x, c, *args, gamma_obs=gamma_obs, gamma_coup=gamma_coup,
                               xaxis='freq', yaxis='chi', is_bpp=False)


class HavriliakNegamiDipolar(_AbstractFCDipolar):
    name = 'Havriliak-Negami (dipolar)'
    params = _AbstractFCDipolar.params + [r'\alpha', r'\gamma']
    bounds = _AbstractFCDipolar.bounds + [(0, 1), (0, 1)]
    relax = Relaxation(distribution=HavriliakNegami)


class KWWDipolar(_AbstractFCDipolar):
    name = 'KWW (dipolar)'
    params = _AbstractFCDipolar.params + [r'\beta']
    bounds = _AbstractFCDipolar.bounds + [(0, 1)]
    relax = Relaxation(distribution=KWW)
