"""
Implement the random and np.random module functions.
"""


import math
import random

import numpy as np

from llvmlite import ir

from numba.core.extending import overload, register_jitable
from numba.core.imputils import (Registry, impl_ret_untracked,
                                    impl_ret_new_ref)
from numba.core.typing import signature
from numba.core import types, utils, cgutils
from numba.np import arrayobj
from numba.core.overload_glue import glue_lowering
from numba.core.errors import NumbaTypeError


POST_PY38 = utils.PYVERSION >= (3, 8)


registry = Registry('randomimpl')
lower = registry.lower

int32_t = ir.IntType(32)
int64_t = ir.IntType(64)
def const_int(x):
    return ir.Constant(int32_t, x)
double = ir.DoubleType()

N = 624
N_const = ir.Constant(int32_t, N)


# This is the same struct as rnd_state_t in _random.c.
rnd_state_t = ir.LiteralStructType([
    # index
    int32_t,
    # mt[N]
    ir.ArrayType(int32_t, N),
    # has_gauss
    int32_t,
    # gauss
    double,
    # is_initialized
    int32_t,
    ])
rnd_state_ptr_t = ir.PointerType(rnd_state_t)

def get_state_ptr(context, builder, name):
    """
    Get a pointer to the given thread-local random state
    (depending on *name*: "py" or "np").
    If the state isn't initialized, it is lazily initialized with
    system entropy.
    """
    assert name in ('py', 'np', 'internal')
    func_name = "numba_get_%s_random_state" % name
    fnty = ir.FunctionType(rnd_state_ptr_t, ())
    fn = cgutils.get_or_insert_function(builder.module, fnty, func_name)
    # These two attributes allow LLVM to hoist the function call
    # outside of loops.
    fn.attributes.add('readnone')
    fn.attributes.add('nounwind')
    return builder.call(fn, ())

def get_py_state_ptr(context, builder):
    """
    Get a pointer to the thread-local Python random state.
    """
    return get_state_ptr(context, builder, 'py')

def get_np_state_ptr(context, builder):
    """
    Get a pointer to the thread-local Numpy random state.
    """
    return get_state_ptr(context, builder, 'np')

def get_internal_state_ptr(context, builder):
    """
    Get a pointer to the thread-local internal random state.
    """
    return get_state_ptr(context, builder, 'internal')

# Accessors
def get_index_ptr(builder, state_ptr):
    return cgutils.gep_inbounds(builder, state_ptr, 0, 0)

def get_array_ptr(builder, state_ptr):
    return cgutils.gep_inbounds(builder, state_ptr, 0, 1)

def get_has_gauss_ptr(builder, state_ptr):
    return cgutils.gep_inbounds(builder, state_ptr, 0, 2)

def get_gauss_ptr(builder, state_ptr):
    return cgutils.gep_inbounds(builder, state_ptr, 0, 3)

def get_rnd_shuffle(builder):
    """
    Get the internal function to shuffle the MT taste.
    """
    fnty = ir.FunctionType(ir.VoidType(), (rnd_state_ptr_t,))
    fn = cgutils.get_or_insert_function(builder.function.module, fnty,
                                        "numba_rnd_shuffle")
    fn.args[0].add_attribute("nocapture")
    return fn


def get_next_int32(context, builder, state_ptr):
    """
    Get the next int32 generated by the PRNG at *state_ptr*.
    """
    idxptr = get_index_ptr(builder, state_ptr)
    idx = builder.load(idxptr)
    need_reshuffle = builder.icmp_unsigned('>=', idx, N_const)
    with cgutils.if_unlikely(builder, need_reshuffle):
        fn = get_rnd_shuffle(builder)
        builder.call(fn, (state_ptr,))
        builder.store(const_int(0), idxptr)
    idx = builder.load(idxptr)
    array_ptr = get_array_ptr(builder, state_ptr)
    y = builder.load(cgutils.gep_inbounds(builder, array_ptr, 0, idx))
    idx = builder.add(idx, const_int(1))
    builder.store(idx, idxptr)
    # Tempering
    y = builder.xor(y, builder.lshr(y, const_int(11)))
    y = builder.xor(y, builder.and_(builder.shl(y, const_int(7)),
                                    const_int(0x9d2c5680)))
    y = builder.xor(y, builder.and_(builder.shl(y, const_int(15)),
                                    const_int(0xefc60000)))
    y = builder.xor(y, builder.lshr(y, const_int(18)))
    return y

def get_next_double(context, builder, state_ptr):
    """
    Get the next double generated by the PRNG at *state_ptr*.
    """
    # a = rk_random(state) >> 5, b = rk_random(state) >> 6;
    a = builder.lshr(get_next_int32(context, builder, state_ptr), const_int(5))
    b = builder.lshr(get_next_int32(context, builder, state_ptr), const_int(6))

    # return (a * 67108864.0 + b) / 9007199254740992.0;
    a = builder.uitofp(a, double)
    b = builder.uitofp(b, double)
    return builder.fdiv(
        builder.fadd(b, builder.fmul(a, ir.Constant(double, 67108864.0))),
        ir.Constant(double, 9007199254740992.0))

def get_next_int(context, builder, state_ptr, nbits, is_numpy):
    """
    Get the next integer with width *nbits*.
    """
    c32 = ir.Constant(nbits.type, 32)
    def get_shifted_int(nbits):
        shift = builder.sub(c32, nbits)
        y = get_next_int32(context, builder, state_ptr)
        if is_numpy:
            # Use the last N bits, to match np.random
            mask = builder.not_(ir.Constant(y.type, 0))
            mask = builder.lshr(mask, builder.zext(shift, y.type))
            return builder.and_(y, mask)
        else:
            # Use the first N bits, to match CPython random
            return builder.lshr(y, builder.zext(shift, y.type))

    ret = cgutils.alloca_once_value(builder, ir.Constant(int64_t, 0))

    is_32b = builder.icmp_unsigned('<=', nbits, c32)
    with builder.if_else(is_32b) as (ifsmall, iflarge):
        with ifsmall:
            low = get_shifted_int(nbits)
            builder.store(builder.zext(low, int64_t), ret)
        with iflarge:
            # XXX This assumes nbits <= 64
            if is_numpy:
                # Get the high bits first to match np.random
                high = get_shifted_int(builder.sub(nbits, c32))
            low = get_next_int32(context, builder, state_ptr)
            if not is_numpy:
                # Get the high bits second to match CPython random
                high = get_shifted_int(builder.sub(nbits, c32))
            total = builder.add(
                builder.zext(low, int64_t),
                builder.shl(builder.zext(high, int64_t), ir.Constant(int64_t, 32)))
            builder.store(total, ret)

    return builder.load(ret)


def _fill_defaults(context, builder, sig, args, defaults):
    """
    Assuming a homogeneous signature (same type for result and all arguments),
    fill in the *defaults* if missing from the arguments.
    """
    ty = sig.return_type
    llty = context.get_data_type(ty)
    args = tuple(args) + tuple(ir.Constant(llty, d) for d in defaults[len(args):])
    sig = signature(*(ty,) * (len(args) + 1))
    return sig, args


@glue_lowering("random.seed", types.uint32)
def seed_impl(context, builder, sig, args):
    res =  _seed_impl(context, builder, sig, args, get_state_ptr(context,
                                                                 builder, "py"))
    return impl_ret_untracked(context, builder, sig.return_type, res)

@glue_lowering("np.random.seed", types.uint32)
def seed_impl(context, builder, sig, args):
    res = _seed_impl(context, builder, sig, args, get_state_ptr(context,
                                                               builder, "np"))
    return impl_ret_untracked(context, builder, sig.return_type, res)

def _seed_impl(context, builder, sig, args, state_ptr):
    seed_value, = args
    fnty = ir.FunctionType(ir.VoidType(), (rnd_state_ptr_t, int32_t))
    fn = cgutils.get_or_insert_function(builder.function.module, fnty,
                                        "numba_rnd_init")
    builder.call(fn, (state_ptr, seed_value))
    return context.get_constant(types.none, None)

@glue_lowering("random.random")
def random_impl(context, builder, sig, args):
    state_ptr = get_state_ptr(context, builder, "py")
    res = get_next_double(context, builder, state_ptr)
    return impl_ret_untracked(context, builder, sig.return_type, res)

@glue_lowering("np.random.random")
@glue_lowering("np.random.random_sample")
@glue_lowering("np.random.sample")
@glue_lowering("np.random.ranf")
def random_impl(context, builder, sig, args):
    state_ptr = get_state_ptr(context, builder, "np")
    res = get_next_double(context, builder, state_ptr)
    return impl_ret_untracked(context, builder, sig.return_type, res)


@glue_lowering("random.gauss", types.Float, types.Float)
@glue_lowering("random.normalvariate", types.Float, types.Float)
def gauss_impl(context, builder, sig, args):
    res = _gauss_impl(context, builder, sig, args, "py")
    return impl_ret_untracked(context, builder, sig.return_type, res)


@glue_lowering("np.random.standard_normal")
@glue_lowering("np.random.normal")
@glue_lowering("np.random.normal", types.Float)
@glue_lowering("np.random.normal", types.Float, types.Float)
def np_gauss_impl(context, builder, sig, args):
    sig, args = _fill_defaults(context, builder, sig, args, (0.0, 1.0))
    res = _gauss_impl(context, builder, sig, args, "np")
    return impl_ret_untracked(context, builder, sig.return_type, res)


def _gauss_pair_impl(_random):
    def compute_gauss_pair():
        """
        Compute a pair of numbers on the normal distribution.
        """
        while True:
            x1 = 2.0 * _random() - 1.0
            x2 = 2.0 * _random() - 1.0
            r2 = x1*x1 + x2*x2
            if r2 < 1.0 and r2 != 0.0:
                break

        # Box-Muller transform
        f = math.sqrt(-2.0 * math.log(r2) / r2)
        return f * x1, f * x2
    return compute_gauss_pair

def _gauss_impl(context, builder, sig, args, state):
    # The type for all computations (either float or double)
    ty = sig.return_type
    llty = context.get_data_type(ty)

    state_ptr = get_state_ptr(context, builder, state)
    _random = {"py": random.random,
               "np": np.random.random}[state]

    ret = cgutils.alloca_once(builder, llty, name="result")

    gauss_ptr = get_gauss_ptr(builder, state_ptr)
    has_gauss_ptr = get_has_gauss_ptr(builder, state_ptr)
    has_gauss = cgutils.is_true(builder, builder.load(has_gauss_ptr))
    with builder.if_else(has_gauss) as (then, otherwise):
        with then:
            # if has_gauss: return it
            builder.store(builder.load(gauss_ptr), ret)
            builder.store(const_int(0), has_gauss_ptr)
        with otherwise:
            # if not has_gauss: compute a pair of numbers using the Box-Muller
            # transform; keep one and return the other
            pair = context.compile_internal(builder,
                                            _gauss_pair_impl(_random),
                                            signature(types.UniTuple(ty, 2)),
                                            ())

            first, second = cgutils.unpack_tuple(builder, pair, 2)
            builder.store(first, gauss_ptr)
            builder.store(second, ret)
            builder.store(const_int(1), has_gauss_ptr)

    mu, sigma = args
    return builder.fadd(mu,
                        builder.fmul(sigma, builder.load(ret)))

@glue_lowering("random.getrandbits", types.Integer)
def getrandbits_impl(context, builder, sig, args):
    nbits, = args
    too_large = builder.icmp_unsigned(">=", nbits, const_int(65))
    too_small = builder.icmp_unsigned("==", nbits, const_int(0))
    with cgutils.if_unlikely(builder, builder.or_(too_large, too_small)):
        msg = "getrandbits() limited to 64 bits"
        context.call_conv.return_user_exc(builder, OverflowError, (msg,))
    state_ptr = get_state_ptr(context, builder, "py")
    res = get_next_int(context, builder, state_ptr, nbits, False)
    return impl_ret_untracked(context, builder, sig.return_type, res)


def _randrange_impl(context, builder, start, stop, step, state):
    state_ptr = get_state_ptr(context, builder, state)
    ty = stop.type
    zero = ir.Constant(ty, 0)
    one = ir.Constant(ty, 1)
    nptr = cgutils.alloca_once(builder, ty, name="n")
    # n = stop - start
    builder.store(builder.sub(stop, start), nptr)

    with builder.if_then(builder.icmp_signed('<', step, zero)):
        # n = (n + step + 1) // step
        w = builder.add(builder.add(builder.load(nptr), step), one)
        n = builder.sdiv(w, step)
        builder.store(n, nptr)
    with builder.if_then(builder.icmp_signed('>', step, one)):
        # n = (n + step - 1) // step
        w = builder.sub(builder.add(builder.load(nptr), step), one)
        n = builder.sdiv(w, step)
        builder.store(n, nptr)

    n = builder.load(nptr)
    with cgutils.if_unlikely(builder, builder.icmp_signed('<=', n, zero)):
        # n <= 0
        msg = "empty range for randrange()"
        context.call_conv.return_user_exc(builder, ValueError, (msg,))

    fnty = ir.FunctionType(ty, [ty, cgutils.true_bit.type])
    fn = cgutils.get_or_insert_function(builder.function.module, fnty,
                                        "llvm.ctlz.%s" % ty)
    # Since the upper bound is exclusive, we need to subtract one before
    # calculating the number of bits. This leads to a special case when
    # n == 1; there's only one possible result, so we don't need bits from
    # the PRNG. This case is handled separately towards the end of this
    # function. CPython's implementation is simpler and just runs another
    # iteration of the while loop when the resulting number is too large
    # instead of subtracting one, to avoid needing to handle a special
    # case. Thus, we only perform this subtraction for the NumPy case.
    nm1 = builder.sub(n, one) if state == "np" else n
    nbits = builder.trunc(builder.call(fn, [nm1, cgutils.true_bit]), int32_t)
    nbits = builder.sub(ir.Constant(int32_t, ty.width), nbits)

    rptr = cgutils.alloca_once(builder, ty, name="r")

    def get_num():
        bbwhile = builder.append_basic_block("while")
        bbend = builder.append_basic_block("while.end")
        builder.branch(bbwhile)

        builder.position_at_end(bbwhile)
        r = get_next_int(context, builder, state_ptr, nbits, state == "np")
        r = builder.trunc(r, ty)
        too_large = builder.icmp_signed('>=', r, n)
        builder.cbranch(too_large, bbwhile, bbend)

        builder.position_at_end(bbend)
        builder.store(r, rptr)

    if state == "np":
        # Handle n == 1 case, per previous comment.
        with builder.if_else(builder.icmp_signed('==', n, one)) as (is_one, is_not_one):
            with is_one:
                builder.store(zero, rptr)
            with is_not_one:
                get_num()
    else:
        get_num()

    return builder.add(start, builder.mul(builder.load(rptr), step))


@glue_lowering("random.randrange", types.Integer)
def randrange_impl_1(context, builder, sig, args):
    stop, = args
    start = ir.Constant(stop.type, 0)
    step = ir.Constant(stop.type, 1)
    res = _randrange_impl(context, builder, start, stop, step, "py")
    return impl_ret_untracked(context, builder, sig.return_type, res)

@glue_lowering("random.randrange", types.Integer, types.Integer)
def randrange_impl_2(context, builder, sig, args):
    start, stop = args
    step = ir.Constant(start.type, 1)
    res = _randrange_impl(context, builder, start, stop, step, "py")
    return impl_ret_untracked(context, builder, sig.return_type, res)

@glue_lowering("random.randrange", types.Integer,
               types.Integer, types.Integer)
def randrange_impl_3(context, builder, sig, args):
    start, stop, step = args
    res = _randrange_impl(context, builder, start, stop, step, "py")
    return impl_ret_untracked(context, builder, sig.return_type, res)

@glue_lowering("random.randint", types.Integer, types.Integer)
def randint_impl_1(context, builder, sig, args):
    start, stop = args
    step = ir.Constant(start.type, 1)
    stop = builder.add(stop, step)
    res = _randrange_impl(context, builder, start, stop, step, "py")
    return impl_ret_untracked(context, builder, sig.return_type, res)

@glue_lowering("np.random.randint", types.Integer)
def randint_impl_2(context, builder, sig, args):
    stop, = args
    start = ir.Constant(stop.type, 0)
    step = ir.Constant(stop.type, 1)
    res = _randrange_impl(context, builder, start, stop, step, "np")
    return impl_ret_untracked(context, builder, sig.return_type, res)

@glue_lowering("np.random.randint", types.Integer, types.Integer)
def randrange_impl_2(context, builder, sig, args):
    start, stop = args
    step = ir.Constant(start.type, 1)
    res = _randrange_impl(context, builder, start, stop, step, "np")
    return impl_ret_untracked(context, builder, sig.return_type, res)

@glue_lowering("random.uniform", types.Float, types.Float)
def uniform_impl(context, builder, sig, args):
    res = uniform_impl(context, builder, sig, args, "py")
    return impl_ret_untracked(context, builder, sig.return_type, res)

@glue_lowering("np.random.uniform", types.Float, types.Float)
def uniform_impl(context, builder, sig, args):
    res = uniform_impl(context, builder, sig, args, "np")
    return impl_ret_untracked(context, builder, sig.return_type, res)


def uniform_impl(context, builder, sig, args, state):
    state_ptr = get_state_ptr(context, builder, state)
    a, b = args
    width = builder.fsub(b, a)
    r = get_next_double(context, builder, state_ptr)
    return builder.fadd(a, builder.fmul(width, r))


@overload(random.triangular)
def triangular_impl_2(low, high):
    def _impl(low, high):
        u = random.random()
        c = 0.5
        if u > c:
            u = 1.0 - u
            low, high = high, low
        return low + (high - low) * math.sqrt(u * c)

    if isinstance(low, types.Number) and isinstance(high, types.Number):
        return _impl


@overload(random.triangular)
def triangular_impl_3(low, high, mode):
    if (isinstance(low, types.Number) and isinstance(high, types.Number) and
       isinstance(mode, types.Number)):
        def _impl(low, high, mode):
            if high == low:
                return low
            u = random.random()
            c = (mode - low) / (high - low)
            if u > c:
                u = 1.0 - u
                c = 1.0 - c
                low, high = high, low
            return low + (high - low) * math.sqrt(u * c)

        return _impl


@overload(np.random.triangular)
def triangular_impl_3(low, mode, high):
    if (isinstance(low, types.Number) and isinstance(mode, types.Number) and
            isinstance(high, types.Number)):
        def _impl(low, mode, high):
            if high == low:
                return low
            u = np.random.random()
            c = (mode - low) / (high - low)
            if u > c:
                u = 1.0 - u
                c = 1.0 - c
                low, high = high, low
            return low + (high - low) * math.sqrt(u * c)

        return _impl


@overload(np.random.triangular)
def triangular_impl(low, high, mode, size):
    if isinstance(size, types.NoneType):
        return lambda low, high, mode, size: np.random.triangular(low, high,
                                                                  mode)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                            isinstance(size.dtype,
                                                       types.Integer))):
        def _impl(low, high, mode, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.triangular(low, high, mode)
            return out
        return _impl


@overload(random.gammavariate)
def gammavariate_impl(alpha, beta):
    if isinstance(alpha, types.Number) and isinstance(beta, types.Number):
        return _gammavariate_impl(random.random)


@overload(np.random.standard_gamma)
@overload(np.random.gamma)
def gammavariate_impl(alpha):
    if isinstance(alpha, types.Number):
        return lambda alpha: np.random.gamma(alpha, 1.0)


@overload(np.random.gamma)
def gammavariate_impl(alpha, beta):
    if isinstance(alpha, types.Number) and isinstance(beta, types.Number):
        return _gammavariate_impl(np.random.random)


def _gammavariate_impl(_random):
    def _impl(alpha, beta):
        """Gamma distribution.  Taken from CPython.
        """
        SG_MAGICCONST = 1.0 + math.log(4.5)
        # alpha > 0, beta > 0, mean is alpha*beta, variance is alpha*beta**2

        # Warning: a few older sources define the gamma distribution in terms
        # of alpha > -1.0
        if alpha <= 0.0 or beta <= 0.0:
            raise ValueError('gammavariate: alpha and beta must be > 0.0')

        if alpha > 1.0:
            # Uses R.C.H. Cheng, "The generation of Gamma
            # variables with non-integral shape parameters",
            # Applied Statistics, (1977), 26, No. 1, p71-74
            ainv = math.sqrt(2.0 * alpha - 1.0)
            bbb = alpha - math.log(4.0)
            ccc = alpha + ainv

            while 1:
                u1 = _random()
                if not 1e-7 < u1 < .9999999:
                    continue
                u2 = 1.0 - _random()
                v = math.log(u1/(1.0-u1))/ainv
                x = alpha*math.exp(v)
                z = u1*u1*u2
                r = bbb+ccc*v-x
                if r + SG_MAGICCONST - 4.5*z >= 0.0 or r >= math.log(z):
                    return x * beta

        elif alpha == 1.0:
            # expovariate(1)

            if POST_PY38:
                # Adjust due to cpython
                # commit 63d152232e1742660f481c04a811f824b91f6790
                return -math.log(1.0 - _random()) * beta
            else:
                u = _random()
                while u <= 1e-7:
                    u = _random()
                return -math.log(u) * beta

        else:   # alpha is between 0 and 1 (exclusive)
            # Uses ALGORITHM GS of Statistical Computing - Kennedy & Gentle
            while 1:
                u = _random()
                b = (math.e + alpha)/math.e
                p = b*u
                if p <= 1.0:
                    x = p ** (1.0/alpha)
                else:
                    x = -math.log((b-p)/alpha)
                u1 = _random()
                if p > 1.0:
                    if u1 <= x ** (alpha - 1.0):
                        break
                elif u1 <= math.exp(-x):
                    break
            return x * beta
    return _impl


@overload(np.random.gamma)
def gamma_impl(alpha, beta, size):
    if isinstance(size, types.NoneType):
        return lambda alpha, beta, size: np.random.gamma(alpha, beta)
    if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                           isinstance(size.dtype,
                                                      types.Integer)):
        def _impl(alpha, beta, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.gamma(alpha, beta)
            return out
        return _impl


@overload(np.random.standard_gamma)
def standard_gamma_impl(alpha, size):
    if isinstance(size, types.NoneType):
        return lambda alpha, size: np.random.standard_gamma(alpha)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                            and isinstance(size.dtype,
                                                           types.Integer))):
        def _impl(alpha, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.standard_gamma(alpha)
            return out
        return _impl


@overload(random.betavariate)
def betavariate_impl(alpha, beta):
    if isinstance(alpha, types.Number) and isinstance(beta, types.Number):
        return _betavariate_impl(random.gammavariate)


@overload(np.random.beta)
def betavariate_impl(alpha, beta):
    if isinstance(alpha, types.Number) and isinstance(beta, types.Number):
        return _betavariate_impl(np.random.gamma)


def _betavariate_impl(gamma):
    def _impl(alpha, beta):
        """Beta distribution.  Taken from CPython.
        """
        # This version due to Janne Sinkkonen, and matches all the std
        # texts (e.g., Knuth Vol 2 Ed 3 pg 134 "the beta distribution").
        y = gamma(alpha, 1.)
        if y == 0.0:
            return 0.0
        else:
            return y / (y + gamma(beta, 1.))
    return _impl


@overload(np.random.beta)
def beta_impl(alpha, beta, size):
    if isinstance(size, types.NoneType):
        return lambda alpha, beta, size: np.random.beta(alpha, beta)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                            and isinstance(size.dtype,
                                                           types.Integer))):
        def _impl(alpha, beta, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.beta(alpha, beta)
            return out
        return _impl


@overload(random.expovariate)
def expovariate_impl(lambd):
    if isinstance(lambd, types.Float):
        def _impl(lambd):
            """Exponential distribution.  Taken from CPython.
            """
            # lambd: rate lambd = 1/mean
            # ('lambda' is a Python reserved word)

            # we use 1-random() instead of random() to preclude the
            # possibility of taking the log of zero.
            return -math.log(1.0 - random.random()) / lambd

        return _impl


@overload(np.random.exponential)
def exponential_impl(scale):
    if isinstance(scale, types.Number):
        def _impl(scale):
            return -math.log(1.0 - np.random.random()) * scale
        return _impl


@overload(np.random.exponential)
def exponential_impl(scale, size):
    if isinstance(size, types.NoneType):
        return lambda scale, size: np.random.exponential(scale)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                            isinstance(size.dtype,
                                                       types.Integer))):
        def _impl(scale, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.exponential(scale)
            return out
        return _impl


@overload(np.random.standard_exponential)
@overload(np.random.exponential)
def exponential_impl():
    def _impl():
        return -math.log(1.0 - np.random.random())
    return _impl


@overload(np.random.standard_exponential)
def standard_exponential_impl(size):
    if isinstance(size, types.NoneType):
        return lambda size: np.random.standard_exponential()
    if (isinstance(size, types.Integer) or
       (isinstance(size, types.UniTuple) and isinstance(size.dtype,
                                                        types.Integer))
    ):
        def _impl(size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.standard_exponential()
            return out
        return _impl


@overload(np.random.lognormal)
def np_lognormal_impl0():
    return lambda: np.random.lognormal(0.0, 1.0)


@overload(np.random.lognormal)
def np_log_normal_impl1(mu):
    if isinstance(mu, types.Number):
        return lambda mu: np.random.lognormal(mu, 1.0)


@overload(np.random.lognormal)
def np_log_normal_impl2(mu, sigma):
    if isinstance(mu, types.Number) and isinstance(sigma, types.Number):
        return _lognormvariate_impl(np.random.normal)


@overload(np.random.lognormal)
def lognormal_impl(mu, sigma, size):
    if isinstance(size, types.NoneType):
        return lambda mu, sigma, size: np.random.lognormal(mu, sigma)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                            isinstance(size.dtype,
                                                       types.Integer))):
        def _impl(mu, sigma, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.lognormal(mu, sigma)
            return out
        return _impl


@overload(random.lognormvariate)
def lognormvariate_impl(mu, sigma):
    if isinstance(mu, types.Float) and isinstance(sigma, types.Float):
        return _lognormvariate_impl(random.gauss)


def _lognormvariate_impl(_gauss):
    return lambda mu, sigma: math.exp(_gauss(mu, sigma))


@overload(random.paretovariate)
def paretovariate_impl(alpha):
    if isinstance(alpha, types.Float):
        def _impl(alpha):
            """Pareto distribution.  Taken from CPython."""
            # Jain, pg. 495
            u = 1.0 - random.random()
            return 1.0 / u ** (1.0/alpha)

        return _impl


@overload(np.random.pareto)
def pareto_impl(alpha):
    if isinstance(alpha, types.Float):
        def _impl(alpha):
            # Same as paretovariate() - 1.
            u = 1.0 - np.random.random()
            return 1.0 / u ** (1.0/alpha) - 1

        return _impl


@overload(np.random.pareto)
def pareto_impl(alpha, size):
    if isinstance(size, types.NoneType):
        return lambda alpha, size: np.random.pareto(alpha)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                            isinstance(size.dtype,
                                                       types.Integer))):
        def _impl(alpha, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.pareto(alpha)
            return out
        return _impl


@overload(random.weibullvariate)
def weibullvariate_impl(alpha, beta):
    if isinstance(alpha, types.Number) and isinstance(beta, types.Number):
        def _impl(alpha, beta):
            """Weibull distribution.  Taken from CPython."""
            # Jain, pg. 499; bug fix courtesy Bill Arms
            u = 1.0 - random.random()
            return alpha * (-math.log(u)) ** (1.0/beta)

        return _impl


@overload(np.random.weibull)
def weibull_impl(beta):
    if isinstance(beta, types.Number):
        def _impl(beta):
            # Same as weibullvariate(1.0, beta)
            u = 1.0 - np.random.random()
            return (-math.log(u)) ** (1.0/beta)

        return _impl


@overload(np.random.weibull)
def weibull_impl2(beta, size):
    if isinstance(size, types.NoneType):
        return lambda beta, size: np.random.weibull(beta)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                            isinstance(size.dtype,
                                                       types.Integer))):
        def _impl(beta, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.weibull(beta)
            return out
        return _impl


@overload(random.vonmisesvariate)
def vonmisesvariate_impl(mu, kappa):
    if isinstance(mu, types.Float) and isinstance(kappa, types.Float):
        return _vonmisesvariate_impl(random.random)


@overload(np.random.vonmises)
def vonmisesvariate_impl(mu, kappa):
    if isinstance(mu, types.Float) and isinstance(kappa, types.Float):
        return _vonmisesvariate_impl(np.random.random)


def _vonmisesvariate_impl(_random):
    def _impl(mu, kappa):
        """Circular data distribution.  Taken from CPython.
        Note the algorithm in Python 2.6 and Numpy is different:
        http://bugs.python.org/issue17141
        """
        # 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.
        if kappa <= 1e-6:
            return 2.0 * math.pi * _random()

        s = 0.5 / kappa
        r = s + math.sqrt(1.0 + s * s)

        while 1:
            u1 = _random()
            z = math.cos(math.pi * u1)

            d = z / (r + z)
            u2 = _random()
            if u2 < 1.0 - d * d or u2 <= (1.0 - d) * math.exp(d):
                break

        q = 1.0 / r
        f = (q + z) / (1.0 + q * z)
        u3 = _random()
        if u3 > 0.5:
            theta = (mu + math.acos(f)) % (2.0 * math.pi)
        else:
            theta = (mu - math.acos(f)) % (2.0 * math.pi)

        return theta

    return _impl


@overload(np.random.vonmises)
def vonmises_impl(mu, kappa, size):
    if isinstance(size, types.NoneType):
        return lambda mu, kappa, size: np.random.vonmises(mu, kappa)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                            and isinstance(size.dtype,
                                                           types.Integer))):
        def _impl(mu, kappa, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.vonmises(mu, kappa)
            return out
        return _impl


@overload(np.random.binomial)
def binomial_impl(n, p):
    if isinstance(n, types.Integer) and isinstance(p, types.Number):
        def _impl(n, p):
            """
            Binomial distribution.  Numpy's variant of the BINV algorithm
            is used.
            (Numpy uses BTPE for n*p >= 30, though)
            """
            if n < 0:
                raise ValueError("binomial(): n <= 0")
            if not (0.0 <= p <= 1.0):
                raise ValueError("binomial(): p outside of [0, 1]")
            if p == 0.0:
                return 0
            if p == 1.0:
                return n

            flipped = p > 0.5
            if flipped:
                p = 1.0 - p
            q = 1.0 - p

            niters = 1
            qn = q ** n
            while qn <= 1e-308:
                # Underflow => split into several iterations
                # Note this is much slower than Numpy's BTPE
                niters <<= 2
                n >>= 2
                qn = q ** n
                assert n > 0

            np_prod = n * p
            bound = min(n, np_prod + 10.0 * math.sqrt(np_prod * q + 1))

            total = 0
            while niters > 0:
                X = 0
                U = np.random.random()
                px = qn
                while X <= bound:
                    if U <= px:
                        total += n - X if flipped else X
                        niters -= 1
                        break
                    U -= px
                    X += 1
                    px = ((n - X + 1) * p * px) / (X * q)

            return total

        return _impl


@overload(np.random.binomial)
def binomial_impl(n, p, size):
    if isinstance(size, types.NoneType):
        return lambda n, p, size: np.random.binomial(n, p)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                            isinstance(size.dtype,
                                                       types.Integer))):
        def _impl(n, p, size):
            out = np.empty(size, dtype=np.intp)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.binomial(n, p)
            return out
        return _impl


@overload(np.random.chisquare)
def chisquare_impl(df):
    if isinstance(df, types.Number):
        def _impl(df):
            return 2.0 * np.random.standard_gamma(df / 2.0)

        return _impl


@overload(np.random.chisquare)
def chisquare_impl2(p, size):
    if isinstance(size, types.NoneType):
        return lambda p, size: np.random.chisquare(p)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                            isinstance(size.dtype,
                                                       types.Integer))):
        def _impl(p, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.chisquare(p)
            return out
        return _impl


@overload(np.random.f)
def f_impl(num, denom):
    if isinstance(num, types.Number) and isinstance(denom, types.Number):
        def _impl(num, denom):
            return ((np.random.chisquare(num) * denom) /
                    (np.random.chisquare(denom) * num))

        return _impl


@overload(np.random.f)
def f_impl(num, denom, size):
    if (isinstance(num, types.Number) and isinstance(denom, types.Number) and
       isinstance(size, types.NoneType)):
        return lambda num, denom, size: np.random.f(num, denom)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                            and isinstance(size.dtype,
                                                           types.Integer))):
        def _impl(num, denom, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.f(num, denom)
            return out
        return _impl


@overload(np.random.geometric)
def geometric_impl(p):
    if isinstance(p, types.Number):
        def _impl(p):
            # Numpy's algorithm.
            if p <= 0.0 or p > 1.0:
                raise ValueError("geometric(): p outside of (0, 1]")
            q = 1.0 - p
            if p >= 0.333333333333333333333333:
                X = int(1)
                sum = prod = p
                U = np.random.random()
                while U > sum:
                    prod *= q
                    sum += prod
                    X += 1
                return X
            else:
                return math.ceil(math.log(1.0 - np.random.random()) /
                                 math.log(q))

        return _impl


@overload(np.random.geometric)
def geometric_impl(p, size):
    if isinstance(size, types.NoneType):
        return lambda p, size: np.random.geometric(p)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                            isinstance(size.dtype,
                                                       types.Integer))):
        def _impl(p, size):
            out = np.empty(size, dtype=np.int64)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.geometric(p)
            return out
        return _impl


@overload(np.random.gumbel)
def gumbel_impl(loc, scale):
    if isinstance(loc, types.Number) and isinstance(scale, types.Number):
        def _impl(loc, scale):
            U = 1.0 - np.random.random()
            return loc - scale * math.log(-math.log(U))

        return _impl


@overload(np.random.gumbel)
def gumbel_impl3(loc, scale, size):
    if isinstance(size, types.NoneType):
        return lambda loc, scale, size: np.random.gumbel(loc, scale)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                            and isinstance(size.dtype,
                                                           types.Integer))):
        def _impl(loc, scale, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.gumbel(loc, scale)
            return out
        return _impl


@overload(np.random.hypergeometric)
def hypergeometric_impl(ngood, nbad, nsamples):
    if (isinstance(ngood, types.Number) and isinstance(nbad, types.Number)
       and isinstance(nsamples, types.Number)):
        def _impl(ngood, nbad, nsamples):
            """Numpy's algorithm for hypergeometric()."""
            d1 = int(nbad) + int(ngood) - int(nsamples)
            d2 = float(min(nbad, ngood))

            Y = d2
            K = int(nsamples)
            while Y > 0.0 and K > 0:
                Y -= math.floor(np.random.random() + Y / (d1 + K))
                K -= 1
            Z = int(d2 - Y)
            if ngood > nbad:
                return int(nsamples) - Z
            else:
                return Z

        return _impl


@overload(np.random.hypergeometric)
def hypergeometric_impl(ngood, nbad, nsamples, size):
    if isinstance(size, types.NoneType):
        return lambda ngood, nbad, nsamples, size:\
            np.random.hypergeometric(ngood, nbad, nsamples)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                            and isinstance(size.dtype,
                                                           types.Integer))):
        def _impl(ngood, nbad, nsamples, size):
            out = np.empty(size, dtype=np.intp)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.hypergeometric(ngood, nbad, nsamples)
            return out
        return _impl


@overload(np.random.laplace)
def laplace_impl0():
    return lambda: np.random.laplace(0.0, 1.0)


@overload(np.random.laplace)
def laplace_impl1(loc):
    if isinstance(loc, types.Number):
        return lambda loc: np.random.laplace(loc, 1.0)


@overload(np.random.laplace)
def laplace_impl2(loc, scale):
    if isinstance(loc, types.Number) and isinstance(scale, types.Number):
        return laplace_impl


@overload(np.random.laplace)
def laplace_impl3(loc, scale, size):
    if isinstance(size, types.NoneType):
        return lambda loc, scale, size: np.random.laplace(loc, scale)
    if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                           isinstance(size.dtype,
                                                      types.Integer)):
        def _impl(loc, scale, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.laplace(loc, scale)
            return out
        return _impl


def laplace_impl(loc, scale):
    U = np.random.random()
    if U < 0.5:
        return loc + scale * math.log(U + U)
    else:
        return loc - scale * math.log(2.0 - U - U)


@overload(np.random.logistic)
def logistic_impl0():
    return lambda: np.random.logistic(0.0, 1.0)


@overload(np.random.logistic)
def logistic_impl1(loc):
    if isinstance(loc, types.Number):
        return lambda loc: np.random.logistic(loc, 1.0)


@overload(np.random.logistic)
def logistic_impl2(loc, scale):
    if isinstance(loc, types.Number) and isinstance(scale, types.Number):
        return logistic_impl


@overload(np.random.logistic)
def logistic_impl3(loc, scale, size):
    if isinstance(size, types.NoneType):
        return lambda loc, scale, size: np.random.logistic(loc, scale)
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                            and isinstance(size.dtype,
                                                           types.Integer))):
        def _impl(loc, scale, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.logistic(loc, scale)
            return out
        return _impl


def logistic_impl(loc, scale):
    U = np.random.random()
    return loc + scale * math.log(U / (1.0 - U))


def _logseries_impl(p):
    """Numpy's algorithm for logseries()."""
    if p <= 0.0 or p > 1.0:
        raise ValueError("logseries(): p outside of (0, 1]")
    r = math.log(1.0 - p)

    while 1:
        V = np.random.random()
        if V >= p:
            return 1
        U = np.random.random()
        q = 1.0 - math.exp(r * U)
        if V <= q * q:
            # XXX what if V == 0.0 ?
            return np.int64(1.0 + math.log(V) / math.log(q))
        elif V >= q:
            return 1
        else:
            return 2


@overload(np.random.logseries)
def logseries_impl(p):
    if isinstance(p, types.Number):
        return _logseries_impl


@overload(np.random.logseries)
def logseries_impl(p, size):
    if isinstance(size, types.NoneType):
        return lambda p, size: np.random.logseries(p)
    if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                           isinstance(size.dtype,
                                                      types.Integer)):
        def _impl(p, size):
            out = np.empty(size, dtype=np.int64)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.logseries(p)
            return out
        return _impl


@overload(np.random.negative_binomial)
def negative_binomial_impl(n, p):
    if isinstance(n, types.Integer) and isinstance(p, types.Number):
        def _impl(n, p):
            if n <= 0:
                raise ValueError("negative_binomial(): n <= 0")
            if p < 0.0 or p > 1.0:
                raise ValueError("negative_binomial(): p outside of [0, 1]")
            Y = np.random.gamma(n, (1.0 - p) / p)
            return np.random.poisson(Y)

        return _impl


@glue_lowering("np.random.poisson")
@glue_lowering("np.random.poisson", types.Float)
def poisson_impl(context, builder, sig, args):
    state_ptr = get_np_state_ptr(context, builder)

    retptr = cgutils.alloca_once(builder, int64_t, name="ret")
    bbcont = builder.append_basic_block("bbcont")
    bbend = builder.append_basic_block("bbend")

    if len(args) == 1:
        lam, = args
        big_lam = builder.fcmp_ordered('>=', lam, ir.Constant(double, 10.0))
        with builder.if_then(big_lam):
            # For lambda >= 10.0, we switch to a more accurate
            # algorithm (see _random.c).
            fnty = ir.FunctionType(int64_t, (rnd_state_ptr_t, double))
            fn = cgutils.get_or_insert_function(builder.function.module, fnty,
                                                "numba_poisson_ptrs")
            ret = builder.call(fn, (state_ptr, lam))
            builder.store(ret, retptr)
            builder.branch(bbend)

    builder.branch(bbcont)
    builder.position_at_end(bbcont)

    _random = np.random.random
    _exp = math.exp

    def poisson_impl(lam):
        """Numpy's algorithm for poisson() on small *lam*.

        This method is invoked only if the parameter lambda of the
        distribution is small ( < 10 ). The algorithm used is described
        in "Knuth, D. 1969. 'Seminumerical Algorithms. The Art of
        Computer Programming' vol 2.
        """
        if lam < 0.0:
            raise ValueError("poisson(): lambda < 0")
        if lam == 0.0:
            return 0
        enlam = _exp(-lam)
        X = 0
        prod = 1.0
        while 1:
            U = _random()
            prod *= U
            if prod <= enlam:
                return X
            X += 1

    if len(args) == 0:
        sig = signature(sig.return_type, types.float64)
        args = (ir.Constant(double, 1.0),)

    ret = context.compile_internal(builder, poisson_impl, sig, args)
    builder.store(ret, retptr)
    builder.branch(bbend)
    builder.position_at_end(bbend)
    res = builder.load(retptr)
    return impl_ret_untracked(context, builder, sig.return_type, res)


@overload(np.random.power)
def power_impl(a):
    if isinstance(a, types.Number):
        def _impl(a):
            if a <= 0.0:
                raise ValueError("power(): a <= 0")
            return math.pow(1 - math.exp(-np.random.standard_exponential()),
                            1./a)

        return _impl


@overload(np.random.power)
def power_impl(a, size):
    if isinstance(size, types.NoneType):
        return lambda a, size: np.random.power(a)
    if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                           isinstance(size.dtype,
                                                      types.Integer)):
        def _impl(a, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.power(a)
            return out
        return _impl


@overload(np.random.rayleigh)
def rayleigh_impl0():
    return lambda: np.random.rayleigh(1.0)


@overload(np.random.rayleigh)
def rayleigh_impl1(mode):
    if isinstance(mode, types.Number):
        return rayleigh_impl


def rayleigh_impl(mode):
    if mode <= 0.0:
        raise ValueError("rayleigh(): mode <= 0")
    return mode * math.sqrt(-2.0 * math.log(1.0 - np.random.random()))


@overload(np.random.rayleigh)
def rayleigh_impl2(mode, size):
    if isinstance(size, types.NoneType):
        return lambda mode, size: np.random.rayleigh(mode)
    if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                           isinstance(size.dtype,
                                                      types.Integer)):
        def _impl(mode, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.rayleigh(mode)
            return out
        return _impl


@overload(np.random.standard_cauchy)
def cauchy_impl():
    def _impl():
        return np.random.standard_normal() / np.random.standard_normal()

    return _impl


@overload(np.random.standard_cauchy)
def standard_cauchy_impl(size):
    if isinstance(size, types.NoneType):
        return lambda size: np.random.standard_cauchy()
    if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                           and isinstance(size.dtype,
                                                          types.Integer)):
        def _impl(size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.standard_cauchy()
            return out
        return _impl


@overload(np.random.standard_t)
def standard_t_impl(df):
    if isinstance(df, types.Number):
        def _impl(df):
            N = np.random.standard_normal()
            G = np.random.standard_gamma(df / 2.0)
            X = math.sqrt(df / 2.0) * N / math.sqrt(G)
            return X

        return _impl


@overload(np.random.standard_t)
def standard_t_impl2(df, size):
    if isinstance(size, types.NoneType):
        return lambda p, size: np.random.standard_t(p)
    if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                           isinstance(size.dtype,
                                                      types.Integer)):
        def _impl(df, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.standard_t(df)
            return out
        return _impl


@overload(np.random.wald)
def wald_impl(mean, scale):
    if isinstance(mean, types.Float) and isinstance(scale, types.Float):
        def _impl(mean, scale):
            if mean <= 0.0:
                raise ValueError("wald(): mean <= 0")
            if scale <= 0.0:
                raise ValueError("wald(): scale <= 0")
            mu_2l = mean / (2.0 * scale)
            Y = np.random.standard_normal()
            Y = mean * Y * Y
            X = mean + mu_2l * (Y - math.sqrt(4 * scale * Y + Y * Y))
            U = np.random.random()
            if U <= mean / (mean + X):
                return X
            else:
                return mean * mean / X

        return _impl


@overload(np.random.wald)
def wald_impl2(mean, scale, size):
    if isinstance(size, types.NoneType):
        return lambda mean, scale, size: np.random.wald(mean, scale)
    if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                           isinstance(size.dtype,
                                                      types.Integer)):
        def _impl(mean, scale, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.wald(mean, scale)
            return out
        return _impl


@overload(np.random.zipf)
def zipf_impl(a):
    if isinstance(a, types.Float):
        def _impl(a):
            if a <= 1.0:
                raise ValueError("zipf(): a <= 1")
            am1 = a - 1.0
            b = 2.0 ** am1
            while 1:
                U = 1.0 - np.random.random()
                V = np.random.random()
                X = int(math.floor(U ** (-1.0 / am1)))
                T = (1.0 + 1.0 / X) ** am1
                if X >= 1 and V * X * (T - 1.0) / (b - 1.0) <= (T / b):
                    return X

        return _impl


@overload(np.random.zipf)
def zipf_impl(a, size):
    if isinstance(size, types.NoneType):
        return lambda a, size: np.random.zipf(a)
    if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                           isinstance(size.dtype,
                                                      types.Integer)):
        def _impl(a, size):
            out = np.empty(size, dtype=np.intp)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.zipf(a)
            return out
        return _impl

def do_shuffle_impl(arr, rng):

    if not isinstance(arr, types.Buffer):
        raise TypeError("The argument to shuffle() should be a buffer type")

    if rng == "np":
        rand = np.random.randint
    elif rng == "py":
        rand = random.randrange

    if arr.ndim == 1:
        def impl(arr):
            i = arr.shape[0] - 1
            while i > 0:
                j = rand(i + 1)
                arr[i], arr[j] = arr[j], arr[i]
                i -= 1
    else:
        def impl(arr):
            i = arr.shape[0] - 1
            while i > 0:
                j = rand(i + 1)
                arr[i], arr[j] = np.copy(arr[j]), np.copy(arr[i])
                i -= 1

    return impl


@overload(random.shuffle)
def shuffle_impl(arr):
    return do_shuffle_impl(arr, "py")


@overload(np.random.shuffle)
def shuffle_impl(arr):
    return do_shuffle_impl(arr, "np")


@overload(np.random.permutation)
def permutation_impl(x):
    if isinstance(x, types.Integer):
        def permutation_impl(x):
            y = np.arange(x)
            np.random.shuffle(y)
            return y
    elif isinstance(x, types.Array):
        def permutation_impl(x):
            arr_copy = x.copy()
            np.random.shuffle(arr_copy)
            return arr_copy
    else:
        permutation_impl = None
    return permutation_impl


# ------------------------------------------------------------------------
# Array-producing variants of scalar random functions

for typing_key, arity in [
    ("np.random.normal", 3),
    ("np.random.poisson", 2),
    ("np.random.random", 1),
    ("np.random.random_sample", 1),
    ("np.random.ranf", 1),
    ("np.random.sample", 1),
    ("np.random.randint", 3),
    ("np.random.standard_normal", 1),
    ("np.random.uniform", 3),
    ]:

    @glue_lowering(typing_key, *(types.Any,) * arity)
    def random_arr(context, builder, sig, args, typing_key=typing_key):

        arrty = sig.return_type
        dtype = arrty.dtype
        scalar_sig = signature(dtype, *sig.args[:-1])
        scalar_args = args[:-1]

        # Allocate array...
        shapes = arrayobj._parse_shape(context, builder, sig.args[-1], args[-1])
        arr = arrayobj._empty_nd_impl(context, builder, arrty, shapes)

        # ... and populate it in natural order
        *mod, fname = typing_key.split('.')
        # Module must be numpy.random
        assert mod == ['np', 'random']
        np_func = getattr(np.random, fname)
        fnty = context.typing_context.resolve_value_type(np_func)
        resolved_sig = fnty.get_call_type(context.typing_context,
                                          scalar_sig.args, {})
        scalar_impl = context.get_function(fnty, resolved_sig)
        with cgutils.for_range(builder, arr.nitems) as loop:
            val = scalar_impl(builder, scalar_args)
            ptr = cgutils.gep(builder, arr.data, loop.index)
            arrayobj.store_item(context, builder, arrty, val, ptr)

        return impl_ret_new_ref(context, builder, sig.return_type, arr._getvalue())


# ------------------------------------------------------------------------
# Irregular aliases: np.random.rand, np.random.randn

@overload(np.random.rand)
def rand(*size):
    if len(size) == 0:
        # Scalar output
        def rand_impl(*size):
            return np.random.random()

    else:
        # Array output
        def rand_impl(*size):
            return np.random.random(size)

    return rand_impl

@overload(np.random.randn)
def randn(*size):
    if len(size) == 0:
        # Scalar output
        def randn_impl(*size):
            return np.random.standard_normal()

    else:
        # Array output
        def randn_impl(*size):
            return np.random.standard_normal(size)

    return randn_impl


# ------------------------------------------------------------------------
# np.random.choice

@overload(np.random.choice)
def choice(a, size=None, replace=True):

    if isinstance(a, types.Array):
        # choice() over an array population
        assert a.ndim == 1
        dtype = a.dtype

        @register_jitable
        def get_source_size(a):
            return len(a)

        @register_jitable
        def copy_source(a):
            return a.copy()

        @register_jitable
        def getitem(a, a_i):
            return a[a_i]

    elif isinstance(a, types.Integer):
        # choice() over an implied arange() population
        dtype = np.intp

        @register_jitable
        def get_source_size(a):
            return a

        @register_jitable
        def copy_source(a):
            return np.arange(a)

        @register_jitable
        def getitem(a, a_i):
            return a_i

    else:
        raise TypeError("np.random.choice() first argument should be "
                        "int or array, got %s" % (a,))

    if size in (None, types.none):
        def choice_impl(a, size=None, replace=True):
            """
            choice() implementation returning a single sample
            (note *replace* is ignored)
            """
            n = get_source_size(a)
            i = np.random.randint(0, n)
            return getitem(a, i)

    else:
        def choice_impl(a, size=None, replace=True):
            """
            choice() implementation returning an array of samples
            """
            n = get_source_size(a)
            if replace:
                out = np.empty(size, dtype)
                fl = out.flat
                for i in range(len(fl)):
                    j = np.random.randint(0, n)
                    fl[i] = getitem(a, j)
                return out
            else:
                # Note we have to construct the array to compute out.size
                # (`size` can be an arbitrary int or tuple of ints)
                out = np.empty(size, dtype)
                if out.size > n:
                    raise ValueError("Cannot take a larger sample than "
                                     "population when 'replace=False'")
                # Get a permuted copy of the source array
                # we need this implementation in order to get the
                # np.random.choice inside numba to match the output
                # of np.random.choice outside numba when np.random.seed
                # is set to the same value
                permuted_a = np.random.permutation(a)
                fl = out.flat
                for i in range(len(fl)):
                    fl[i] = permuted_a[i]
                return out

    return choice_impl


# ------------------------------------------------------------------------
# np.random.multinomial

@overload(np.random.multinomial)
def multinomial(n, pvals, size=None):

    dtype = np.intp

    @register_jitable
    def multinomial_inner(n, pvals, out):
        # Numpy's algorithm for multinomial()
        fl = out.flat
        sz = out.size
        plen = len(pvals)

        for i in range(0, sz, plen):
            # Loop body: take a set of n experiments and fill up
            # fl[i:i + plen] with the distribution of results.

            # Current sum of outcome probabilities
            p_sum = 1.0
            # Current remaining number of experiments
            n_experiments = n
            # For each possible outcome `j`, compute the number of results
            # with this outcome.  This is done by considering the
            # conditional probability P(X=j | X>=j) and running a binomial
            # distribution over the remaining number of experiments.
            for j in range(0, plen - 1):
                p_j = pvals[j]
                n_j = fl[i + j] = np.random.binomial(n_experiments, p_j / p_sum)
                n_experiments -= n_j
                if n_experiments <= 0:
                    # Note the output was initialized to zero
                    break
                p_sum -= p_j
            if n_experiments > 0:
                # The remaining experiments end up in the last bucket
                fl[i + plen - 1] = n_experiments

    if not isinstance(n, types.Integer):
        raise TypeError("np.random.multinomial(): n should be an "
                        "integer, got %s" % (n,))

    if not isinstance(pvals, (types.Sequence, types.Array)):
        raise TypeError("np.random.multinomial(): pvals should be an "
                        "array or sequence, got %s" % (pvals,))

    if size in (None, types.none):
        def multinomial_impl(n, pvals, size=None):
            """
            multinomial(..., size=None)
            """
            out = np.zeros(len(pvals), dtype)
            multinomial_inner(n, pvals, out)
            return out

    elif isinstance(size, types.Integer):
        def multinomial_impl(n, pvals, size=None):
            """
            multinomial(..., size=int)
            """
            out = np.zeros((size, len(pvals)), dtype)
            multinomial_inner(n, pvals, out)
            return out

    elif isinstance(size, types.BaseTuple):
        def multinomial_impl(n, pvals, size=None):
            """
            multinomial(..., size=tuple)
            """
            out = np.zeros(size + (len(pvals),), dtype)
            multinomial_inner(n, pvals, out)
            return out

    else:
        raise TypeError("np.random.multinomial(): size should be int or "
                        "tuple or None, got %s" % (size,))

    return multinomial_impl

# ------------------------------------------------------------------------
# np.random.dirichlet


@overload(np.random.dirichlet)
def dirichlet(alpha):
    if isinstance(alpha, (types.Sequence, types.Array)):
        def dirichlet_impl(alpha):
            out = np.empty(len(alpha))
            dirichlet_arr(alpha, out)
            return out
        return dirichlet_impl


@overload(np.random.dirichlet)
def dirichlet(alpha, size=None):
    if not isinstance(alpha, (types.Sequence, types.Array)):
        raise NumbaTypeError(
            "np.random.dirichlet(): alpha should be an "
            "array or sequence, got %s" % (alpha,)
        )

    if size in (None, types.none):

        def dirichlet_impl(alpha, size=None):
            out = np.empty(len(alpha))
            dirichlet_arr(alpha, out)
            return out

    elif isinstance(size, types.Integer):

        def dirichlet_impl(alpha, size=None):
            """
            dirichlet(..., size=int)
            """
            out = np.empty((size, len(alpha)))
            dirichlet_arr(alpha, out)
            return out

    elif isinstance(size, types.UniTuple) and isinstance(size.dtype,
                                                         types.Integer):
        def dirichlet_impl(alpha, size=None):
            """
            dirichlet(..., size=tuple)
            """
            out = np.empty(size + (len(alpha),))
            dirichlet_arr(alpha, out)
            return out

    else:
        raise NumbaTypeError(
            "np.random.dirichlet(): size should be int or "
            "tuple of ints or None, got %s" % size
        )

    return dirichlet_impl


@register_jitable
def dirichlet_arr(alpha, out):

    # Gamma distribution method to generate a Dirichlet distribution

    for a_val in iter(alpha):
        if a_val <= 0:
            raise ValueError("dirichlet: alpha must be > 0.0")

    a_len = len(alpha)
    size = out.size
    flat = out.flat
    for i in range(0, size, a_len):
        # calculate gamma random numbers per alpha specifications
        norm = 0  # use this to normalize every the group total to 1
        for k, w in enumerate(alpha):
            flat[i + k] = np.random.gamma(w, 1)
            norm += flat[i + k].item()
        for k, w in enumerate(alpha):
            flat[i + k] /= norm


# ------------------------------------------------------------------------
# np.random.noncentral_chisquare


@overload(np.random.noncentral_chisquare)
def noncentral_chisquare(df, nonc):
    if isinstance(df, types.Number) and isinstance(nonc, types.Number):
        def noncentral_chisquare_impl(df, nonc):
            validate_noncentral_chisquare_input(df, nonc)
            return noncentral_chisquare_single(df, nonc)

        return noncentral_chisquare_impl


@overload(np.random.noncentral_chisquare)
def noncentral_chisquare(df, nonc, size=None):
     if size in (None, types.none):
         def noncentral_chisquare_impl(df, nonc, size=None):
             validate_noncentral_chisquare_input(df, nonc)
             return noncentral_chisquare_single(df, nonc)
         return noncentral_chisquare_impl
     elif isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                              and isinstance(size.dtype,
                                                             types.Integer)):

         def noncentral_chisquare_impl(df, nonc, size=None):
             validate_noncentral_chisquare_input(df, nonc)
             out = np.empty(size)
             out_flat = out.flat
             for idx in range(out.size):
                 out_flat[idx] = noncentral_chisquare_single(df, nonc)
             return out
         return noncentral_chisquare_impl
     else:
         raise NumbaTypeError(
             "np.random.noncentral_chisquare(): size should be int or "
             "tuple of ints or None, got %s" % size
         )


@register_jitable
def noncentral_chisquare_single(df, nonc):
    # identical to numpy implementation from distributions.c
    # https://github.com/numpy/numpy/blob/c65bc212ec1987caefba0ea7efe6a55803318de9/numpy/random/src/distributions/distributions.c#L797

    if np.isnan(nonc):
        return np.nan

    if 1 < df:
        chi2 = np.random.chisquare(df-1)
        n = np.random.standard_normal() + np.sqrt(nonc)
        return chi2 + n * n

    else:
        i = np.random.poisson(nonc/2.0)
        return np.random.chisquare(df + 2 * i)


@register_jitable
def validate_noncentral_chisquare_input(df, nonc):
    if df <= 0:
        raise ValueError("df <= 0")
    if nonc < 0:
        raise ValueError("nonc < 0")
