import numpy as np

from . import PowerLaw
from ..distributions import Debye, ColeCole, ColeDavidson, KWW, HavriliakNegami
from ..utils.constants import epsilon0


class _AbstractBDS:
    name = 'Abstract'
    type = 'Dielectric Spectroscopy'
    equation = ''
    params = [r'\Delta\epsilon', r'\tau_{0}']
    bounds = [(0, None), (0, None)]
    choices = [
        ('x axis', 'xaxis', {'Frequency': 'freq', 'Omega': 'omega'})
    ]
    susceptibility = None
    iscomplex = True

    @classmethod
    def func(cls, x, *args, complex_mode: int = 0, xaxis: str = 'freq', **kwargs):
        # args[0] : Delta epsilon
        # args[1:] : every other parameter

        _w = _convert_x_to_omega(x, xaxis=xaxis)

        chi = args[0] * cls.susceptibility(_w, *args[1:], **kwargs)
        if complex_mode == 0:
            return chi
        elif complex_mode == 1:
            return chi.real
        elif complex_mode == 2:
            return chi.imag
        else:
            raise ValueError(f'{complex_mode!r} is not 0, 1, 2')


class DebyeBDS(_AbstractBDS):
    name = 'Debye'
    equation = r'\Delta\epsilon / (1-i\omega\tau)'
    susceptibility = Debye.susceptibility


class ColeColeBDS(_AbstractBDS):
    name = 'Cole-Cole'
    equation = r'\Delta\epsilon / (1-i\omega\tau)^{\alpha}'
    params = _AbstractBDS.params + [r'\alpha']
    bounds = _AbstractBDS.bounds + [(0, 1)]
    susceptibility = ColeCole.susceptibility


class ColeDavidsonBDS(_AbstractBDS):
    name = 'Cole-Davidson'
    equation = r'\Delta\epsilon / [1-(i\omega\tau)^{\gamma})'
    params = _AbstractBDS.params + [r'\gamma']
    bounds = _AbstractBDS.bounds + [(0, 1)]
    susceptibility = ColeDavidson.susceptibility


class HavriliakNegamiBDS(_AbstractBDS):
    name = 'Havriliak-Negami'
    equation = r'\Delta\epsilon / [1-(i\omega\tau)^{\gamma}]^{\alpha}'
    params = _AbstractBDS.params + [r'\alpha', r'\gamma']
    bounds = _AbstractBDS.bounds + [(0, 1), (0, None)]
    susceptibility = HavriliakNegami.susceptibility


class HavriliakNegamiAlphaGammaBDS:
    type = 'Dielectric Spectroscopy'
    name = 'Havriliak-Negami (ind. slopes)'
    equation = r'\Delta\epsilon / [1-(i\omega\tau)^{\gamma}]^{\alpha}'
    params = [r'\Delta\epsilon', r'\tau_{0}', r'\alpha', r'\alpha\gamma']
    bounds = [(0, None), (0, None), (0, 1), (0, 1)]
    choices = [
        ('x axis', 'xaxis', {'Frequency': 'freq', 'Omega': 'omega'})
    ]
    iscomplex = True

    @staticmethod
    def func(x, deps, tau, alpha, alphagamma, complex_mode: int = 0, xaxis: str = 'freq', **kwargs):
        _w = _convert_x_to_omega(x, xaxis=xaxis)

        chi = deps * HavriliakNegami.susceptibility(_w, tau, alpha, alphagamma/alpha, **kwargs)
        if complex_mode == 0:
            return chi
        elif complex_mode == 1:
            return chi.real
        elif complex_mode == 2:
            return chi.imag
        else:
            raise ValueError(f'{complex_mode!r} is not 0, 1, 2')


class KWWBDS(_AbstractBDS):
    name = 'KWW'
    params = _AbstractBDS.params + [r'\beta']
    bounds = _AbstractBDS.bounds + [(0.1, 1)]
    susceptibility = KWW.susceptibility


class EpsInfty:
    name = 'Epsilon Infinity'
    type = 'Dielectric Spectroscopy'
    equation = r'\epsilon_{\infty}'
    params = [r'\epsilon_{\infty}']
    bounds = [(None, None)]
    iscomplex = True

    @staticmethod
    def func(x, eps, complex_mode: int = 0):
        if complex_mode == 0:
            ret_val = np.zeros(x.shape, dtype=complex)
            ret_val += eps
        elif complex_mode == 1:
            ret_val = eps * np.ones(x.shape)
        elif complex_mode == 2:
            ret_val = np.zeros(x.shape)
        else:
            raise ValueError(f'{complex_mode!r} is not 0, 1, 2')

        return ret_val


class HNWithHF:
    name = 'HN + HF wing'
    type = 'Dielectric Spectroscopy'
    equation = r'\Delta\epsilon HN(\omega, \tau, \alpha, \gamma) / CD(\omega, \tau_{c}, \alpha\gamma-\delta)'
    params = [r'\Delta\epsilon', r'\tau', r'\alpha', r'\gamma', r'\tau_{c}', '\delta']
    bounds = [(0, None), (0, None), (0, 1), (0, 1), (0, None), (0, 1)]
    choices = [
        ('x axis', 'xaxis', {'Frequency': 'freq', 'Omega': 'omega'})
    ]
    iscomplex = True

    @staticmethod
    def func(x, deps, tau, alpha, gamma, tauc, delta, complex_mode: int = 0, xaxis: str = 'freq'):
        _w = _convert_x_to_omega(x, xaxis=xaxis)

        numerator = (1 + 1j*_w*tauc) ** (gamma-delta)
        denominator = (1 + (1j*_w*tau)**alpha) ** gamma

        epsilon = deps * np.conjugate(numerator / denominator)

        if complex_mode == 0:
            return epsilon
        elif complex_mode == 1:
            return epsilon.real
        elif complex_mode == 2:
            return epsilon.imag
        else:
            raise ValueError(f'{complex_mode!r} has not value 0, 1, or 2')


class _CCWithHF:
    name = 'CC + HF wing'
    type = 'Dielectric Spectroscopy'
    equation = r'\Delta\epsilon CC(\omega, \tau, \alpha) / CD(\omega, \tau_{c}, \beta-\delta)'
    params = [r'\Delta\epsilon', r'\tau', r'\alpha', r'\tau_{c}', '\delta']
    bounds = [(0, None), (0, None), (0, 1), (0, None), (0, 1)]
    iscomplex = True

    @staticmethod
    def func(x, deps, tau, alpha, tauc, delta, complex_mode: int = 0, xaxis: str = 'freq'):
        return HNWithHF.func(x, deps, tau, alpha, 1, tauc, delta, complex_mode=complex_mode, xaxis=xaxis)


class _CDWithHF:
    name = 'CD + HF wing'
    type = 'Dielectric Spectroscopy'
    equation = r'\Delta\epsilon CD(\omega, \tau, \gamma) / CD(\omega, \tau_{c}, \gamma-\delta)'
    params = [r'\Delta\epsilon', r'\tau', r'\gamma', r'\tau_{c}', '\delta']
    bounds = [(0, None), (0, None), (0, 1), (0, None), (0, 1)]
    iscomplex = True

    @staticmethod
    def func(x, deps, tau, gamma, tauc, delta, complex_mode: int = 0, xaxis: str = 'freq'):
        return HNWithHF.func(x, deps, tau, 1, gamma, tauc, delta, complex_mode=complex_mode, xaxis=xaxis)


class PowerLawBDS:
    name = 'Power Law'
    type = 'Dielectric Spectroscopy'
    equation = r'A / (i\omega)^{n}'
    params = ['A', 'n']
    bounds = [(None, None), (None, None)]
    choices = [
        ('x axis', 'xaxis', {'Frequency': 'freq', 'Omega': 'omega'})
    ]
    iscomplex = True

    @staticmethod
    def func(x, a, n, complex_mode: int = 0, xaxis: str = 'freq'):
        _w = _convert_x_to_omega(x, xaxis=xaxis)

        if complex_mode == 0:
            ret_val = np.exp(1j*n*np.pi/2) * a / _w**n
        elif complex_mode == 1:
            ret_val = np.cos(n*np.pi/2) * a / _w**n
        elif complex_mode == 2:
            ret_val = np.sin(n*np.pi/2) * a / _w**n
        else:
            raise ValueError(f'{complex_mode!r} is not 0, 1, 2')

        return ret_val


class DCCondBDS:
    name = 'DC cond.'
    type = 'Dielectric Spectroscopy'
    equation = r'i\sigma_{dc}/\epsilon_{0}\omega'
    params = [r'\sigma_{dc}']
    bounds = [(0, None)]
    choices = [
        ('x axis', 'xaxis', {'Frequency': 'freq', 'Omega': 'omega'})
    ]
    iscomplex = True

    @staticmethod
    def func(x, sigma, complex_mode: int = 0, xaxis: str = 'freq'):
        _w = _convert_x_to_omega(x, xaxis=xaxis)

        if complex_mode == 0:
            ret_val = np.zeros(x.shape, dtype=complex)
            ret_val += 1j * sigma / _w / epsilon0
        elif complex_mode == 1:
            ret_val = np.zeros(x.shape)
        elif complex_mode == 2:
            ret_val = sigma / _w / epsilon0
        else:
            raise ValueError(f'{complex_mode!r} is not 0, 1, 2')

        return ret_val


class DerivativeHavriliakNegami:
    name = 'Havriliak-Negami (der.)'
    type = 'Dielectric Spectroscopy (derivative)'
    params = [r'\Delta\epsilon', r'\tau', r'\alpha', r'\gamma']
    choices = [
        ('x axis', 'xaxis', {'Frequency': 'freq', 'Omega': 'omega'})
    ]
    bounds = [(0, None), (0, None), (0, 1), (0, 1)]

    @staticmethod
    def func(x, eps, tau, a, g, xaxis: str = 'freq'):
        w = _convert_x_to_omega(x, xaxis=xaxis)

        omtau = w * tau
        theta = np.arctan2(np.sin(np.pi*a/2.), omtau**(-a) + np.cos(np.pi*a/2.))

        numer = a*g * omtau**a * np.cos(np.pi*a/2. - (1+g)*theta)
        denom = (1 + 2*omtau**a * np.cos(np.pi*a/2.) + omtau**(2*a))**((1+g)/2.)

        return eps * np.pi * numer / denom / 2.


class DerivativeColeCole:
    name = 'Cole-Cole (der.)'
    type = 'Dielectric Spectroscopy (derivative)'
    params = [r'\Delta\epsilon', r'\tau', r'\alpha']
    bounds = [(0, None), (0, None), (0, 1)]
    choices = [
        ('x axis', 'xaxis', {'Frequency': 'freq', 'Omega': 'omega'})
    ]

    @staticmethod
    def func(x, eps, tau, alpha, xaxis: str = 'freq'):
        w = _convert_x_to_omega(x, xaxis=xaxis)

        omtau = w * tau
        theta = np.arctan2(np.sin(np.pi*alpha/2.), omtau**(-alpha) + np.cos(np.pi*alpha/2.))

        numer = alpha * omtau**alpha * np.cos(np.pi*alpha/2. - 2*theta)
        denom = 1 + 2*omtau**alpha * np.cos(np.pi*alpha/2.) + omtau**(2*alpha)

        return eps * np.pi * numer / denom / 2.


class DerivativeColeDavidson:
    name = 'Cole-Davidson (der.)'
    type = 'Dielectric Spectroscopy (derivative)'
    params = [r'\Delta\epsilon', r'\tau', r'\gamma']
    bounds = [(0, None), (0, None), (0, 1)]
    choices = [
        ('x axis', 'xaxis', {'Frequency': 'freq', 'Omega': 'omega'})
    ]

    @staticmethod
    def func(x, eps, tau, g, xaxis: str = 'freq'):
        w = _convert_x_to_omega(x, xaxis=xaxis)
        omtau = w * tau
        numer = g * omtau * np.sin((1+g)*np.arctan(omtau))
        denom = (1 + omtau**2)**((1+g)/2.)

        return eps * np.pi * numer / denom / 2.


class _DerivativeHNWithHF:
    name = 'HN + HF wing (der.)'
    type = 'Dielectric Spectroscopy (derivative)'
    params = [r'\Delta\epsilon', r'\tau', r'\alpha', r'\gamma', r'\tau_{c}', r'\delta']
    bounds = [(0, None), (0, None), (0, 1), (0, 1), (0, None), (0, 1)]
    choices = [
        ('x axis', 'xaxis', {'Frequency': 'freq', 'Omega': 'omega'})
    ]

    @staticmethod
    def func(x, deps, tau, alpha, gamma, tauc, delta, xaxis: str = 'freq'):
        w = _convert_x_to_omega(x, xaxis=xaxis)

        a_pi2 = alpha*np.pi/2
        w_lamb = (w*tau)**alpha
        w_mu = w*tauc

        ag_d = alpha*gamma - delta

        phi = np.sin(a_pi2) * w_lamb / (1+np.cos(a_pi2) * w_lamb)
        chi = ag_d*np.arctan(w_mu) - gamma*np.arctan(phi)
        m_squared = 1 + w_mu**2
        l_squared = 1 + 2*np.cos(a_pi2)*w_lamb + w_lamb**2

        term1 = (ag_d * w_mu**2 / m_squared - alpha*gamma * w_lamb * (w_lamb + np.cos(a_pi2)) / l_squared) * np.cos(chi)
        term2 = (ag_d * w_mu / m_squared - alpha*gamma*np.sin(a_pi2) * w_lamb / l_squared) * np.sin(chi)

        return deps * (term2 - term1) * m_squared**(ag_d/2) / l_squared**(gamma/2)


class DerivativeCCWithHF:
    name = 'CC + HF wing (der.)'
    type = 'Dielectric Spectroscopy (derivative)'
    params = [r'\Delta\epsilon', r'\tau', r'\alpha', r'\tau_{c}', r'\delta']
    bounds = [(0, None), (0, None), (0, 1), (0, None), (0, 1)]
    choices = [
        ('x axis', 'xaxis', {'Frequency': 'freq', 'Omega': 'omega'})
    ]

    @staticmethod
    def func(x, deps, tau, alpha, tauc, delta, xaxis: str = 'freq'):
        return _DerivativeHNWithHF.func(x, deps, tau, alpha, 1, tauc, delta, xaxis=xaxis)


class DerivativeCDWithHF:
    name = 'CD + HF wing (der.)'
    type = 'Dielectric Spectroscopy (derivative)'
    params = [r'\Delta\epsilon', r'\tau', r'\gamma', r'\tau_{c}', r'\delta']
    bounds = [(0, None), (0, None), (0, 1), (0, None), (0, 1)]
    choices = [
        ('x axis', 'xaxis', {'Frequency': 'freq', 'Omega': 'omega'})
    ]

    @staticmethod
    def func(x, deps, tau, gamma, tauc, delta, xaxis: str = 'freq'):
        return _DerivativeHNWithHF.func(x, deps, tau, 1, gamma, tauc, delta, xaxis=xaxis)


class PowerLawBDSDer:
    name = 'Power Law'
    type = 'Dielectric Spectroscopy (derivative)'
    equation = r'A * \omega^{n}'
    params = ['A', 'n']
    bounds = [(None, None), (None, None)]
    choices = [
        ('x axis', 'xaxis', {'Frequency': 'freq', 'Omega': 'omega'})
    ]

    @staticmethod
    def func(x, a, n, xaxis: str = 'freq'):
        _w = _convert_x_to_omega(x, xaxis=xaxis)
        return a * _w ** n


def _convert_x_to_omega(x, xaxis: str = 'freq'):
    if xaxis not in ['freq', 'omega']:
        raise ValueError(f'Argument `xaxis` is `freq` or `omega`, given is {xaxis!r}')

    if xaxis == 'freq':
        _w = 2 * np.pi * x
    else:
        _w = x

    return _w
