import numpy as np
import pytest
from scipy import stats, integrate

import numpy as np

import autofit.graphical.messages.normal
from autofit import graphical as graph
from autofit.graphical.messages import transform
from autofit.graphical.utils import numerical_jacobian


def check_dist_norm(dist):
    norm, err = integrate.quad(dist.pdf, *dist._support[0])
    assert norm == pytest.approx(1, abs=err), dist


def check_dist_norms(dist):
    vals, err = integrate.quad_vec(
        lambda x: dist.pdf(np.full(dist.shape, x)),
        *dist._support[0]
    )
    assert np.allclose(vals, 1, atol=err), dist


def check_log_normalisation(ms):
    m1, *m2s = ms
    A = np.exp(m1.log_normalisation(*m2s))

    # Calculate normalisation numerically
    i12, ierr = integrate.quad(
        lambda x: np.exp(sum(m.logpdf(x) for m in ms)),
        *m1._support[0])

    # verify within tolerance
    assert np.abs(A - i12) < ierr < 1e-1, m1 


def check_numerical_gradient_hessians(message, x=None):
    x = message.sample() if x is None else x 

    res = message.logpdf_gradient(x)
    nres = message.numerical_logpdf_gradient(x)
    for i, (x1, x2) in enumerate(zip(res, nres)):
        assert np.allclose(x1, x2, rtol=1e-3, atol=1e-2), (i, x1, x2, message)

    res = message.logpdf_gradient_hessian(x)
    nres = message.numerical_logpdf_gradient_hessian(x)
    for i, (x1, x2) in enumerate(zip(res, nres)):
        assert np.allclose(x1, x2, rtol=1e-3, atol=1e-2), (i, x1, x2, message)


def test_message_norm():
    messages = [
        tuple(
            map(graph.NormalMessage,
                [0.5, 0.1], [0.2, 0.3])),
        tuple(
            map(graph.NormalMessage,
                [0.5, 0.1, -0.5], [0.2, 0.3, 1.3])),
        tuple(
            map(graph.GammaMessage,
                [0.5, 1.1], [0.2, 1.3])),
        tuple(
            map(graph.GammaMessage,
                [0.5, 1.1, 2], [0.2, 1.3, 1])),
        tuple(
            map(graph.BetaMessage,
                [2., 3.2, 1.5], [4.1, 2.3, 3])),
    ]
    for ms in messages:
        check_log_normalisation(ms)
        for m in ms:
            check_dist_norm(m)
            check_numerical_gradient_hessians(m)


def test_numerical_gradient_hessians():
    N = graph.NormalMessage
    UN = graph.UniformNormalMessage
    SUN = UN.shifted(shift=0.3, scale=0.8)
    LN = graph.LogNormalMessage
    MLN = graph.MultiLogitNormalMessage
    # test doubly transformed distributions
    WN = graph.NormalMessage.transformed(
        transform.log_transform
    ).transformed(
        transform.exp_transform,
    )
    test_cases = [
        (N, 1., 0.5, 0.3),
        (N, 1., 0.5, [0.3, 2.1]),
        (N, [0.1, 1., 2.], [2., 0.5, 3.], [0.1, 0.2, 0.3]),
        (N, [0.1, 1., 2.], [2., 0.5, 3.], [[0.1, 0.2, 0.3], [2., 1., -1]]),
        (UN, 1., 0.5, None),
        (UN, [0.1, 1., 2.], [.1, 0.5, 0.2], None),
        (SUN, 1., 0.5, None),
        (SUN, [0.1, 1., 2.], [.1, 0.5, 0.2], None),
        (LN, 1., 0.5, None),
        (LN, [0.1, 1., 2.], [2., 0.5, 3.], None),
        (LN, [0.1, 1., 2.], [2., 0.5, 3.], None),
        (MLN, [0.1, 1., 2.], [.1, 0.5, 0.2], None),
        (WN, 1., 0.5, 0.3),
        (WN, 1., 0.5, [0.3, 2.1]),
        (WN, [0.1, 1., 2.], [2., 0.5, 3.], [0.1, 0.2, 0.3]),
        (WN, [0.1, 1., 2.], [2., 0.5, 3.], [[0.1, 0.2, 0.3], [2., 1., -1]]),
    ]
    for M, m, s, x in test_cases:
        check_numerical_gradient_hessians(M(m, s), x)

def test_meanfield_gradients():
    n1, n2, n3 = 2, 3, 5
    p1, p2, p3 = [graph.Plate() for i in range(3)]

    v1 = graph.Variable('v1', p1, p2)
    v2 = graph.Variable('v2', p2, p3)
    v3 = graph.Variable('v3', p3, p1)

    mean_field = graph.MeanField({
        v1: graph.NormalMessage(
            np.random.randn(n1, n2),
            np.random.exponential(size=(n1, n2))),
        v2: graph.NormalMessage(
            np.random.randn(n2, n3),
            np.random.exponential(size=(n2, n3))),
        v3: graph.NormalMessage(
            np.random.randn(n3, n1),
            np.random.exponential(size=(n3, n1)))})

    values = mean_field.sample()
    l0 = mean_field(values, axis=None)
    logl = mean_field(values, axis=False)
    assert logl.sum() == pytest.approx(l0, abs=1e-5)
    logl = mean_field(values, axis=1)
    assert logl.sum() == pytest.approx(l0, abs=1e-5)
    logl = mean_field(values, axis=2)
    assert logl.sum() == pytest.approx(l0, abs=1e-5)
    logl = mean_field(values, axis=(0, 2))
    assert logl.sum() == pytest.approx(l0, abs=1e-5)

    njac0 = mean_field._numerical_func_jacobian(
        values, axis=None, _eps=1e-8)[1]
    njac1 = mean_field._numerical_func_jacobian(
        values, axis=1, _eps=1e-8)[1]
    njac2 = mean_field._numerical_func_jacobian(
        values, axis=(0, 1), _eps=1e-8)[1]
    njac = mean_field._numerical_func_jacobian_hessian(
        values, axis=False, _eps=1e-8)[1]
    grad = mean_field.logpdf_gradient(values, axis=False)[1]
    for v in grad:
        norm = np.linalg.norm(grad[v] - njac[v].sum((0, 1, 2)))
        assert norm == pytest.approx(0, abs=1e-2)
        norm = np.linalg.norm(grad[v] - njac0[v])
        assert norm == pytest.approx(0, abs=1e-2)
        norm = np.linalg.norm(grad[v] - njac1[v].sum((0, 1)))
        assert norm == pytest.approx(0, abs=1e-2)
        norm = np.linalg.norm(grad[v] - njac2[v].sum(0))
        assert norm == pytest.approx(0, abs=1e-2)


def test_beta():
    a = b = np.r_[2., 3.2, 1.5]
    beta = graph.BetaMessage(a, b[::-1])
    check_dist_norms(beta)

    betas = [
        graph.BetaMessage(a, b)
        for a, b in (np.random.poisson(5, size=(10, 2)) + 1)
    ]
    for b in betas:
        check_dist_norm(b)

    check_log_normalisation(betas)


def check_transforms(transform, x):
    y, logd, logd_grad, jac = transform.transform_det_jac(x)

    njac = numerical_jacobian(x, transform.transform)
    nlogd_grad = numerical_jacobian(x, transform.log_det)

    assert np.allclose(transform.inv_transform(y), x)
    assert np.allclose(jac.to_dense(), njac)
    assert np.isclose(logd.sum(), np.linalg.slogdet(njac)[1])
    assert np.allclose(nlogd_grad.sum(0), logd_grad)


def test_transforms():
    tests = [
        (transform.log_transform, np.r_[10, 2, 0.1]),
        (transform.exp_transform, np.r_[10, -2, 0.1]),
        (transform.logistic_transform, np.r_[0.22, 0.51, 0.1]),
        (transform.phi_transform, np.r_[0.22, 0.51, 0.1]),
        (transform.shifted_logistic(shift=11, scale=5.1), np.r_[11.1, 12, 16]),
    ]
    for args in tests:
        check_transforms(*args)


def test_multinomial_logit():
    mult_logit = transform.multinomial_logit_transform

    d = 3
    p = np.random.dirichlet(np.ones(d+1))[:d]

    x, logd, logd_grad, jac = mult_logit.transform_det_jac(p)

    njac = numerical_jacobian(p, mult_logit.transform)
    nlogd_grad = numerical_jacobian(p, mult_logit.log_det)

    assert np.allclose(mult_logit.inv_transform(x), p)
    assert np.allclose(jac.to_dense(), njac)
    assert np.isclose(logd.sum(), np.linalg.slogdet(njac)[1])
    assert np.allclose(nlogd_grad.sum(0), logd_grad)

    n = 5

    ps = np.random.dirichlet(np.ones(d+1), size=n)[:, :d]
    xs, logd, logd_grad, jac = mult_logit.transform_det_jac(ps)
    njac = numerical_jacobian(ps, mult_logit.transform).reshape(jac.shape)
    nlogd_grad = numerical_jacobian(ps, mult_logit.log_det)

    assert np.allclose(mult_logit.inv_transform(xs), ps)
    assert np.allclose(njac, jac.to_dense())
    assert np.isclose(
        logd.sum(), 
        np.linalg.slogdet(njac.reshape(jac.lsize, jac.rsize))[1]
    )
    assert np.allclose(nlogd_grad.sum((0, 1)), logd_grad, 1e-5, 1e-3)

    assert np.allclose(xs[0], mult_logit.transform(ps[0]))
    assert np.allclose(logd[0], mult_logit.log_det(ps[0]))
    assert np.allclose(logd_grad[0], mult_logit.log_det_grad(ps[0])[1])


def test_normal_simplex():
    mult_logit = transform.MultinomialLogitTransform()
    NormalSimplex = graph.messages.NormalMessage.transformed(mult_logit)

    message = NormalSimplex([-1, 2], [.3, .3])
    
    check_numerical_gradient_hessians(message, message.sample())

    def func(*p):
        return message.pdf(p).prod()

    def simplex_lims(*args):
        return [0, 1 - sum(args)]

    # verify transformation normalises correctly
    res, err = integrate.nquad(
        func,
        [simplex_lims] * message.size
    )
    assert res == pytest.approx(1, rel=err)
