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


import math
import random

from llvmlite import ir

from numba.core.extending import intrinsic, overload, register_jitable
from numba.core.imputils import Registry
from numba.core.typing import signature
from numba.core import types, cgutils
from numba.core.errors import NumbaTypeError
from collections import namedtuple

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)

GuassArgs = namedtuple('GuassArgs', ['state', 'random_function'])
cpython_guass_args = GuassArgs("py", random.random)


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_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("captures(none)")
    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


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(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(cpython_guass_args, loc_preprocessor, scale_preprocessor)
        return lambda mu, sigma: _impl(mu, sigma)



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_args, 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 = state_args.random_function

        state_ptr = get_state_ptr(context, builder, state_args.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(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)



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(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(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)


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(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)


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(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

def do_shuffle_impl(x, rand):

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

    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] = x[j].copy(), x[i].copy()
                i -= 1

    return impl

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


@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(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(random.vonmisesvariate)
def vonmisesvariate_impl(mu, kappa):
    if isinstance(mu, types.Float) and isinstance(kappa, types.Float):
        return _vonmisesvariate_impl(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