from functools import partial

import jax
from jax import lax
from jax import numpy as np
from jax import random
from jax.scipy.special import xlogy

from mcx.distributions import constraints
from mcx.distributions.distribution import Distribution
from mcx.distributions.shapes import broadcast_batch_shape


class Poisson(Distribution):
    parameters = {"lambda": constraints.positive}
    support = constraints.positive_integer

    def __init__(self, lmbda):
        self.event_shape = ()
        self.batch_shape = broadcast_batch_shape(np.shape(lmbda))
        self.lmbda = lmbda

    def sample(self, rng_key, sample_shape):
        shape = sample_shape + self.batch_shape + self.event_shape
        return _random_poisson(rng_key, self.lmbda, shape)

    @constraints.limit_to_support
    def logpdf(self, x):
        x = x * 1.0
        return lax.add(
            xlogy(x, lax.log(self.lmbda)),
            lax.add(lax.neg(self.lmbda), lax.neg(lax.lgamma(x))),
        )


@partial(jax.jit, static_argnums=(2,))
def _random_poisson(rng_key, lmbda, shape):
    """
    References
    ----------
    .. [1] Knuth, Donald E. Art of computer programming, volume 2:
           Seminumerical algorithms. Addison-Wesley Professional, 2014 (p 137).
    """
    L = lax.exp(lax.neg(lmbda))
    k = np.zeros(shape=shape)
    p = np.ones(shape=shape)

    is_done = p < L
    while not is_done.all():
        _, rng_key = random.split(rng_key)
        u = random.uniform(rng_key, shape=shape)
        p = np.where(is_done, p, u * p)
        k = np.where(is_done, k, k + 1)
        is_done = p < L

    return k
