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)

[((297, 340), 'collections.OrderedDict', 'OrderedDict', (["[('a', 2), ('b', 2), ('c', 2)]"], {}), "([('a', 2), ('b', 2), ('c', 2)])\n", (308, 340), False, 'from collections import OrderedDict\n'), ((794, 837), 'collections.OrderedDict', 'OrderedDict', (["[('a', 2), ('b', 2), ('c', 2)]"], {}), "([('a', 2), ('b', 2), ('c', 2)])\n", (805, 837), False, 'from collections import OrderedDict\n'), ((851, 885), 'torch.tensor', 'torch.tensor', (['[1.0, 2, 3, 4, 5, 6]'], {}), '([1.0, 2, 3, 4, 5, 6])\n', (863, 885), False, 'import torch\n'), ((1424, 1440), 'pyro.contrib.util.lexpand', 'lexpand', (['A', '(5)', '(4)'], {}), '(A, 5, 4)\n', (1431, 1440), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((1457, 1473), 'pyro.contrib.util.lexpand', 'lexpand', (['b', '(5)', '(4)'], {}), '(b, 5, 4)\n', (1464, 1473), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((1772, 1788), 'pyro.contrib.util.lexpand', 'lexpand', (['a', '(5)', '(4)'], {}), '(a, 5, 4)\n', (1779, 1788), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((1805, 1821), 'pyro.contrib.util.lexpand', 'lexpand', (['b', '(5)', '(4)'], {}), '(b, 5, 4)\n', (1812, 1821), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((1969, 2006), 'torch.tensor', 'torch.tensor', (['[[1.0, 2.0], [-2.0, 0]]'], {}), '([[1.0, 2.0], [-2.0, 0]])\n', (1981, 2006), False, 'import torch\n'), ((2207, 2244), 'torch.tensor', 'torch.tensor', (['[[1.0, 2.0], [-2.0, 0]]'], {}), '([[1.0, 2.0], [-2.0, 0]])\n', (2219, 2244), False, 'import torch\n'), ((2485, 2522), 'torch.tensor', 'torch.tensor', (['[[1.0, 2.0], [-2.0, 0]]'], {}), '([[1.0, 2.0], [-2.0, 0]])\n', (2497, 2522), False, 'import torch\n'), ((2588, 2604), 'pyro.contrib.util.lexpand', 'lexpand', (['A', '(5)', '(4)'], {}), '(A, 5, 4)\n', (2595, 2604), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((2732, 2762), 'torch.tensor', 'torch.tensor', (['[1.0, 2.0, -1.0]'], {}), '([1.0, 2.0, -1.0])\n', (2744, 2762), False, 'import torch\n'), ((2828, 2844), 'pyro.contrib.util.lexpand', 'lexpand', (['v', '(5)', '(4)'], {}), '(v, 5, 4)\n', (2835, 2844), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((2982, 3000), 'torch.randn', 'torch.randn', (['(3)', '(10)'], {}), '(3, 10)\n', (2993, 3000), False, 'import torch\n'), ((3104, 3138), 'torch.randn', 'torch.randn', (['(3)'], {'requires_grad': '(True)'}), '(3, requires_grad=True)\n', (3115, 3138), False, 'import torch\n'), ((3264, 3298), 'torch.randn', 'torch.randn', (['(3)'], {'requires_grad': '(True)'}), '(3, requires_grad=True)\n', (3275, 3298), False, 'import torch\n'), ((3307, 3341), 'torch.randn', 'torch.randn', (['(3)'], {'requires_grad': '(True)'}), '(3, requires_grad=True)\n', (3318, 3341), False, 'import torch\n'), ((3387, 3405), 'pyro.contrib.util.hessian', 'hessian', (['y', '(x, z)'], {}), '(y, (x, z))\n', (3394, 3405), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((3565, 3590), 'tests.common.assert_equal', 'assert_equal', (['H', 'target_H'], {}), '(H, target_H)\n', (3577, 3590), False, 'from tests.common import assert_equal\n'), ((358, 389), 'pyro.contrib.util.get_indices', 'get_indices', (["['b']"], {'sizes': 'sizes'}), "(['b'], sizes=sizes)\n", (369, 389), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((391, 411), 'torch.tensor', 'torch.tensor', (['[2, 3]'], {}), '([2, 3])\n', (403, 411), False, 'import torch\n'), ((430, 466), 'pyro.contrib.util.get_indices', 'get_indices', (["['b', 'c']"], {'sizes': 'sizes'}), "(['b', 'c'], sizes=sizes)\n", (441, 466), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((468, 494), 'torch.tensor', 'torch.tensor', (['[2, 3, 4, 5]'], {}), '([2, 3, 4, 5])\n', (480, 494), False, 'import torch\n'), ((607, 642), 'pyro.contrib.util.get_indices', 'get_indices', (["['b']"], {'tensors': 'tensors'}), "(['b'], tensors=tensors)\n", (618, 642), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((644, 664), 'torch.tensor', 'torch.tensor', (['[2, 3]'], {}), '([2, 3])\n', (656, 664), False, 'import torch\n'), ((683, 723), 'pyro.contrib.util.get_indices', 'get_indices', (["['b', 'c']"], {'tensors': 'tensors'}), "(['b', 'c'], tensors=tensors)\n", (694, 723), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((725, 751), 'torch.tensor', 'torch.tensor', (['[2, 3, 4, 5]'], {}), '([2, 3, 4, 5])\n', (737, 751), False, 'import torch\n'), ((902, 931), 'pyro.contrib.util.tensor_to_dict', 'tensor_to_dict', (['sizes', 'vector'], {}), '(sizes, vector)\n', (916, 931), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((1137, 1180), 'pyro.contrib.util.tensor_to_dict', 'tensor_to_dict', (['sizes', 'vector'], {'subset': "['b']"}), "(sizes, vector, subset=['b'])\n", (1151, 1180), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((1377, 1386), 'pyro.contrib.util.rmv', 'rmv', (['A', 'b'], {}), '(A, b)\n', (1380, 1386), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((1532, 1557), 'pyro.contrib.util.rmv', 'rmv', (['batched_A', 'batched_b'], {}), '(batched_A, batched_b)\n', (1535, 1557), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((1717, 1726), 'pyro.contrib.util.rvv', 'rvv', (['a', 'b'], {}), '(a, b)\n', (1720, 1726), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((1728, 1743), 'torch.dot', 'torch.dot', (['a', 'b'], {}), '(a, b)\n', (1737, 1743), False, 'import torch\n'), ((1848, 1863), 'torch.dot', 'torch.dot', (['a', 'b'], {}), '(a, b)\n', (1857, 1863), False, 'import torch\n'), ((1888, 1913), 'pyro.contrib.util.rvv', 'rvv', (['batched_a', 'batched_b'], {}), '(batched_a, batched_b)\n', (1891, 1913), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((2021, 2031), 'pyro.contrib.util.lexpand', 'lexpand', (['A'], {}), '(A)\n', (2028, 2031), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((2064, 2077), 'pyro.contrib.util.lexpand', 'lexpand', (['A', '(4)'], {}), '(A, 4)\n', (2071, 2077), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((2126, 2142), 'pyro.contrib.util.lexpand', 'lexpand', (['A', '(4)', '(2)'], {}), '(A, 4, 2)\n', (2133, 2142), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((2259, 2269), 'pyro.contrib.util.rexpand', 'rexpand', (['A'], {}), '(A)\n', (2266, 2269), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((2302, 2315), 'pyro.contrib.util.rexpand', 'rexpand', (['A', '(4)'], {}), '(A, 4)\n', (2309, 2315), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((2378, 2394), 'pyro.contrib.util.rexpand', 'rexpand', (['A', '(4)', '(2)'], {}), '(A, 4, 2)\n', (2385, 2394), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((2537, 2545), 'pyro.contrib.util.rtril', 'rtril', (['A'], {}), '(A)\n', (2542, 2545), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((2547, 2560), 'torch.tril', 'torch.tril', (['A'], {}), '(A)\n', (2557, 2560), False, 'import torch\n'), ((2628, 2641), 'torch.tril', 'torch.tril', (['A'], {}), '(A)\n', (2638, 2641), False, 'import torch\n'), ((2666, 2681), 'pyro.contrib.util.rtril', 'rtril', (['expanded'], {}), '(expanded)\n', (2671, 2681), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((2777, 2785), 'pyro.contrib.util.rdiag', 'rdiag', (['v'], {}), '(v)\n', (2782, 2785), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((2787, 2800), 'torch.diag', 'torch.diag', (['v'], {}), '(v)\n', (2797, 2800), False, 'import torch\n'), ((2869, 2882), 'torch.diag', 'torch.diag', (['v'], {}), '(v)\n', (2879, 2882), False, 'import torch\n'), ((2907, 2922), 'pyro.contrib.util.rdiag', 'rdiag', (['expanded'], {}), '(expanded)\n', (2912, 2922), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((3180, 3193), 'pyro.contrib.util.hessian', 'hessian', (['y', 'x'], {}), '(y, x)\n', (3187, 3193), False, 'from pyro.contrib.util import get_indices, tensor_to_dict, rmv, rvv, lexpand, rexpand, rdiag, rtril, hessian\n'), ((939, 963), 'torch.tensor', 'torch.tensor', (['[1.0, 2.0]'], {}), '([1.0, 2.0])\n', (951, 963), False, 'import torch\n'), ((1017, 1041), 'torch.tensor', 'torch.tensor', (['[3.0, 4.0]'], {}), '([3.0, 4.0])\n', (1029, 1041), False, 'import torch\n'), ((1095, 1119), 'torch.tensor', 'torch.tensor', (['[5.0, 6.0]'], {}), '([5.0, 6.0])\n', (1107, 1119), False, 'import torch\n'), ((1205, 1229), 'torch.tensor', 'torch.tensor', (['[3.0, 4.0]'], {}), '([3.0, 4.0])\n', (1217, 1229), False, 'import torch\n'), ((1271, 1310), 'torch.tensor', 'torch.tensor', (['[[1.0, 2.0], [2.0, -3.0]]'], {}), '([[1.0, 2.0], [2.0, -3.0]])\n', (1283, 1310), False, 'import torch\n'), ((1308, 1333), 'torch.tensor', 'torch.tensor', (['[-1.0, 2.0]'], {}), '([-1.0, 2.0])\n', (1320, 1333), False, 'import torch\n'), ((1624, 1648), 'torch.tensor', 'torch.tensor', (['[1.0, 2.0]'], {}), '([1.0, 2.0])\n', (1636, 1648), False, 'import torch\n'), ((1648, 1673), 'torch.tensor', 'torch.tensor', (['[-1.0, 2.0]'], {}), '([-1.0, 2.0])\n', (1660, 1673), False, 'import torch\n'), ((3507, 3528), 'torch.cat', 'torch.cat', (['[Hxx, Hxz]'], {}), '([Hxx, Hxz])\n', (3516, 3528), False, 'import torch\n'), ((3530, 3551), 'torch.cat', 'torch.cat', (['[Hxz, Hzz]'], {}), '([Hxz, Hzz])\n', (3539, 3551), False, 'import torch\n'), ((529, 542), 'torch.ones', 'torch.ones', (['(2)'], {}), '(2)\n', (539, 542), False, 'import torch\n'), ((551, 564), 'torch.ones', 'torch.ones', (['(2)'], {}), '(2)\n', (561, 564), False, 'import torch\n'), ((573, 586), 'torch.ones', 'torch.ones', (['(2)'], {}), '(2)\n', (583, 586), False, 'import torch\n')]