import numbers

import numpy as np

from .base import Distribution
from ..lib.decorator import adjust_dims
from ..lib.utils import ArrayLike


class Debye(Distribution):
    name = 'Debye'

    @staticmethod
    def distribution(tau: ArrayLike, tau0: float) -> ArrayLike:
        if isinstance(tau, numbers.Number):
            return 1 if tau == tau0 else 0

        ret_val = np.zeros_like(tau)
        ret_val[tau == tau0] = 1.
        return ret_val

    @staticmethod
    @adjust_dims
    def correlation(t, tau, *args):
        return np.exp(-t/tau)

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

        .. math::
            \chi(\omega, \tau) = \frac{1}{1-i\omega\tau}

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

        """
        return 1/(1 - 1j*omega*tau)

    @staticmethod
    @adjust_dims
    def specdens(omega: ArrayLike, tau: ArrayLike) ->  ArrayLike:
        return tau / (1 + (omega*tau)**2)
