randomimpl.py

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


import math
import random

import numpy as np

from llvmlite import ir

from numba.core.cgutils import is_nonelike, is_empty_tuple
from numba.core.extending import intrinsic, 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, cgutils
from numba.core.errors import NumbaTypeError
from numba.np.random._constants import LONG_MAX

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)

        # This truncation/extension is safe because 0 < nbits <= 64
        if nbits.type.width < y.type.width:
            shift = builder.zext(shift, y.type)
        elif nbits.type.width > y.type.width:
            shift = builder.trunc(shift, y.type)

        if is_numpy:
            # Use the last N bits, to match np.random
            mask = builder.not_(ir.Constant(y.type, 0))
            mask = builder.lshr(mask, shift)
            return builder.and_(y, mask)
        else:
            # Use the first N bits, to match CPython random
            return builder.lshr(y, shift)

    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)


@overload(random.seed)
def seed_impl(a):
    if isinstance(a, types.Integer):
        fn = register_jitable(_seed_impl('py'))
        def impl(a):
            return fn(a)
        return impl


@overload(np.random.seed)
def seed_impl(seed):
    if isinstance(seed, types.Integer):
        return _seed_impl('np')


def _seed_impl(state_type):
    @intrinsic
    def _impl(typingcontext, seed):
        def codegen(context, builder, sig, args):
            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, (get_state_ptr(context, builder, state_type),
                              seed_value))
            return context.get_constant(types.none, None)
        return signature(types.void, types.uint32), codegen
    return lambda seed: _impl(seed)


@overload(random.random)
def random_impl():
    @intrinsic
    def _impl(typingcontext):
        def codegen(context, builder, sig, args):
            state_ptr = get_state_ptr(context, builder, "py")
            return get_next_double(context, builder, state_ptr)
        return signature(types.double), codegen
    return lambda: _impl()


@overload(np.random.random)
@overload(np.random.random_sample)
@overload(np.random.sample)
@overload(np.random.ranf)
def random_impl0():
    @intrinsic
    def _impl(typingcontext):
        def codegen(context, builder, sig, args):
            state_ptr = get_state_ptr(context, builder, "np")
            return get_next_double(context, builder, state_ptr)
        return signature(types.float64), codegen
    return lambda: _impl()


@overload(np.random.random)
@overload(np.random.random_sample)
@overload(np.random.sample)
@overload(np.random.ranf)
def random_impl1(size=None):
    if is_nonelike(size):
        return lambda size=None: np.random.random()
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda size=None: np.array(np.random.random())
    if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                           and isinstance(size.dtype,
                                                          types.Integer)):
        def _impl(size=None):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.random()
            return out
        return _impl


@overload(random.gauss)
@overload(random.normalvariate)
def gauss_impl(mu, sigma):
    if isinstance(mu, (types.Float, types.Integer)) and isinstance(
            sigma, (types.Float, types.Integer)):
        @intrinsic
        def _impl(typingcontext, mu, sigma):
            loc_preprocessor = _double_preprocessor(mu)
            scale_preprocessor = _double_preprocessor(sigma)
            return signature(types.float64, mu, sigma),\
                   _gauss_impl("py", loc_preprocessor, scale_preprocessor)
        return lambda mu, sigma: _impl(mu, sigma)


@overload(np.random.standard_normal)
@overload(np.random.normal)
def np_gauss_impl0():
    return lambda: np.random.normal(0.0, 1.0)


@overload(np.random.normal)
def np_gauss_impl1(loc):
    if isinstance(loc, (types.Float, types.Integer)):
        return lambda loc: np.random.normal(loc, 1.0)


@overload(np.random.normal)
def np_gauss_impl2(loc, scale):
    if isinstance(loc, (types.Float, types.Integer)) and isinstance(
            scale, (types.Float, types.Integer)):
        @intrinsic
        def _impl(typingcontext, loc, scale):
            loc_preprocessor = _double_preprocessor(loc)
            scale_preprocessor = _double_preprocessor(scale)
            return signature(types.float64, loc, scale),\
                   _gauss_impl("np", loc_preprocessor, scale_preprocessor)
        return lambda loc, scale: _impl(loc, scale)


@overload(np.random.standard_normal)
def standard_normal_impl1(size):
    if is_nonelike(size):
        return lambda size: np.random.standard_normal()
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda size: np.array(np.random.standard_normal())
    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_normal()
            return out
        return _impl


@overload(np.random.normal)
def np_gauss_impl3(loc, scale, size):
    if (isinstance(loc, (types.Float, types.Integer)) and isinstance(
            scale, (types.Float, types.Integer)) and
       is_nonelike(size)):
        return lambda loc, scale, size: np.random.normal(loc, scale)
    if (isinstance(loc, (types.Float, types.Integer)) and isinstance(
            scale, (types.Float, types.Integer)) and
       is_empty_tuple(size)):
        # Handle size = ()
        return lambda loc, scale, size: np.array(np.random.normal(loc, scale))
    if (isinstance(loc, (types.Float, types.Integer)) and isinstance(
            scale, (types.Float, types.Integer)) and
       (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.normal(loc, scale)
            return out
        return _impl


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(state, loc_preprocessor, scale_preprocessor):
    def _impl(context, builder, sig, args):
        # The type for all computations (either float or double)
        ty = sig.return_type
        llty = context.get_data_type(ty)
        _random = {"py": random.random,
                   "np": np.random.random}[state]

        state_ptr = get_state_ptr(context, builder, 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(loc_preprocessor(builder, mu),
                            builder.fmul(scale_preprocessor(builder, sigma),
                                         builder.load(ret)))
    return _impl


def _double_preprocessor(value):
    ty = ir.types.DoubleType()

    if isinstance(value, types.Integer):
        if value.signed:
            return lambda builder, v: builder.sitofp(v, ty)
        else:
            return lambda builder, v: builder.uitofp(v, ty)
    elif isinstance(value, types.Float):
        if value.bitwidth != 64:
            return lambda builder, v: builder.fpext(v, ty)
        else:
            return lambda _builder, v: v
    else:
        raise NumbaTypeError("Cannot convert {} to floating point type" % value)


@overload(random.getrandbits)
def getrandbits_impl(k):
    if isinstance(k, types.Integer):
        @intrinsic
        def _impl(typingcontext, k):
            def codegen(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")
                return get_next_int(context, builder, state_ptr, nbits, False)
            return signature(types.uint64, k), codegen
        return lambda k: _impl(k)


def _randrange_impl(context, builder, start, stop, step, ty, signed, state):
    state_ptr = get_state_ptr(context, builder, state)
    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))


@overload(random.randrange)
def randrange_impl_1(start):
    if isinstance(start, types.Integer):
        return lambda start: random.randrange(0, start, 1)


@overload(random.randrange)
def randrange_impl_2(start, stop):
    if isinstance(start, types.Integer) and isinstance(stop, types.Integer):
        return lambda start, stop: random.randrange(start, stop, 1)


def _randrange_preprocessor(bitwidth, ty):
    if ty.bitwidth != bitwidth:
        return (ir.IRBuilder.sext if ty.signed
                else ir.IRBuilder.zext)
    else:
        return lambda _builder, v, _ty: v


@overload(random.randrange)
def randrange_impl_3(start, stop, step):
    if (isinstance(start, types.Integer) and isinstance(stop, types.Integer) and
       isinstance(step, types.Integer)):
        signed = max(start.signed, stop.signed, step.signed)
        bitwidth = max(start.bitwidth, stop.bitwidth, step.bitwidth)
        int_ty = types.Integer.from_bitwidth(bitwidth, signed)
        llvm_type = ir.IntType(bitwidth)

        start_preprocessor = _randrange_preprocessor(bitwidth, start)
        stop_preprocessor = _randrange_preprocessor(bitwidth, stop)
        step_preprocessor = _randrange_preprocessor(bitwidth, step)

        @intrinsic
        def _impl(typingcontext, start, stop, step):
            def codegen(context, builder, sig, args):
                start, stop, step = args

                start = start_preprocessor(builder, start, llvm_type)
                stop = stop_preprocessor(builder, stop, llvm_type)
                step = step_preprocessor(builder, step, llvm_type)
                return _randrange_impl(context, builder, start, stop, step,
                                       llvm_type, signed, 'py')
            return signature(int_ty, start, stop, step), codegen
        return lambda start, stop, step: _impl(start, stop, step)


@overload(random.randint)
def randint_impl_1(a, b):
    if isinstance(a, types.Integer) and isinstance(b, types.Integer):
        return lambda a, b: random.randrange(a, b + 1, 1)


@overload(np.random.randint)
def np_randint_impl_1(low):
    if isinstance(low, types.Integer):
        return lambda low: np.random.randint(0, low)


@overload(np.random.randint)
def np_randint_impl_2(low, high):
    if isinstance(low, types.Integer) and isinstance(high, types.Integer):
        signed = max(low.signed, high.signed)
        bitwidth = max(low.bitwidth, high.bitwidth)
        int_ty = types.Integer.from_bitwidth(bitwidth, signed)
        llvm_type = ir.IntType(bitwidth)

        start_preprocessor = _randrange_preprocessor(bitwidth, low)
        stop_preprocessor = _randrange_preprocessor(bitwidth, high)

        @intrinsic
        def _impl(typingcontext, low, high):
            def codegen(context, builder, sig, args):
                start, stop = args

                start = start_preprocessor(builder, start, llvm_type)
                stop = stop_preprocessor(builder, stop, llvm_type)
                step = ir.Constant(llvm_type, 1)
                return _randrange_impl(context, builder, start, stop, step,
                                       llvm_type, signed, 'np')
            return signature(int_ty, low, high), codegen
        return lambda low, high: _impl(low, high)


@overload(np.random.randint)
def np_randint_impl_3(low, high, size):
    if (isinstance(low, types.Integer) and isinstance(high, types.Integer) and
       is_nonelike(size)):
        return lambda low, high, size: np.random.randint(low, high)
    if (isinstance(low, types.Integer) and isinstance(high, types.Integer) and
       is_empty_tuple(size)):
        # Handle size = ()
        return lambda low, high, size: np.array(np.random.randint(low, high))
    if (isinstance(low, types.Integer) and isinstance(high, types.Integer) and
       (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                            and isinstance(size.dtype,
                                                           types.Integer)))):
        bitwidth = max(low.bitwidth, high.bitwidth)
        result_type = getattr(np, f'int{bitwidth}')

        def _impl(low, high, size):
            out = np.empty(size, dtype=result_type)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.randint(low, high)
            return out
        return _impl


@overload(np.random.uniform)
def np_uniform_impl0():
    return lambda: np.random.uniform(0.0, 1.0)


@overload(random.uniform)
def uniform_impl2(a, b):
    if isinstance(a, (types.Float, types.Integer)) and isinstance(
            b, (types.Float, types.Integer)):
        @intrinsic
        def _impl(typingcontext, a, b):
            low_preprocessor = _double_preprocessor(a)
            high_preprocessor = _double_preprocessor(b)
            return signature(types.float64, a, b), uniform_impl(
                'py', low_preprocessor, high_preprocessor)
        return lambda a, b: _impl(a, b)


@overload(np.random.uniform)
def np_uniform_impl2(low, high):
    if isinstance(low, (types.Float, types.Integer)) and isinstance(
            high, (types.Float, types.Integer)):
        @intrinsic
        def _impl(typingcontext, low, high):
            low_preprocessor = _double_preprocessor(low)
            high_preprocessor = _double_preprocessor(high)
            return signature(types.float64, low, high), uniform_impl(
                'np', low_preprocessor, high_preprocessor)
        return lambda low, high: _impl(low, high)


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


@overload(np.random.uniform)
def np_uniform_impl3(low, high, size):
    if (isinstance(low, (types.Float, types.Integer)) and isinstance(
            high, (types.Float, types.Integer)) and
       is_nonelike(size)):
        return lambda low, high, size: np.random.uniform(low, high)
    if (isinstance(low, (types.Float, types.Integer)) and isinstance(
            high, (types.Float, types.Integer)) and
       is_empty_tuple(size)):
        # When calling np.random.uniform with size = (), the returned value isn't a
        # float like when size = None. Instead, it's an array of shape ()
        return lambda low, high, size: np.array(np.random.uniform(low, high))
    if (isinstance(low, (types.Float, types.Integer)) and isinstance(
            high, (types.Float, types.Integer)) and
       (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                            and isinstance(size.dtype,
                                                           types.Integer)))):
        def _impl(low, high, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.uniform(low, high)
            return out
        return _impl


@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.Float, types.Integer)) and isinstance(
            high, (types.Float, types.Integer)):
        return _impl


@overload(random.triangular)
def triangular_impl_3(low, high, mode):
    if (isinstance(low, (types.Float, types.Integer)) and isinstance(
            high, (types.Float, types.Integer)) and
       isinstance(mode, (types.Float, types.Integer))):
        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(left, mode, right):
    if (isinstance(left, (types.Float, types.Integer)) and isinstance(
            mode, (types.Float, types.Integer)) and
            isinstance(right, (types.Float, types.Integer))):
        def _impl(left, mode, right):
            if right == left:
                return left
            u = np.random.random()
            c = (mode - left) / (right - left)
            if u > c:
                u = 1.0 - u
                c = 1.0 - c
                left, right = right, left
            return left + (right - left) * math.sqrt(u * c)

        return _impl


@overload(np.random.triangular)
def triangular_impl(left, mode, right, size=None):
    if is_nonelike(size):
        return lambda left, mode, right, size=None: np.random.triangular(left,
                                                                         mode,
                                                                         right)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda left, mode, right, size=None: np.array(
            np.random.triangular(left, mode, right)
        )
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                            isinstance(size.dtype,
                                                       types.Integer))):
        def _impl(left, mode, right, size=None):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.triangular(left, mode, right)
            return out
        return _impl


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


@overload(np.random.standard_gamma)
@overload(np.random.gamma)
def ol_np_random_gamma1(shape):
    if isinstance(shape, (types.Float, types.Integer)):
        return lambda shape: np.random.gamma(shape, 1.0)


@overload(np.random.gamma)
def ol_np_random_gamma2(shape, scale):
    if isinstance(shape, (types.Float, types.Integer)) and isinstance(
            scale, (types.Float, types.Integer)):
        fn = register_jitable(_gammavariate_impl(np.random.random))
        def impl(shape, scale):
            return fn(shape, scale)
        return impl


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)

            # Adjust due to cpython
            # commit 63d152232e1742660f481c04a811f824b91f6790
            return -math.log(1.0 - _random()) * 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(shape, scale, size):
    if is_nonelike(size):
        return lambda shape, scale, size: np.random.gamma(shape, scale)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda shape, scale, size: np.array(np.random.gamma(shape, scale))
    if isinstance(size, types.Integer) or (isinstance(size, types.UniTuple) and
                                           isinstance(size.dtype,
                                                      types.Integer)):
        def _impl(shape, scale, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.gamma(shape, scale)
            return out
        return _impl


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


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


@overload(np.random.beta)
def ol_np_random_beta(a, b):
    if isinstance(a, (types.Float, types.Integer)) and isinstance(
            b, (types.Float, types.Integer)):
        fn =  register_jitable(_betavariate_impl(np.random.gamma))
        def impl(a, b):
            return fn(a, b)
        return impl


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(a, b, size):
    if is_nonelike(size):
        return lambda a, b, size: np.random.beta(a, b)
    if is_empty_tuple(size):
        # When calling np.random.beta with size = (), the returned value isn't a
        # float like when size = None. Instead, it's an array of shape ()
        return lambda a, b, size: np.array(np.random.beta(a, b))
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                            and isinstance(size.dtype,
                                                           types.Integer))):
        def _impl(a, b, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.beta(a, b)
            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.Float, types.Integer)):
        def _impl(scale):
            return -math.log(1.0 - np.random.random()) * scale
        return _impl


@overload(np.random.exponential)
def exponential_impl(scale, size):
    if is_nonelike(size):
        return lambda scale, size: np.random.exponential(scale)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda scale, size: np.array(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 is_nonelike(size):
        return lambda size: np.random.standard_exponential()
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda size: np.array(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(mean):
    if isinstance(mean, (types.Float, types.Integer)):
        return lambda mean: np.random.lognormal(mean, 1.0)


@overload(np.random.lognormal)
def np_log_normal_impl2(mean, sigma):
    if isinstance(mean, (types.Float, types.Integer)) and isinstance(
            sigma, (types.Float, types.Integer)):
        fn = register_jitable(_lognormvariate_impl(np.random.normal))
        return lambda mean, sigma: fn(mean, sigma)


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


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


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(a):
    if isinstance(a, types.Float):
        def _impl(a):
            # Same as paretovariate() - 1.
            u = 1.0 - np.random.random()
            return 1.0 / u ** (1.0/a) - 1

        return _impl


@overload(np.random.pareto)
def pareto_impl(a, size):
    if is_nonelike(size):
        return lambda a, size: np.random.pareto(a)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda a, size: np.array(np.random.pareto(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.pareto(a)
            return out
        return _impl


@overload(random.weibullvariate)
def weibullvariate_impl(alpha, beta):
    if isinstance(alpha, (types.Float, types.Integer)) and isinstance(
            beta, (types.Float, types.Integer)):
        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(a):
    if isinstance(a, (types.Float, types.Integer)):
        def _impl(a):
            # Same as weibullvariate(1.0, a)
            u = 1.0 - np.random.random()
            return (-math.log(u)) ** (1.0/a)

        return _impl


@overload(np.random.weibull)
def weibull_impl2(a, size):
    if is_nonelike(size):
        return lambda a, size: np.random.weibull(a)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda a, size: np.array(np.random.weibull(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.weibull(a)
            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 is_nonelike(size):
        return lambda mu, kappa, size: np.random.vonmises(mu, kappa)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda mu, kappa, size: np.array(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.Float, types.Integer)):
        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 is_nonelike(size):
        return lambda n, p, size: np.random.binomial(n, p)
    if is_empty_tuple(size):
        # When calling np.random.binomial with size = (), the returned value isn't a
        # float like when size = None. Instead, it's an array of shape ()
        return lambda n, p, size: np.array(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.Float, types.Integer)):
        def _impl(df):
            return 2.0 * np.random.standard_gamma(df / 2.0)

        return _impl


@overload(np.random.chisquare)
def chisquare_impl2(df, size):
    if is_nonelike(size):
        return lambda df, size: np.random.chisquare(df)
    if is_empty_tuple(size):
        return lambda df, size: np.array(np.random.chisquare(df))
    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.chisquare(df)
            return out
        return _impl


@overload(np.random.f)
def f_impl(dfnum, dfden):
    if isinstance(dfnum, (types.Float, types.Integer)) and isinstance(
            dfden, (types.Float, types.Integer)):
        def _impl(dfnum, dfden):
            return ((np.random.chisquare(dfnum) * dfden) /
                    (np.random.chisquare(dfden) * dfnum))

        return _impl


@overload(np.random.f)
def f_impl(dfnum, dfden, size):
    if (isinstance(dfnum, (types.Float, types.Integer)) and isinstance(
            dfden, (types.Float, types.Integer)) and
       is_nonelike(size)):
        return lambda dfnum, dfden, size: np.random.f(dfnum, dfden)
    if (isinstance(dfnum, (types.Float, types.Integer)) and isinstance(
            dfden, (types.Float, types.Integer)) and
       is_empty_tuple(size)):
        # Handle size = ()
        return lambda dfnum, dfden, size: np.array(np.random.f(dfnum, dfden))
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                            and isinstance(size.dtype,
                                                           types.Integer))):
        def _impl(dfnum, dfden, size):
            out = np.empty(size)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.f(dfnum, dfden)
            return out
        return _impl


@overload(np.random.geometric)
def geometric_impl(p):
    if isinstance(p, (types.Float, types.Integer)):
        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 is_nonelike(size):
        return lambda p, size: np.random.geometric(p)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda p, size: np.array(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.Float, types.Integer)) and isinstance(
            scale, (types.Float, types.Integer)):
        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 is_nonelike(size):
        return lambda loc, scale, size: np.random.gumbel(loc, scale)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda loc, scale, size: np.array(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, nsample):
    if (isinstance(ngood, (types.Float, types.Integer)) and isinstance(
            nbad, (types.Float, types.Integer))
       and isinstance(nsample, (types.Float, types.Integer))):
        def _impl(ngood, nbad, nsample):
            """Numpy's algorithm for hypergeometric()."""
            d1 = int(nbad) + int(ngood) - int(nsample)
            d2 = float(min(nbad, ngood))

            Y = d2
            K = int(nsample)
            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(nsample) - Z
            else:
                return Z

        return _impl


@overload(np.random.hypergeometric)
def hypergeometric_impl(ngood, nbad, nsample, size):
    if is_nonelike(size):
        return lambda ngood, nbad, nsample, size:\
            np.random.hypergeometric(ngood, nbad, nsample)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda ngood, nbad, nsample, size:\
            np.array(np.random.hypergeometric(ngood, nbad, nsample))
    if (isinstance(size, types.Integer) or (isinstance(size, types.UniTuple)
                                            and isinstance(size.dtype,
                                                           types.Integer))):
        def _impl(ngood, nbad, nsample, 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, nsample)
            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.Float, types.Integer)):
        return lambda loc: np.random.laplace(loc, 1.0)


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


@overload(np.random.laplace)
def laplace_impl3(loc, scale, size):
    if is_nonelike(size):
        return lambda loc, scale, size: np.random.laplace(loc, scale)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda loc, scale, size: np.array(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.Float, types.Integer)):
        return lambda loc: np.random.logistic(loc, 1.0)


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


@overload(np.random.logistic)
def logistic_impl3(loc, scale, size):
    if is_nonelike(size):
        return lambda loc, scale, size: np.random.logistic(loc, scale)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda loc, scale, size: np.array(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.Float, types.Integer)):
        return _logseries_impl


@overload(np.random.logseries)
def logseries_impl(p, size):
    if is_nonelike(size):
        return lambda p, size: np.random.logseries(p)
    if is_empty_tuple(size):
        return lambda p, size: np.array(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.Float, types.Integer)):
        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


@overload(np.random.poisson)
def poisson_impl0():
    return lambda: np.random.poisson(1.0)


@overload(np.random.poisson)
def poisson_impl1(lam):
    if isinstance(lam, (types.Float, types.Integer)):
        @intrinsic
        def _impl(typingcontext, lam):
            lam_preprocessor = _double_preprocessor(lam)

            def codegen(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")

                lam, = args
                lam = lam_preprocessor(builder, lam)
                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

                ret = context.compile_internal(builder, poisson_impl, sig, args)
                builder.store(ret, retptr)
                builder.branch(bbend)
                builder.position_at_end(bbend)
                return builder.load(retptr)
            return signature(types.int64, lam), codegen
        return lambda lam: _impl(lam)


@overload(np.random.poisson)
def poisson_impl2(lam, size):
    if isinstance(lam, (types.Float, types.Integer)) and is_nonelike(size):
        return lambda lam, size: np.random.poisson(lam)
    if isinstance(lam, (types.Float, types.Integer)) and is_empty_tuple(size):
        # Handle size = ()
        return lambda lam, size: np.array(np.random.poisson(lam))
    if isinstance(lam, (types.Float, types.Integer)) and (
            isinstance(size, types.Integer) or
       (isinstance(size, types.UniTuple) and isinstance(size.dtype,
                                                        types.Integer))
    ):
        def _impl(lam, size):
            out = np.empty(size, dtype=np.intp)
            out_flat = out.flat
            for idx in range(out.size):
                out_flat[idx] = np.random.poisson(lam)
            return out
        return _impl


@overload(np.random.power)
def power_impl(a):
    if isinstance(a, (types.Float, types.Integer)):
        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 is_nonelike(size):
        return lambda a, size: np.random.power(a)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda a, size: np.array(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(scale):
    if isinstance(scale, (types.Float, types.Integer)):
        def impl(scale):
            if scale <= 0.0:
                raise ValueError("rayleigh(): scale <= 0")
            return scale * math.sqrt(-2.0 * math.log(1.0 - np.random.random()))
        return impl




@overload(np.random.rayleigh)
def rayleigh_impl2(scale, size):
    if is_nonelike(size):
        return lambda scale, size: np.random.rayleigh(scale)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda scale, size: np.array(np.random.rayleigh(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.rayleigh(scale)
            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 is_nonelike(size):
        return lambda size: np.random.standard_cauchy()
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda size: np.array(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.Float, types.Integer)):
        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 is_nonelike(size):
        return lambda df, size: np.random.standard_t(df)
    if is_empty_tuple(size):
        return lambda df, size: np.array(np.random.standard_t(df))
    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 is_nonelike(size):
        return lambda mean, scale, size: np.random.wald(mean, scale)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda mean, scale, size: np.array(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)))
                
                if (X > LONG_MAX or X < 1.0):
                    continue

                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 is_nonelike(size):
        return lambda a, size: np.random.zipf(a)
    if is_empty_tuple(size):
        # Handle size = ()
        return lambda a, size: np.array(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(x, rng):

    if not isinstance(x, types.Buffer):
        msg = "The argument to shuffle() should be a buffer type"
        raise NumbaTypeError(msg)

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

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

    return impl


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


@overload(np.random.shuffle)
def shuffle_impl(x):
    return do_shuffle_impl(x, "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


# ------------------------------------------------------------------------
# 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 NumbaTypeError("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 NumbaTypeError("np.random.multinomial(): n should be an "
                             "integer, got %s" % (n,))

    if not isinstance(pvals, (types.Sequence, types.Array)):
        raise NumbaTypeError("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 NumbaTypeError("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) or is_empty_tuple(size):

        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.Float, types.Integer)) and isinstance(
            nonc, (types.Float, types.Integer)):
        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
     if is_empty_tuple(size):
         # Handle size = ()
         def noncentral_chisquare_impl(df, nonc, size=None):
             validate_noncentral_chisquare_input(df, nonc)
             return np.array(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")