import torch

from ceem.opt_criteria import *
from ceem.systems import LorenzAttractor
from ceem.dynamics import *
from ceem.smoother import *
from ceem import utils


def test_smoother():

    utils.set_rng_seed(1)

    torch.set_default_dtype(torch.float64)

    sigma = torch.tensor([10.])
    rho = torch.tensor([28.])
    beta = torch.tensor([8. / 3.])

    C = torch.randn(2, 3)

    dt = 0.04

    sys = LorenzAttractor(sigma, rho, beta, C, dt, method='midpoint')

    B = 1
    T = 200
    xs = [torch.randn(B, 1, 3)]
    for t in range(T - 1):
        xs.append(sys.step(torch.tensor([0.] * B), xs[-1]))

    x = torch.cat(xs, dim=1).detach()
    x.requires_grad = True
    y = sys.observe(0., x).detach()
    # y += torch.rand_like(y) * 0.01

    t = torch.stack([torch.arange(T), torch.arange(T)]).to(torch.get_default_dtype())
    
    x0 = torch.zeros_like(x)

    obscrit = GaussianObservationCriterion(torch.ones(2), t, y)

    dyncrit = GaussianDynamicsCriterion(torch.ones(3), t)

    # Test GroupSOSCriterion
    crit = GroupSOSCriterion([obscrit, dyncrit])

    xsm, metrics = NLSsmoother(x0, crit, sys, solver_kwargs={'verbose': 2, 'tr_rho': 0.})

    err = float((xsm - x).norm())
    assert err < 1e-8, 'Smoothing Error: %.3e' % err

    print('Passed.')

    # Test BlockSparseGroupSOSCriterion
    crit = BlockSparseGroupSOSCriterion([obscrit, dyncrit])

    xsm, metrics = NLSsmoother(torch.zeros_like(x), crit, sys)

    err = float((xsm - x).norm())
    assert err < 1e-8, 'Smoothing Error: %.3e' % err

    print('Passed.')

if __name__ == '__main__':
    test_smoother()

[((194, 215), 'ceem.utils.set_rng_seed', 'utils.set_rng_seed', (['(1)'], {}), '(1)\n', (212, 215), False, 'from ceem import utils\n'), ((221, 259), 'torch.set_default_dtype', 'torch.set_default_dtype', (['torch.float64'], {}), '(torch.float64)\n', (244, 259), False, 'import torch\n'), ((273, 293), 'torch.tensor', 'torch.tensor', (['[10.0]'], {}), '([10.0])\n', (285, 293), False, 'import torch\n'), ((303, 323), 'torch.tensor', 'torch.tensor', (['[28.0]'], {}), '([28.0])\n', (315, 323), False, 'import torch\n'), ((334, 359), 'torch.tensor', 'torch.tensor', (['[8.0 / 3.0]'], {}), '([8.0 / 3.0])\n', (346, 359), False, 'import torch\n'), ((367, 384), 'torch.randn', 'torch.randn', (['(2)', '(3)'], {}), '(2, 3)\n', (378, 384), False, 'import torch\n'), ((411, 470), 'ceem.systems.LorenzAttractor', 'LorenzAttractor', (['sigma', 'rho', 'beta', 'C', 'dt'], {'method': '"""midpoint"""'}), "(sigma, rho, beta, C, dt, method='midpoint')\n", (426, 470), False, 'from ceem.systems import LorenzAttractor\n'), ((853, 872), 'torch.zeros_like', 'torch.zeros_like', (['x'], {}), '(x)\n', (869, 872), False, 'import torch\n'), ((504, 524), 'torch.randn', 'torch.randn', (['B', '(1)', '(3)'], {}), '(B, 1, 3)\n', (515, 524), False, 'import torch\n'), ((812, 837), 'torch.get_default_dtype', 'torch.get_default_dtype', ([], {}), '()\n', (835, 837), False, 'import torch\n'), ((917, 930), 'torch.ones', 'torch.ones', (['(2)'], {}), '(2)\n', (927, 930), False, 'import torch\n'), ((979, 992), 'torch.ones', 'torch.ones', (['(3)'], {}), '(3)\n', (989, 992), False, 'import torch\n'), ((1410, 1429), 'torch.zeros_like', 'torch.zeros_like', (['x'], {}), '(x)\n', (1426, 1429), False, 'import torch\n'), ((622, 642), 'torch.cat', 'torch.cat', (['xs'], {'dim': '(1)'}), '(xs, dim=1)\n', (631, 642), False, 'import torch\n'), ((580, 603), 'torch.tensor', 'torch.tensor', (['([0.0] * B)'], {}), '([0.0] * B)\n', (592, 603), False, 'import torch\n'), ((774, 789), 'torch.arange', 'torch.arange', (['T'], {}), '(T)\n', (786, 789), False, 'import torch\n'), ((791, 806), 'torch.arange', 'torch.arange', (['T'], {}), '(T)\n', (803, 806), False, 'import torch\n')]