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()

[((12, 4, 12, 25), 'ceem.utils.set_rng_seed', 'utils.set_rng_seed', ({(12, 23, 12, 24): '(1)'}, {}), '(1)', False, 'from ceem import utils\n'), ((14, 4, 14, 42), 'torch.set_default_dtype', 'torch.set_default_dtype', ({(14, 28, 14, 41): 'torch.float64'}, {}), '(torch.float64)', False, 'import torch\n'), ((16, 12, 16, 31), 'torch.tensor', 'torch.tensor', ({(16, 25, 16, 30): '[10.0]'}, {}), '([10.0])', False, 'import torch\n'), ((17, 10, 17, 29), 'torch.tensor', 'torch.tensor', ({(17, 23, 17, 28): '[28.0]'}, {}), '([28.0])', False, 'import torch\n'), ((18, 11, 18, 34), 'torch.tensor', 'torch.tensor', ({(18, 24, 18, 33): '[8.0 / 3.0]'}, {}), '([8.0 / 3.0])', False, 'import torch\n'), ((20, 8, 20, 25), 'torch.randn', 'torch.randn', ({(20, 20, 20, 21): '2', (20, 23, 20, 24): '3'}, {}), '(2, 3)', False, 'import torch\n'), ((24, 10, 24, 69), 'ceem.systems.LorenzAttractor', 'LorenzAttractor', (), '', False, 'from ceem.systems import LorenzAttractor\n'), ((39, 9, 39, 28), 'torch.zeros_like', 'torch.zeros_like', ({(39, 26, 39, 27): 'x'}, {}), '(x)', False, 'import torch\n'), ((28, 10, 28, 30), 'torch.randn', 'torch.randn', ({(28, 22, 28, 23): 'B', (28, 25, 28, 26): '(1)', (28, 28, 28, 29): '(3)'}, {}), '(B, 1, 3)', False, 'import torch\n'), ((37, 59, 37, 84), 'torch.get_default_dtype', 'torch.get_default_dtype', ({}, {}), '()', False, 'import torch\n'), ((41, 43, 41, 56), 'torch.ones', 'torch.ones', ({(41, 54, 41, 55): '2'}, {}), '(2)', False, 'import torch\n'), ((43, 40, 43, 53), 'torch.ones', 'torch.ones', ({(43, 51, 43, 52): '3'}, {}), '(3)', False, 'import torch\n'), ((58, 31, 58, 50), 'torch.zeros_like', 'torch.zeros_like', ({(58, 48, 58, 49): 'x'}, {}), '(x)', False, 'import torch\n'), ((32, 8, 32, 28), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((30, 27, 30, 49), 'torch.tensor', 'torch.tensor', ({(30, 40, 30, 48): '([0.0] * B)'}, {}), '([0.0] * B)', False, 'import torch\n'), ((37, 21, 37, 36), 'torch.arange', 'torch.arange', ({(37, 34, 37, 35): 'T'}, {}), '(T)', False, 'import torch\n'), ((37, 38, 37, 53), 'torch.arange', 'torch.arange', ({(37, 51, 37, 52): 'T'}, {}), '(T)', False, 'import torch\n')]