"""Functions related to math.

Fast exponential code from:
Schraudolph, 1999, A fast, compact approximation of the exponential function.
"""

import numpy as np
cimport numpy as np
from scipy.special import gammaln

from libc.math cimport cos, sin, acos, log, exp, sqrt, lgamma, nearbyint
from libc.math cimport M_PI

from cython_gsl cimport *
gsl_rng_env_setup()
cdef:
    const gsl_rng_type * T
    gsl_rng * r
T = gsl_rng_default
r = gsl_rng_alloc (T)

cdef:
    double TWOPI = 2 * M_PI

ctypedef struct TwoInts:
    int j, i

cdef:
    double log_2 = 0.69314718055994529
    double exp_A = 1048576 / log_2
    int exp_C = 45799#60801

    union Eco:
        double d
        TwoInts n

    Eco eco


cdef double exp_macro(double y):

    eco.n.i = int(exp_A * y + (1072693248 - exp_C))

    return eco.d



cpdef double exp_fast(double x):

    return exp_macro(x)

cpdef double exp_fast2(double x):
    """Slightly slower version, but more accurate."""

    cdef double y = x / 2.0

    return exp_macro(y) / exp_macro(0.0 - y)


def log_sum_exp(x, axis=0):
    """Returns the logarithm of the sum of exponentials.

    Args:
        x: array
    """
    xmax = x.max(axis)

    return np.log(np.exp(x - xmax).sum(axis)) + xmax

cdef double log_sum_exp_c(double[::1] x):

    cdef:
        int i
        double x_max, sum_exp

    x_max = x[0]
    for i in range(x.shape[0]):
        if x[i] > x_max:
            x_max = x[i]
    sum_exp = 0.0
    for i in range(x.shape[0]):
        sum_exp += exp(x[i] - x_max)

    return log(sum_exp) + x_max



cpdef double vonmisesvariate(double mu, double kappa):
    """Circular data distribution.

    mu is the mean angle, expressed in radians between 0 and 2*pi, and
    kappa is the concentration parameter, which must be greater than or
    equal to zero.  If kappa is equal to zero, this distribution reduces
    to a uniform random angle over the range 0 to 2*pi.

    Note: Code taken from the python module random.py, and converted
    to cython.
    """
    # mu:    mean angle (in radians between 0 and 2*pi)
    # kappa: concentration parameter kappa (>= 0)
    # if kappa = 0 generate uniform random angle

    # Based upon an algorithm published in: Fisher, N.I.,
    # "Statistical Analysis of Circular Data", Cambridge
    # University Press, 1993.

    # Thanks to Magnus Kessler for a correction to the
    # implementation of step 4.

    cdef:
        double s, r1, u1, z, d, u2, q, f, u3, theta

    if kappa <= 1e-6:
        return TWOPI * gsl_ran_flat(r, 0, 1)

    s = 0.5 / kappa
    r1 = s + sqrt(1.0 + s * s)

    while 1:
        u1 = gsl_ran_flat(r, 0, 1)
        z = cos(M_PI * u1)

        d = z / (r1 + z)
        u2 = gsl_ran_flat(r, 0, 1)
        if u2 < 1.0 - d * d or u2 <= (1.0 - d) * exp(d):
            break

    q = 1.0 / r1
    f = (q + z) / (1.0 + q * z)
    u3 = gsl_ran_flat(r, 0, 1)
    if u3 > 0.5:
        theta = (mu + acos(f)) % TWOPI
    else:
        theta = (mu - acos(f)) % TWOPI

    return theta

cpdef int gaussian_poisson_brute_force(
    double mean, double variance, double rate, int max_k=1000):
    """Sample from product of Gaussian and Poisson distributions."""

    cdef:
        int i, index
        long[:] k
        double[::1] logf, f
        double std, c1, log_rate, cumsum, p

    k = np.arange(max_k)
    N = len(k)
    std = sqrt(variance)
    logf = np.empty((N, ), dtype='double')
    f = np.empty((N, ), dtype='double')
    c1 = -log(std)
    log_rate = log(rate)
    for i in range(N):
        logf[i] = c1 - (k[i] - mean)**2 / (2 * variance)
        logf[i] += k[i] * log_rate - gammaln(k[i] + 1)

    c2 = log_sum_exp_c(logf)
    for i in range(N):
        f[i] = exp(logf[i] - c2)

    # f is now the normalised, discrete distribution.
    # Sample from f.
    cumsum = 0.0
    ## p = gsl_ran_flat(r, 0, 1) # gives the same value every time
    p = np.random.rand()
    index = 0
    while index < N and cumsum < p:
        cumsum += f[index]
        index += 1

    return index - 1


