import numpy as np
from scipy.special import psi

from .base import Distribution
from ..lib.utils import ArrayLike
from ..lib.decorator import adjust_dims
from ..math.mittagleffler import mlf
from ..utils import Eu


class HavriliakNegami(Distribution):
    """
    Functions based on Havriliak-Negami distribution
    """

    name = 'Havriliak-Negami'
    parameter = [r'\alpha', r'\gamma']
    bounds = [(0, 1), (0, 1)]

    @staticmethod
    def distribution(tau: ArrayLike, tau0: float, alpha: float, gamma: float) -> ArrayLike:
        r"""
        Distribution of correlation times :math:`G(\ln\tau)`.

        .. math::
            G(\ln\tau) = \frac{(\tau/\tau_0)^{\alpha\gamma}
            \sin\left\lbrace\gamma\arctan\left[\frac{\sin\alpha\pi}{(\tau/\tau_0)^\alpha + \cos\alpha\pi}\right] \right\rbrace}{\pi\left[1+2(\tau/\tau_0)^\alpha\cos\alpha\pi + (\tau/\tau_0)^{2\gamma}\right]}

        Args:
            tau (array_like) :
            tau0 (array_like) :
            alpha:
            gamma:

        Returns:

        """
        if alpha == 1:
            from .coledavidson import ColeDavidson
            return ColeDavidson.distribution(tau, tau0, gamma)

        elif gamma == 1:
            from .colecole import ColeCole
            return ColeCole.distribution(tau, tau0, alpha)

        else:
            _y = tau/tau0
            om_y = (1 + 2*np.cos(np.pi*alpha)*(_y**alpha) + _y**(2*alpha))**(-gamma/2)
            theta_y = np.arctan2(np.sin(np.pi * alpha), _y**alpha + np.cos(np.pi*alpha))

            return np.sin(gamma*theta_y) * om_y * (_y**(alpha*gamma)) / np.pi

    @staticmethod
    @adjust_dims
    def correlation(t: ArrayLike, tau: ArrayLike, alpha: float, gamma: float) -> ArrayLike:
        r"""
        Correlation function

        .. math::
            C(t, \tau_\text{HN}, \alpha, \gamma) = 1 - \left(\frac{t}{\tau_\text{HN}}\right)^{\alpha\gamma}
            \text{E}_{\alpha,\alpha\gamma+1}^\gamma\left[-\left(\frac{t}{\tau_\text{HN}}\right)^\alpha\right]

        with three-parameter Mittag-Leffler function

        .. math::
            \text{E}_{a,b}^g (z) = \frac{1}{\Gamma(g)} \sum_{k=0}^\infty \frac{\Gamma(g+k) z^k}{k!\Gamma(ak+b)}

        Args:
            t (array_like): Time axis in s.
            tau (array_like): Correlation times :math:`\tau_\text{HN}` in s.
            alpha (float): Cole-Cole shape parameter.
            gamma (float): Cole-Davidson shape parameter.

        References:
            R. Garrappa: Models of dielectric relaxation based on completely monotone functions.
            Frac. Calc Appl. Anal. 19 (2016) https://arxiv.org/abs/1611.04028

        """
        return 1 - (t/tau)**(alpha*gamma) * mlf(-(t/tau)**alpha, alpha, alpha*gamma+1, gamma)

    @staticmethod
    @adjust_dims
    def susceptibility(omega: ArrayLike, tau: ArrayLike, alpha: float, gamma: float) -> ArrayLike:
        r"""
        Complex susceptibility

        .. math::
            \chi(\omega, \tau, \alpha, \gamma) = \frac{1}{[1-(i\omega\tau)^\alpha]^\gamma}

        Args:
            omega (array_like): Frequency axis in 1/s (not Hz).
            tau (array_like): Correlation times in s.
            alpha (float): Shape parameter.
            gamma (float): Even more shape parameter.

        """
        return np.conjugate(1/(1 + (1j*omega*tau)**alpha)**gamma)

    @staticmethod
    @adjust_dims
    def specdens(omega: ArrayLike, tau: ArrayLike, alpha: float, gamma: float) ->  ArrayLike:
        r"""
        Spectral  density :math:`J(\omega)`

        .. math::
            J(\omega) = \frac{\sin\left\lbrace \gamma\arctan\left[ \frac{(\omega\tau)^\alpha\sin(\alpha\pi/2)}{1+(\omega\tau)^\alpha\cos(\alpha\pi/2)} \right] \right\rbrace}
            {\omega \left[ 1+(\omega\tau)^\alpha \cos(\alpha\pi/2) + (\omega\tau)^{2\alpha} \right]^{\gamma/2}}

        Args:
            omega (array_like): Frequency axis in 1/s (not Hz).
            tau (array_like): Correlation times in s.
            alpha (float): Cole-Cole shape parameter.
            gamma (float): Cole-Davidson shape parameter.

        Returns:

        """
        omtau = (omega*tau)**alpha

        zaehler = np.sin(gamma * np.arctan2(omtau*np.sin(0.5*alpha*np.pi), 1 + omtau*np.cos(0.5*alpha*np.pi)))
        nenner = (1 + 2*omtau * np.cos(0.5*alpha*np.pi) + omtau**2)**(0.5*gamma)

        result = (1 / omega) * (zaehler/nenner)

        return result

    @staticmethod
    def mean(tau: ArrayLike, alpha: float, gamma: float) -> ArrayLike:
        r"""

        .. math::
            \langle \tau \rangle \approx \tau_\text{HN} \alpha \gamma

        This is an approximation given in `[1]`_

        Args:
            tau (array_like): Function parameter in s.
            alpha (float): Shape parameter.
            gamma (float): Even more shape parameter.

        Reference:
            _`[1]` Bauer, Th., Köhler, M., Lunkenheimer, P., Angell, C.A.:
            Relaxation dynamics and ionic conductivity in a fragile plastic crystal.
            J. Chem. Phys. 133, 144509 (2010). https://doi.org/10.1063/1.3487521


        """
        return tau * alpha*gamma

    @staticmethod
    def logmean(tau: ArrayLike, alpha: float, gamma: float) -> ArrayLike:
        r"""
        Calculate mean logarithm of tau

        .. math::
            \langle \ln \tau \rangle = \ln \tau + \frac{\Psi(\gamma) + \text{Eu}}{\alpha}

        Args:
            tau:
            alpha (float):
            gamma (float):

        """
        return np.log(tau) + (psi(gamma)+Eu)/alpha

    @staticmethod
    def max(tau: ArrayLike, alpha: float, gamma: float) -> ArrayLike:
        return tau*(np.sin(0.5*np.pi*alpha*gamma/(gamma+1)) / np.sin(0.5*np.pi*alpha/(gamma+1)))**(1/alpha)
