from __future__ import annotations

import numpy as np
from scipy.special import psi, gammaincc

from .base import Distribution
from ..lib.decorator import adjust_dims
from ..utils.constants import Eu


class ColeDavidson(Distribution):
    """
    Functions based on Cole-Davidson distribution
    """
    name = 'Cole-Davidson'
    parameter = [r'\gamma']
    bounds = [(0, 1)]

    @staticmethod
    def distribution(tau, tau0, gamma):
        """
        Distribution of correlation times

        .. math::
            G(\ln\\tau) =
               \\begin{cases}
                   \\frac{\sin(\pi\gamma)}{\pi} \\Big(\\frac{\\tau}{\\tau_0 - \\tau}\\Big)^\gamma & \\tau < \\tau_0 \\\\
                   0 & \\text{else}
               \end{cases}

        Args:
            tau:
            tau0:
            gamma:
        """

        ratio = tau0/tau

        ret_val = np.zeros_like(ratio)
        ret_val[ratio > 1] = np.sin(np.pi*gamma) / np.pi * (1 / (ratio[ratio>1] - 1))**gamma

        return ret_val

    @staticmethod
    @adjust_dims
    def susceptibility(omega: float | np.ndarray, tau: float | np.ndarray, gamma: float) -> float | np.ndarray:
        r"""
        Complex susceptibility

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

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

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

    @staticmethod
    @adjust_dims
    def specdens(omega, tau, gamma):
        """
        Spectral density

        .. math::
            J(\omega, \\tau, \\gamma) = \\frac{\sin[\\gamma\\arctan(\\omega\\tau)]}{\omega[1+(\omega\\tau)^2]^{\\gamma/2}}

        Args:
            omega:
            tau:
            gamma:
        """
        omtau = omega*tau

        ret_val = np.sin(gamma*np.arctan2(omtau, 1)) / (1 + omtau**2)**(gamma/2.) / omega
        ret_val[np.argwhere(omega == 0)] = gamma*tau

        return ret_val

    @staticmethod
    def mean(tau, gamma):
        """
        Mean tau

        .. math::
            \\langle \\tau \\rangle = \\gamma \\tau

        Args:
            tau:
            gamma:

        """

        return tau*gamma

    @staticmethod
    def logmean(tau, gamma):
        """
        Mean logarithm of tau

        .. math::
            \\langle \ln \\tau \\rangle = \ln\\tau + \psi(\gamma) + \mathrm{Eu}

        with the Euler's constant :math:`\mathrm{Eu} \\approx 0.57721567\ldots` and the digamma function

        .. math::
            \psi(x) = \\frac{d}{dx} \ln(\Gamma(x))

        Args:
            tau:
            gamma:

        References:
            R. Zorn: Logartihmic moments of relaxation time distribution. J. Chem. Phys. (2002). http://dx.doi.org/10.1063/1.1446035

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

    @staticmethod
    def max(tau, gamma):
        """
        Calculate :math:`\\tau_\\text{peak}`, i.e. the inverse of the maximum frequency of the imaginary part of the susceptibility

        .. math::
            \\tau_\\text{peak} = \\frac{\\tau}{\\tan[\\pi/(2\gamma+2)]}

        Args:
            tau:
            gamma:

        References:
            R. Dı́az-Calleja: Comment on the Maximum in the Loss Permittivity for the Havriliak-Negami Equation.
            Macromolecules (2000). https://doi.org/10.1021/ma991082i

        """

        return tau/np.tan(np.pi/(2*gamma+2))

    @staticmethod
    @adjust_dims
    def correlation(t, tau, gamma):
        r"""
        Correlation function

        .. math::
            C(t, \tau_0, \gamma) = \frac{\Gamma(\gamma, t/\tau_0)}{\Gamma(\gamma)}

        with incomplete Gamma function

        .. math::
            \Gamma(\gamma, t/\tau_0) = \frac{1}{\Gamma(\gamma)}\int_{t/\tau_0}^\infty x^{\gamma-1}\text{e}^{-x}\,\mathrm{d}x

        Args:
            t:
            tau:
            gamma:

        References:
            R. Hilfer: H-function representations for stretched exponential relaxation and non-Debye susceptibilities in glassy systems. Phys. Rev. E (2002) https://doi.org/10.1103/PhysRevE.65.061510

        """
        return gammaincc(gamma, t / tau)
