import pickle
from functools import partial

import numpy as np
import scipy.stats as stats
from pytest import fixture, importorskip, raises

import theano.tensor as tt
from theano import change_flags, config
from theano.gof.fg import FunctionGraph
from theano.gof.graph import Variable, graph_inputs
from theano.gof.op import get_test_value
from theano.tensor.random.basic import (
    bernoulli,
    beta,
    betabinom,
    binomial,
    categorical,
    cauchy,
    choice,
    dirichlet,
    exponential,
    gamma,
    halfcauchy,
    halfnormal,
    invgamma,
    multinomial,
    multivariate_normal,
    nbinom,
    normal,
    permutation,
    poisson,
    polyagamma,
    randint,
    truncexpon,
    uniform,
)


@fixture(scope="module", autouse=True)
def set_theano_flags():
    with change_flags(cxx="", compute_test_value="raise"):
        yield


def rv_numpy_tester(rv, *params, **kwargs):
    """Test for correspondence between `RandomVariable` and NumPy shape and
    broadcast dimensions.
    """
    test_fn = kwargs.pop("test_fn", None)

    if test_fn is None:
        name = getattr(rv, "name", None)

        if name is None:
            name = rv.__name__

        test_fn = getattr(np.random, name)

    theano_res = rv(*params, **kwargs)

    param_vals = [get_test_value(p) if isinstance(p, Variable) else p for p in params]
    kwargs_vals = {
        k: get_test_value(v) if isinstance(v, Variable) else v
        for k, v in kwargs.items()
    }

    if "size" in kwargs:
        kwargs["size"] = get_test_value(kwargs["size"])

    numpy_res = np.asarray(test_fn(*param_vals, **kwargs_vals))

    assert theano_res.type.numpy_dtype.kind == numpy_res.dtype.kind

    numpy_shape = np.shape(numpy_res)
    numpy_bcast = [s == 1 for s in numpy_shape]
    np.testing.assert_array_equal(theano_res.type.broadcastable, numpy_bcast)

    theano_res_val = theano_res.get_test_value()
    np.testing.assert_array_equal(theano_res_val.shape, numpy_res.shape)


def test_uniform_samples():

    rv_numpy_tester(uniform)
    rv_numpy_tester(uniform, size=())

    test_low = np.array(10, dtype=config.floatX)
    test_high = np.array(20, dtype=config.floatX)

    rv_numpy_tester(uniform, test_low, test_high)
    rv_numpy_tester(uniform, test_low, test_high, size=[3])


def test_beta_samples():

    test_a = np.array(0.5, dtype=config.floatX)
    test_b = np.array(0.5, dtype=config.floatX)

    rv_numpy_tester(beta, test_a, test_b)
    rv_numpy_tester(beta, test_a, test_b, size=[3])


def test_normal_infer_shape():
    M_tt = tt.iscalar("M")
    M_tt.tag.test_value = 3
    sd_tt = tt.scalar("sd")
    sd_tt.tag.test_value = np.array(1.0, dtype=config.floatX)

    test_params = [
        ([tt.as_tensor_variable(np.array(1.0, dtype=config.floatX)), sd_tt], None),
        ([tt.as_tensor_variable(np.array(1.0, dtype=config.floatX)), sd_tt], (M_tt,)),
        ([tt.as_tensor_variable(np.array(1.0, dtype=config.floatX)), sd_tt], (2, M_tt)),
        ([tt.zeros((M_tt,)), sd_tt], None),
        ([tt.zeros((M_tt,)), sd_tt], (M_tt,)),
        ([tt.zeros((M_tt,)), sd_tt], (2, M_tt)),
        ([tt.zeros((M_tt,)), tt.ones((M_tt,))], None),
        ([tt.zeros((M_tt,)), tt.ones((M_tt,))], (2, M_tt)),
        (
            [
                np.array([[-1, 20], [300, -4000]], dtype=config.floatX),
                np.array([[1e-6, 2e-6]], dtype=config.floatX),
            ],
            (3, 2, 2),
        ),
        (
            [np.array([1], dtype=config.floatX), np.array([10], dtype=config.floatX)],
            (1, 2),
        ),
    ]
    for args, size in test_params:
        rv = normal(*args, size=size)
        rv_shape = tuple(normal._infer_shape(size or (), args, None))
        assert tuple(get_test_value(rv_shape)) == tuple(get_test_value(rv).shape)


def test_normal_ShapeFeature():
    M_tt = tt.iscalar("M")
    M_tt.tag.test_value = 3
    sd_tt = tt.scalar("sd")
    sd_tt.tag.test_value = np.array(1.0, dtype=config.floatX)

    d_rv = normal(tt.ones((M_tt,)), sd_tt, size=(2, M_tt))
    d_rv.tag.test_value

    fg = FunctionGraph(
        [i for i in graph_inputs([d_rv]) if not isinstance(i, tt.Constant)],
        [d_rv],
        clone=False,
        features=[tt.opt.ShapeFeature()],
    )
    s1, s2 = fg.shape_feature.shape_of[d_rv]

    assert get_test_value(s1) == get_test_value(d_rv).shape[0]
    assert get_test_value(s2) == get_test_value(d_rv).shape[1]


def test_normal_samples():

    rv_numpy_tester(normal)

    test_mean = np.array(0, dtype=config.floatX)
    test_stddev = np.array(1, dtype=config.floatX)

    rv_numpy_tester(normal, test_mean, test_stddev)
    rv_numpy_tester(normal, test_mean, test_stddev, size=[3])

    # Broadcast sd over independent means...
    test_mean = np.array([0, 1, 2], dtype=config.floatX)
    test_stddev = np.array(1, dtype=config.floatX)
    rv_numpy_tester(normal, test_mean, test_stddev)
    rv_numpy_tester(normal, test_mean, test_stddev, size=[3, 3])

    test_mean = np.array([0], dtype=config.floatX)
    test_stddev = np.array([1], dtype=config.floatX)
    rv_numpy_tester(normal, test_mean, test_stddev, size=[1])
    rv_numpy_tester(normal, tt.as_tensor(test_mean), test_stddev, size=[1])
    rv_numpy_tester(
        normal,
        tt.as_tensor_variable(test_mean),
        test_stddev,
        size=tt.as_tensor_variable([1]),
    )


def test_halfnormal_samples():
    test_mean = np.array(0, dtype=config.floatX)
    test_stddev = np.array(1, dtype=config.floatX)

    rv_numpy_tester(halfnormal, test_fn=stats.halfnorm.rvs)
    rv_numpy_tester(halfnormal, test_mean, test_stddev, test_fn=stats.halfnorm.rvs)
    rv_numpy_tester(
        halfnormal,
        test_mean,
        test_stddev,
        size=[2, 3],
        test_fn=stats.halfnorm.rvs,
    )


def test_gamma_samples():
    test_a = np.array(0.5, dtype=config.floatX)
    test_b = np.array(0.5, dtype=config.floatX)

    rv_numpy_tester(gamma, test_a, test_b, test_fn=stats.gamma.rvs)
    rv_numpy_tester(gamma, test_a, test_b, size=[2, 3], test_fn=stats.gamma.rvs)


def test_exponential_samples():

    rv_numpy_tester(exponential)

    test_lambda = np.array(10, dtype=config.floatX)

    rv_numpy_tester(exponential, test_lambda)
    rv_numpy_tester(exponential, test_lambda, size=[2, 3])


def test_mvnormal_samples():
    def test_fn(mean=None, cov=None, size=None, rng=None):
        if mean is None:
            mean = np.array([0.0], dtype=config.floatX)
        if cov is None:
            cov = np.array([[1.0]], dtype=config.floatX)
        if size is None:
            size = ()
        return multivariate_normal.rng_fn(rng, mean, cov, size)

    rv_numpy_tester(multivariate_normal, test_fn=test_fn)

    test_mean = np.array([0], dtype=config.floatX)
    test_covar = np.diag(np.array([1], dtype=config.floatX))
    rv_numpy_tester(multivariate_normal, test_mean, test_covar, test_fn=test_fn)
    rv_numpy_tester(
        multivariate_normal, test_mean, test_covar, size=[1], test_fn=test_fn
    )
    rv_numpy_tester(
        multivariate_normal, test_mean, test_covar, size=[4], test_fn=test_fn
    )
    rv_numpy_tester(
        multivariate_normal, test_mean, test_covar, size=[4, 1], test_fn=test_fn
    )
    rv_numpy_tester(
        multivariate_normal, test_mean, test_covar, size=[4, 1, 1], test_fn=test_fn
    )
    rv_numpy_tester(
        multivariate_normal, test_mean, test_covar, size=[1, 5, 8], test_fn=test_fn
    )
    test_mean = np.array(
        [0, 1, 2],
        dtype=config.floatX,
    )
    test_covar = np.diag(np.array([1, 10, 100], dtype=config.floatX))
    rv_numpy_tester(multivariate_normal, test_mean, test_covar, test_fn=test_fn)

    # Test parameter broadcasting
    rv_numpy_tester(
        multivariate_normal,
        np.array([[0, 1, 2], [4, 5, 6]], dtype=config.floatX),
        test_covar,
        test_fn=test_fn,
    )

    test_covar = np.stack([test_covar, test_covar * 10.0])
    rv_numpy_tester(
        multivariate_normal,
        np.array([0, 1, 2], dtype=config.floatX),
        test_covar,
        size=[2, 3],
        test_fn=test_fn,
    )

    test_covar = np.stack([test_covar, test_covar * 10.0])
    rv_numpy_tester(
        multivariate_normal,
        np.array([[0, 1, 2]], dtype=config.floatX),
        test_covar,
        size=[2, 3],
        test_fn=test_fn,
    )

    rv_numpy_tester(
        multivariate_normal,
        np.array([[0], [10], [100]], dtype=config.floatX),
        np.diag(np.array([1e-6], dtype=config.floatX)),
        size=[2, 3],
        test_fn=test_fn,
    )


def test_mvnormal_ShapeFeature():
    M_tt = tt.iscalar("M")
    M_tt.tag.test_value = 2

    d_rv = multivariate_normal(tt.ones((M_tt,)), tt.eye(M_tt), size=2)

    fg = FunctionGraph(
        [i for i in graph_inputs([d_rv]) if not isinstance(i, tt.Constant)],
        [d_rv],
        clone=False,
        features=[tt.opt.ShapeFeature()],
    )

    s1, s2 = fg.shape_feature.shape_of[d_rv]

    assert get_test_value(s1) == 2
    assert M_tt in graph_inputs([s2])

    # Test broadcasted shapes
    mean = tt.tensor(config.floatX, [True, False])
    mean.tag.test_value = np.array([[0, 1, 2]], dtype=config.floatX)

    test_covar = np.diag(np.array([1, 10, 100], dtype=config.floatX))
    test_covar = np.stack([test_covar, test_covar * 10.0])
    cov = tt.as_tensor(test_covar).type()
    cov.tag.test_value = test_covar

    d_rv = multivariate_normal(mean, cov, size=[2, 3])

    fg = FunctionGraph(
        [i for i in graph_inputs([d_rv]) if not isinstance(i, tt.Constant)],
        [d_rv],
        clone=False,
        features=[tt.opt.ShapeFeature()],
    )

    s1, s2, s3, s4 = fg.shape_feature.shape_of[d_rv]

    assert s1.get_test_value() == 2
    assert s2.get_test_value() == 3
    assert s3.get_test_value() == 2
    assert s4.get_test_value() == 3


def test_dirichlet_samples():

    alphas = np.array([[100, 1, 1], [1, 100, 1], [1, 1, 100]], dtype=config.floatX)

    res = get_test_value(dirichlet(alphas))
    assert np.all(np.diag(res) >= res)

    res = get_test_value(dirichlet(alphas, size=2))
    assert res.shape == (2, 3, 3)
    assert all(np.all(np.diag(r) >= r) for r in res)

    for i in range(alphas.shape[0]):
        res = get_test_value(dirichlet(alphas[i]))
        assert np.all(res[i] > np.delete(res, [i]))

        res = get_test_value(dirichlet(alphas[i], size=2))
        assert res.shape == (2, 3)
        assert all(np.all(r[i] > np.delete(r, [i])) for r in res)

    rng_state = np.random.RandomState(np.random.MT19937(np.random.SeedSequence(1234)))

    alphas = np.array([[1000, 1, 1], [1, 1000, 1], [1, 1, 1000]], dtype=config.floatX)

    assert dirichlet.rng_fn(rng_state, alphas, None).shape == alphas.shape
    assert dirichlet.rng_fn(rng_state, alphas, size=10).shape == (10,) + alphas.shape
    assert (
        dirichlet.rng_fn(rng_state, alphas, size=(10, 2)).shape
        == (10, 2) + alphas.shape
    )


def test_dirichlet_infer_shape():
    M_tt = tt.iscalar("M")
    M_tt.tag.test_value = 3

    test_params = [
        ([tt.ones((M_tt,))], None),
        ([tt.ones((M_tt,))], (M_tt + 1,)),
        ([tt.ones((M_tt,))], (2, M_tt)),
        ([tt.ones((M_tt, M_tt + 1))], None),
        ([tt.ones((M_tt, M_tt + 1))], (M_tt + 2,)),
        ([tt.ones((M_tt, M_tt + 1))], (2, M_tt + 2, M_tt + 3)),
    ]
    for args, size in test_params:
        rv = dirichlet(*args, size=size)
        rv_shape = tuple(dirichlet._infer_shape(size or (), args, None))
        assert tuple(get_test_value(rv_shape)) == tuple(get_test_value(rv).shape)


def test_dirichlet_ShapeFeature():
    """Make sure `RandomVariable.infer_shape` works with `ShapeFeature`."""
    M_tt = tt.iscalar("M")
    M_tt.tag.test_value = 2
    N_tt = tt.iscalar("N")
    N_tt.tag.test_value = 3

    d_rv = dirichlet(tt.ones((M_tt, N_tt)), name="Gamma")

    fg = FunctionGraph(
        [i for i in graph_inputs([d_rv]) if not isinstance(i, tt.Constant)],
        [d_rv],
        clone=False,
        features=[tt.opt.ShapeFeature()],
    )

    s1, s2 = fg.shape_feature.shape_of[d_rv]

    assert M_tt in graph_inputs([s1])
    assert N_tt in graph_inputs([s2])


def test_poisson_samples():

    rv_numpy_tester(poisson)
    rv_numpy_tester(poisson, size=tt.as_tensor((2, 3)))

    test_lambda = np.array(10, dtype="int64")

    rv_numpy_tester(poisson, test_lambda)
    rv_numpy_tester(poisson, test_lambda, size=[2, 3])


def test_cauchy_samples():
    rv_numpy_tester(cauchy, test_fn=stats.cauchy.rvs)

    test_loc = np.array(10, dtype=config.floatX)
    test_scale = np.array(0.1, dtype=config.floatX)

    rv_numpy_tester(cauchy, test_loc, test_scale, test_fn=stats.cauchy.rvs)
    rv_numpy_tester(cauchy, test_loc, test_scale, size=[2, 3], test_fn=stats.cauchy.rvs)


def test_halfcauchy_samples():
    rv_numpy_tester(halfcauchy, test_fn=stats.halfcauchy.rvs)

    test_loc = np.array(10, dtype=config.floatX)
    test_scale = np.array(0.1, dtype=config.floatX)

    rv_numpy_tester(halfcauchy, test_loc, test_scale, test_fn=stats.halfcauchy.rvs)
    rv_numpy_tester(
        halfcauchy,
        test_loc,
        test_scale,
        size=[2, 3],
        test_fn=stats.halfcauchy.rvs,
    )


def test_invgamma_samples():
    test_loc = np.array(2, dtype=config.floatX)
    test_scale = np.array(2, dtype=config.floatX)

    rv_numpy_tester(
        invgamma, test_loc, test_scale, test_fn=partial(invgamma.rng_fn, None)
    )
    rv_numpy_tester(
        invgamma,
        test_loc,
        test_scale,
        size=[2, 3],
        test_fn=partial(invgamma.rng_fn, None),
    )


def test_truncexpon_samples():
    test_b = np.array(5, dtype=config.floatX)
    test_loc = np.array(0, dtype=config.floatX)
    test_scale = np.array(1, dtype=config.floatX)

    rv_numpy_tester(
        truncexpon,
        test_b,
        test_loc,
        test_scale,
        test_fn=partial(truncexpon.rng_fn, None),
    )
    rv_numpy_tester(
        truncexpon,
        test_b,
        test_loc,
        test_scale,
        size=[2, 3],
        test_fn=partial(truncexpon.rng_fn, None),
    )


def test_bernoulli_samples():
    test_p = np.array(0.5, dtype=config.floatX)

    rv_numpy_tester(bernoulli, test_p, test_fn=partial(bernoulli.rng_fn, None))
    rv_numpy_tester(
        bernoulli,
        test_p,
        size=[2, 3],
        test_fn=partial(bernoulli.rng_fn, None),
    )


def test_binomial_samples():
    test_M = np.array(10, dtype="int64")
    test_p = np.array(0.5, dtype=config.floatX)

    rv_numpy_tester(binomial, test_M, test_p)
    rv_numpy_tester(binomial, test_M, test_p, size=[2, 3])


def test_nbinom_samples():
    test_M = np.array(10, dtype="int64")
    test_p = np.array(0.5, dtype=config.floatX)

    rv_numpy_tester(nbinom, test_M, test_p, test_fn=partial(nbinom.rng_fn, None))
    rv_numpy_tester(
        nbinom,
        test_M,
        test_p,
        size=[2, 3],
        test_fn=partial(nbinom.rng_fn, None),
    )


def test_betabinom_samples():
    test_M = np.array(10, dtype="int64")
    test_a = np.array(0.5, dtype=config.floatX)
    test_b = np.array(0.5, dtype=config.floatX)

    rv_numpy_tester(
        betabinom, test_M, test_a, test_b, test_fn=partial(betabinom.rng_fn, None)
    )
    rv_numpy_tester(
        betabinom,
        test_M,
        test_a,
        test_b,
        size=[2, 3],
        test_fn=partial(betabinom.rng_fn, None),
    )


def test_multinomial_samples():
    test_M = np.array(10, dtype="int64")
    test_p = np.array([0.7, 0.3], dtype=config.floatX)

    rv_numpy_tester(multinomial, test_M, test_p)
    rv_numpy_tester(
        multinomial,
        test_M,
        test_p,
        size=[2, 3],
    )


def test_categorical_samples():

    rng_state = np.random.RandomState(np.random.MT19937(np.random.SeedSequence(1234)))

    p = np.array([[100000, 1, 1], [1, 100000, 1], [1, 1, 100000]], dtype=config.floatX)
    p = p / p.sum(axis=-1)

    assert categorical.rng_fn(rng_state, p, size=None).shape == p.shape[:-1]

    with raises(ValueError):
        categorical.rng_fn(rng_state, p, size=10)

    assert categorical.rng_fn(rng_state, p, size=(10, 3)).shape == (10, 3)
    assert categorical.rng_fn(rng_state, p, size=(10, 2, 3)).shape == (10, 2, 3)

    res = categorical(p)
    assert np.array_equal(get_test_value(res), np.arange(3))

    res = categorical(p, size=(10, 3))
    exp_res = np.tile(np.arange(3), (10, 1))
    assert np.array_equal(get_test_value(res), exp_res)

    res = categorical(p, size=(10, 2, 3))
    exp_res = np.tile(np.arange(3), (10, 2, 1))
    assert np.array_equal(get_test_value(res), exp_res)


def test_polyagamma_samples():

    _ = importorskip("pypolyagamma")

    # Sampled values should be scalars
    a = np.array(1.1, dtype=config.floatX)
    b = np.array(-10.5, dtype=config.floatX)
    pg_rv = polyagamma(a, b)
    assert get_test_value(pg_rv).shape == ()

    pg_rv = polyagamma(a, b, size=[1])
    assert get_test_value(pg_rv).shape == (1,)

    pg_rv = polyagamma(a, b, size=[2, 3])
    bcast_smpl = get_test_value(pg_rv)
    assert bcast_smpl.shape == (2, 3)
    # Make sure they're not all equal
    assert np.all(np.abs(np.diff(bcast_smpl.flat)) > 0.0)

    a = np.array([1.1, 3], dtype=config.floatX)
    b = np.array(-10.5, dtype=config.floatX)
    pg_rv = polyagamma(a, b)
    bcast_smpl = get_test_value(pg_rv)
    assert bcast_smpl.shape == (2,)
    assert np.all(np.abs(np.diff(bcast_smpl.flat)) > 0.0)

    pg_rv = polyagamma(a, b, size=(3, 2))
    bcast_smpl = get_test_value(pg_rv)
    assert bcast_smpl.shape == (3, 2)
    assert np.all(np.abs(np.diff(bcast_smpl.flat)) > 0.0)


def test_random_integer_samples():

    rv_numpy_tester(randint, 10, None)
    rv_numpy_tester(randint, 0, 1)
    rv_numpy_tester(randint, 0, 1, size=[3])
    rv_numpy_tester(randint, [0, 1, 2], 5)
    rv_numpy_tester(randint, [0, 1, 2], 5, size=[3, 3])
    rv_numpy_tester(randint, [0], [5], size=[1])
    rv_numpy_tester(randint, tt.as_tensor_variable([-1]), [1], size=[1])
    rv_numpy_tester(
        randint, tt.as_tensor_variable([-1]), [1], size=tt.as_tensor_variable([1])
    )


def test_choice_samples():
    with raises(NotImplementedError):
        choice._shape_from_params(np.asarray(5))

    rv_numpy_tester(choice, np.asarray(5))
    rv_numpy_tester(choice, np.array([1.0, 5.0], dtype=config.floatX))
    rv_numpy_tester(choice, np.asarray(5), 3)

    with raises(ValueError):
        rv_numpy_tester(choice, np.array([[1, 2], [3, 4]]))

    rv_numpy_tester(choice, [1, 2, 3], 1)
    rv_numpy_tester(choice, [1, 2, 3], 1, p=tt.as_tensor([1 / 3.0, 1 / 3.0, 1 / 3.0]))
    rv_numpy_tester(choice, [1, 2, 3], (10, 2), replace=True)
    rv_numpy_tester(choice, tt.as_tensor_variable([1, 2, 3]), 2, replace=True)


def test_permutation_samples():
    rv_numpy_tester(
        permutation, np.asarray(5), test_fn=lambda x: np.random.permutation(x.item())
    )
    rv_numpy_tester(permutation, [1, 2, 3])
    rv_numpy_tester(permutation, [[1, 2], [3, 4]])
    rv_numpy_tester(permutation, np.array([1.0, 2.0, 3.0], dtype=config.floatX))


@change_flags(compute_test_value="off")
def test_pickle():
    # This is an interesting `Op` case, because it has `None` types and a
    # conditional dtype
    sample_a = choice(5, size=(2, 3))

    a_pkl = pickle.dumps(sample_a)
    a_unpkl = pickle.loads(a_pkl)

    assert a_unpkl.owner.op._props() == sample_a.owner.op._props()
