import numpy as np
from scipy.interpolate import interp1d
from scipy.special import gamma

from .base import Distribution
from ..lib.decorator import adjust_dims
from ..math.kww import kww_cos, kww_sin
from ..utils.constants import Eu


class KWW(Distribution):
    name = 'KWW'
    parameter = [r'\beta']
    boounds = [(0, 1)]

    @staticmethod
    def distribution(taus, tau, beta):
        if not (0.1 <= beta <= 0.9):
            raise ValueError('KWW distribution is only defined between 0.1 and 0.9')
        b_supp = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
        B_supp = [0.145, 0.197, 0.243, 0.285, 0.382, 0.306, 0.360, 0.435, 0.700]
        C_supp = [0.89, 0.50, 0.35, 0.25, 0, 0.13, 0.22, 0.4015, 0.32]

        B = interp1d(b_supp, B_supp)(beta)
        C = interp1d(b_supp, C_supp)(beta)

        tt = tau/taus

        delta = beta * abs(beta-0.5) / (1-beta)
        if beta > 0.5:
            f = 1 + C * tt**delta
        else:
            f = 1 / (1 + C * tt**delta)
        ret_val = tau * B * np.exp(-(1-beta) * beta**(beta/(1-beta)) / tt**(beta/(1-beta))) / tt**((1-0.5*beta)/(1-beta))

        return ret_val * f / taus

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

    @staticmethod
    @adjust_dims
    def susceptibility(omega, tau, beta):
        return 1-omega*kww_sin(omega, tau, beta) + 1j*omega*kww_cos(omega, tau, beta)

    @staticmethod
    @adjust_dims
    def specdens(omega, tau, beta):
        return kww_cos(omega, tau, beta)

    @staticmethod
    def mean(tau, beta):
        return tau/beta * gamma(1 / beta)

    @staticmethod
    def logmean(tau, beta):
        return (1-1/beta) * Eu + np.log(tau)

    @staticmethod
    def max(tau, beta):
        return (1.7851 - 0.87052*beta - 0.028836*beta**2 + 0.11391*beta**3) * tau
