from collections import OrderedDict
import pytest
import torch

import pyro.distributions as dist
from pyro.contrib.util import (
    get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian
)
from tests.common import assert_equal


def test_get_indices_sizes():
    sizes = OrderedDict([("a", 2), ("b", 2), ("c", 2)])
    assert_equal(get_indices(["b"], sizes=sizes), torch.tensor([2, 3]))
    assert_equal(get_indices(["b", "c"], sizes=sizes), torch.tensor([2, 3, 4, 5]))
    tensors = OrderedDict([("a", torch.ones(2)), ("b", torch.ones(2)), ("c", torch.ones(2))])
    assert_equal(get_indices(["b"], tensors=tensors), torch.tensor([2, 3]))
    assert_equal(get_indices(["b", "c"], tensors=tensors), torch.tensor([2, 3, 4, 5]))


def test_tensor_to_dict():
    sizes = OrderedDict([("a", 2), ("b", 2), ("c", 2)])
    vector = torch.tensor([1., 2, 3, 4, 5, 6])
    assert_equal(tensor_to_dict(sizes, vector), {"a": torch.tensor([1., 2.]),
                                                 "b": torch.tensor([3., 4.]),
                                                 "c": torch.tensor([5., 6.])})
    assert_equal(tensor_to_dict(sizes, vector, subset=["b"]),
                 {"b": torch.tensor([3., 4.])})


@pytest.mark.parametrize("A,b", [
    (torch.tensor([[1., 2.], [2., -3.]]), torch.tensor([-1., 2.]))
    ])
def test_rmv(A, b):
    assert_equal(rmv(A, b), A.mv(b), prec=1e-8)
    batched_A = lexpand(A, 5, 4)
    batched_b = lexpand(b, 5, 4)
    expected_Ab = lexpand(A.mv(b), 5, 4)
    assert_equal(rmv(batched_A, batched_b), expected_Ab, prec=1e-8)


@pytest.mark.parametrize("a,b", [
    (torch.tensor([1., 2.]), torch.tensor([-1., 2.]))
    ])
def test_rvv(a, b):
    assert_equal(rvv(a, b), torch.dot(a, b), prec=1e-8)
    batched_a = lexpand(a, 5, 4)
    batched_b = lexpand(b, 5, 4)
    expected_ab = lexpand(torch.dot(a, b), 5, 4)
    assert_equal(rvv(batched_a, batched_b), expected_ab, prec=1e-8)


def test_lexpand():
    A = torch.tensor([[1., 2.], [-2., 0]])
    assert_equal(lexpand(A), A, prec=1e-8)
    assert_equal(lexpand(A, 4), A.expand(4, 2, 2), prec=1e-8)
    assert_equal(lexpand(A, 4, 2), A.expand(4, 2, 2, 2), prec=1e-8)


def test_rexpand():
    A = torch.tensor([[1., 2.], [-2., 0]])
    assert_equal(rexpand(A), A, prec=1e-8)
    assert_equal(rexpand(A, 4), A.unsqueeze(-1).expand(2, 2, 4), prec=1e-8)
    assert_equal(rexpand(A, 4, 2), A.unsqueeze(-1).unsqueeze(-1).expand(2, 2, 4, 2), prec=1e-8)


def test_rtril():
    A = torch.tensor([[1., 2.], [-2., 0]])
    assert_equal(rtril(A), torch.tril(A), prec=1e-8)
    expanded = lexpand(A, 5, 4)
    expected = lexpand(torch.tril(A), 5, 4)
    assert_equal(rtril(expanded), expected, prec=1e-8)


def test_rdiag():
    v = torch.tensor([1., 2., -1.])
    assert_equal(rdiag(v), torch.diag(v), prec=1e-8)
    expanded = lexpand(v, 5, 4)
    expeceted = lexpand(torch.diag(v), 5, 4)
    assert_equal(rdiag(expanded), expeceted, prec=1e-8)


def test_hessian_mvn():
    tmp = torch.randn(3, 10)
    cov = torch.matmul(tmp, tmp.t())
    mvn = dist.MultivariateNormal(cov.new_zeros(3), cov)

    x = torch.randn(3, requires_grad=True)
    y = mvn.log_prob(x)
    assert_equal(hessian(y, x), -mvn.precision_matrix)


def test_hessian_multi_variables():
    x = torch.randn(3, requires_grad=True)
    z = torch.randn(3, requires_grad=True)
    y = (x ** 2 * z + z ** 3).sum()

    H = hessian(y, (x, z))
    Hxx = (2 * z).diag()
    Hxz = (2 * x).diag()
    Hzz = (6 * z).diag()
    target_H = torch.cat([torch.cat([Hxx, Hxz]), torch.cat([Hxz, Hzz])], dim=1)
    assert_equal(H, target_H)

[(13, 'collections.OrderedDict', 'OrderedDict', False, 'from collections import OrderedDict\n'), (22, 'collections.OrderedDict', 'OrderedDict', False, 'from collections import OrderedDict\n'), (23, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (36, 'pyro.contrib.util.lexpand', 'lexpand', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (37, 'pyro.contrib.util.lexpand', 'lexpand', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (47, 'pyro.contrib.util.lexpand', 'lexpand', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (48, 'pyro.contrib.util.lexpand', 'lexpand', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (54, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (61, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (68, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (70, 'pyro.contrib.util.lexpand', 'lexpand', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (76, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (78, 'pyro.contrib.util.lexpand', 'lexpand', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (84, 'torch.randn', 'torch.randn', False, 'import torch\n'), (88, 'torch.randn', 'torch.randn', False, 'import torch\n'), (94, 'torch.randn', 'torch.randn', False, 'import torch\n'), (95, 'torch.randn', 'torch.randn', False, 'import torch\n'), (98, 'pyro.contrib.util.hessian', 'hessian', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (103, 'tests.common.assert_equal', 'assert_equal', False, 'from tests.common import assert_equal\n'), (14, 'pyro.contrib.util.get_indices', 'get_indices', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (14, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (15, 'pyro.contrib.util.get_indices', 'get_indices', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (15, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (17, 'pyro.contrib.util.get_indices', 'get_indices', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (17, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (18, 'pyro.contrib.util.get_indices', 'get_indices', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (18, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (24, 'pyro.contrib.util.tensor_to_dict', 'tensor_to_dict', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (27, 'pyro.contrib.util.tensor_to_dict', 'tensor_to_dict', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (35, 'pyro.contrib.util.rmv', 'rmv', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (39, 'pyro.contrib.util.rmv', 'rmv', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (46, 'pyro.contrib.util.rvv', 'rvv', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (46, 'torch.dot', 'torch.dot', False, 'import torch\n'), (49, 'torch.dot', 'torch.dot', False, 'import torch\n'), (50, 'pyro.contrib.util.rvv', 'rvv', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (55, 'pyro.contrib.util.lexpand', 'lexpand', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (56, 'pyro.contrib.util.lexpand', 'lexpand', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (57, 'pyro.contrib.util.lexpand', 'lexpand', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (62, 'pyro.contrib.util.rexpand', 'rexpand', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (63, 'pyro.contrib.util.rexpand', 'rexpand', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (64, 'pyro.contrib.util.rexpand', 'rexpand', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (69, 'pyro.contrib.util.rtril', 'rtril', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (69, 'torch.tril', 'torch.tril', False, 'import torch\n'), (71, 'torch.tril', 'torch.tril', False, 'import torch\n'), (72, 'pyro.contrib.util.rtril', 'rtril', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (77, 'pyro.contrib.util.rdiag', 'rdiag', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (77, 'torch.diag', 'torch.diag', False, 'import torch\n'), (79, 'torch.diag', 'torch.diag', False, 'import torch\n'), (80, 'pyro.contrib.util.rdiag', 'rdiag', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (90, 'pyro.contrib.util.hessian', 'hessian', False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), (24, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (25, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (26, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (28, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (32, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (32, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (43, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (43, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (102, 'torch.cat', 'torch.cat', False, 'import torch\n'), (102, 'torch.cat', 'torch.cat', False, 'import torch\n'), (16, 'torch.ones', 'torch.ones', False, 'import torch\n'), (16, 'torch.ones', 'torch.ones', False, 'import torch\n'), (16, 'torch.ones', 'torch.ones', False, 'import torch\n')]