import numpy as np
import scipy
import scipy.special

__all__ = ['kww_cos', 'kww_sin']

__kww_delta = 2.2e-16
__kww_eps = 5.5e-20
__max_terms = 150
__max_iter_int = 10


def _kww_low(w, beta, kappa):
    k = np.atleast_2d(np.arange(kappa, 2*__max_terms+kappa, step=2))

    sign = np.empty(__max_terms)
    sign[::2] = 1
    sign[1::2] = -1

    ln_w = np.log(w[:, np.newaxis])
    ln_a_k = scipy.special.gammaln((k+1) / beta) - scipy.special.gammaln(k+1) + k*ln_w
    with np.errstate(over='ignore'):
        a_k = np.exp(ln_a_k)

    y_n = (sign * a_k).cumsum(axis=1)
    y_n = np.nan_to_num(y_n)

    z_n = a_k.cumsum(axis=1)
    z_n = np.nan_to_num(z_n)

    mask = (__kww_eps * z_n[:, :-1] + a_k[:, 1:] <= __kww_delta * y_n[:, : -1]).argmax(axis=1) - 1

    y = y_n[list(range(0, len(w))), mask]
    y[mask == -1] = np.nan

    return y / beta


def kwwc_low(w, beta):
    return _kww_low(w, beta, 0)


def kwws_low(w, beta):
    return _kww_low(w, beta, 1)


def _kww_high(w, beta, kappa):
    if beta < 0.1 or beta > 2.0:
        raise ValueError("invalid call to kww_high: beta out of range")

    ln_omega = np.log(w[:, np.newaxis])
    k = np.atleast_2d(np.arange(kappa, __max_terms+kappa))

    if kappa:
        c_k = np.sin(k*beta*np.pi/2)
    else:
        c_k = np.cos(k*beta*np.pi/2)

    ln_u_k = scipy.special.gammaln(k*beta+1) - scipy.special.gammaln(k+1) - (k*beta+1) * ln_omega
    u_k = np.exp(ln_u_k)

    a_k = c_k*u_k

    if beta > 1:
        r_k = np.sin(np.pi/2/beta)**(-1-k*beta)
    else:
        r_k = np.ones_like(k)
    sign = 1-2*((k-kappa) % 2)

    y_n = (sign*a_k).cumsum(axis=1)
    y_n = np.nan_to_num(y_n)
    z_n = np.abs(a_k).cumsum(axis=1)
    z_n = np.nan_to_num(z_n)

    mask = (__kww_eps*z_n[:, :-1] + u_k[:, 1:]*r_k[:, 1:] <= __kww_delta*np.abs(y_n[:, : -1])).argmax(axis=1)-1
    y = y_n[list(range(0, len(w))), mask]
    y[mask == -1] = np.nan

    return y


def kwws_high(w, beta):
    return _kww_high(w, beta, 0)


def kwwc_high(w, beta):
    return _kww_high(w, beta, 1)


def _kww_mid(w, beta, kappa):
    if beta < 0.1 or beta > 2.0:
        raise ValueError("invalid call to kww_mid: beta out of range")
    elif np.any(w < 0):
        raise ValueError("invalid call to kww_mid: w out of range")

    if kappa and beta == 2:
        return np.sqrt(np.pi) / 2 * np.exp(-np.sqrt(w) / 4)

    if beta < 0.15:
        p = 1.8
        q = 0.2
    elif beta < 0.25:
        p = 1.6
        q = 0.4
    elif beta < 1:
        p = 1.4
        q = 0.6
    elif beta < 1.75:
        p = 1.0
        q = 0.2
    elif beta < 1.95:
        p = 0.75
        q = 0.2
    else:
        p = 0.15
        q = 0.4

    n = 10
    y = np.zeros(w.shape, dtype=float)
    y_min = 2e-20
    for _ in range(__max_iter_int):
        y_last = y
        h = np.log(np.log((np.pi+1)*10*n/__kww_delta/y_min) / p) / n

        k = np.arange(-n, n+1, dtype=float)
        sign = 1 - 2*(k % 2)
        hk = h*k

        eta_k = 2*p*np.sinh(hk) + 2*q*hk
        eta_k_prime = 2*p*np.cosh(hk) + 2*q

        exp_eta = np.exp(-eta_k)
        expm1_eta = -np.expm1(-eta_k)
        expm1_eta[n] = 1.
        a_k = np.pi * k / expm1_eta
        a_k[n] = np.pi/h/eta_k_prime[n]

        b_k_2 = sign * np.sin(np.pi*k*exp_eta / expm1_eta)
        mask_exp = exp_eta > 1
        if kappa:
            b_k_2[mask_exp] = np.cos(np.pi*k[mask_exp]/expm1_eta[mask_exp])
        else:
            b_k_2[mask_exp] = np.sin(np.pi*k[mask_exp]/expm1_eta[mask_exp])

        phi_k_prime = 1. / expm1_eta - hk*exp_eta*eta_k_prime / expm1_eta / expm1_eta
        b_k = phi_k_prime * b_k_2
        b_k[n] = np.sin(a_k[n]) / 2

        g_k = np.exp(-(a_k / w[:, np.newaxis])**beta)
        diff_mode = kappa and (beta > 1.75)
        if diff_mode:
            g_k -= np.exp(-(a_k / w[:, np.newaxis])**2)

        y_k = b_k*g_k

        y = y_k.sum(axis=1)
        z = np.abs(y_k).sum(axis=1)

        if diff_mode:
            y += w/np.sqrt(np.pi)/2 * np.exp(-w**2 / 4)

        if all(np.abs(y-y_last) + __kww_eps*z < __kww_delta * np.abs(y)):
            return y * np.pi / w

        n *= 2

    return y * np.pi / w


def kwwc_mid(w, beta):
    return _kww_mid(w, beta, True)


def kwws_mid(w, beta):
    return _kww_mid(w, beta, False)


def _kwwc_lim_low(b):
    if b > 1.024:
        return -0.8774954 * b + 3.5873 * b**2 - 2.083 * b**3 + 0.3796 * b**4
    else:
        return np.exp(-0.02194 / b**2 - 4.130 / b + 2.966189 + 0.030104 * b + 1.062 * b**2)


def _kwws_lim_low(b):
    if b > 1.024:
        return -1.68725 * b + 4.8108 * b**2 - 2.561 * b**3 + 0.442 * b**4
    else:
        return np.exp(-0.03208 / b**2 - 4.314 / b + 3.516200 - 0.50287 * b + 1.240 * b**2)


def _kwwc_lim_hig(b):
    if b < 0.82:
        return np.exp(0.006923209 / b ** 2 - 1.321692 / b - 1.44582 + 2.516339 * b + 0.2973773 * b ** 2)
    else:
        return np.exp(-0.746496154631 + 6.057558 * (b-0.82) - 3.41052 * (b-0.82)**2 + 0.7932314 * (b-0.82)**3)


def _kwws_lim_hig(b):
    if b < 0.82:
        return np.exp(0.07847516 / b**2 - 2.585876 / b + 4.999414 - 8.460926 * b + 6.289183 * b**2)
    else:
        return np.exp(-0.962597724393 + 5.818057 * (b - 0.82) - 3.026212 * (b-0.82)**2 + 0.5485754 * (b-0.82)**3)


def kww_cos(w, tau, beta):
    beta = np.asanyarray(beta)
    if beta.ndim == 0:
        return _kww_cos_single(w, tau, beta)

    else:
        tau = np.asanyarray(tau)
        if tau.size != beta.size:
            raise ValueError('KWW mismatched number of beta and tau')
        else:
            beta.reshape(tau.shape)
            ret_val = np.zeros((w.size, tau.size))
            for i, (tau_i, beta_i) in enumerate(zip(tau.ravel(), beta.ravel())):
                ret_val[:, i] = _kww_cos_single(w, tau_i, beta_i)

            return ret_val.reshape(-1, *tau.shape)


def _kww_cos_single(w, tau, beta):
    """

    Args:
        w:
        tau:
        beta:

    Returns:

    """
    # check input data
    if not (0.1 <= beta <= 2.):
        raise ValueError('KWW beta is not inside range 0.1-2')

    if beta == 1:
        return tau/(1 + (w*tau)**2)

    if beta == 2:
        return tau*np.sqrt(np.pi)/2 * np.exp(-(w*tau)**2 / 4)

    _w = np.atleast_1d(np.abs(w*tau))

    low_approx = _w < _kwwc_lim_low(beta)
    high_approx = _w > _kwwc_lim_hig(beta)
    mid_integration = ~(low_approx | high_approx)

    ret_val = np.zeros_like(_w, dtype=float)
    ret_val[high_approx] = kwwc_high(_w[high_approx], beta)
    ret_val[low_approx] = kwwc_low(_w[low_approx], beta)
    ret_val[mid_integration] = kwwc_mid(_w[mid_integration], beta)
    ret_val[np.isnan(ret_val)] = kwwc_mid(_w[np.isnan(ret_val)], beta)

    ret_val[np.where(_w == 0)] = scipy.special.gamma(1 / beta) / beta

    if np.isscalar(_w):
        return tau*ret_val[0]

    return tau*ret_val


# \int_0^\infty dt sin(w*t) exp(-(t/tau)^beta)
def kww_sin(w, tau, beta):
    beta = np.asanyarray(beta)
    if beta.ndim == 0:
        return _kww_sin_single(w, tau, beta)

    else:
        tau = np.asanyarray(tau)
        if tau.size != beta.size:
            raise ValueError('KWW mismatched number of beta and tau')
        else:
            beta.reshape(tau.shape)
            ret_val = np.zeros((w.size, tau.size))
            for i, (tau_i, beta_i) in enumerate(zip(tau.ravel(), beta.ravel())):
                ret_val[:, i] = _kww_sin_single(w, tau_i, beta_i)

            return ret_val.reshape(-1, *tau.shape)


def _kww_sin_single(w, tau, beta):
    # check input data
    if not (0.1 <= beta <= 2.):
        raise ValueError('KWW beta is not inside range 0.1-2')

    if beta == 1:
        return w*tau**2/(1+(w*tau)**2)

    if beta == 2:
        return tau*scipy.special.dawsn(w*tau/2.0)

    _w = np.atleast_1d(np.abs(w*tau))

    sign = np.sign(_w)

    low_approx = _w < _kwws_lim_low(beta)
    high_approx = _w > _kwws_lim_hig(beta)
    mid_integration = ~(low_approx | high_approx)

    ret_val = np.zeros_like(_w, dtype=float)
    ret_val[high_approx] = kwws_high(_w[high_approx], beta)
    ret_val[low_approx] = kwws_low(_w[low_approx], beta)
    ret_val[mid_integration] = kwws_mid(_w[mid_integration], beta)

    ret_val[np.isnan(ret_val)] = kwws_mid(_w[np.isnan(ret_val)], beta)

    ret_val[_w == 0] = 0.0

    if np.isscalar(_w):
        return sign*tau*ret_val[0]

    return sign*ret_val*tau
