from __future__ import annotations

import numpy as np

from .base import Distribution
from ..lib.decorator import adjust_dims
from ..math.mittagleffler import mlf


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

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

        Args:
            tau (array_like):
            tau0 (array_like):
            alpha (float):
        """
        z = np.log(tau/tau0)
        return np.sin(np.pi*alpha)/(np.cosh(z*alpha) + np.cos(alpha*np.pi))/(2*np.pi)

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

        Args:
            omega (array_like): Frequency axis in 1/s (not Hz).
            tau (array_like): Correlation times :math:`\tau_\text{CC}` in s.
            alpha (float): Shape parameter.

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

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

        Args:
            omega (array_like):
            tau (array_like):
            alpha (float):
        """
        omtau = (omega*tau)**alpha
        return np.sin(alpha*np.pi/2) * omtau / (1 + omtau**2 + 2*np.cos(alpha*np.pi/2)*omtau) / omega

    @staticmethod
    @adjust_dims
    def correlation(t, tau0, alpha):
        """
        Correlation function :math:`C(t)`

        with Mittag-Leffler function :math:`E_a(x)`

        Args:
            t (array_like):
            tau0 (array_like):
            alpha (float):
        """
        return mlf(-(t/tau0)**alpha, alpha)
