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, 'ceem.utils.set_rng_seed', 'utils.set_rng_seed', (['(1)'], {}), False, 'from ceem import utils\n'), (14, 'torch.set_default_dtype', 'torch.set_default_dtype', (['torch.float64'], {}), False, 'import torch\n'), (16, 'torch.tensor', 'torch.tensor', (['[10.0]'], {}), False, 'import torch\n'), (17, 'torch.tensor', 'torch.tensor', (['[28.0]'], {}), False, 'import torch\n'), (18, 'torch.tensor', 'torch.tensor', (['[8.0 / 3.0]'], {}), False, 'import torch\n'), (20, 'torch.randn', 'torch.randn', (['(2)', '(3)'], {}), False, 'import torch\n'), (24, 'ceem.systems.LorenzAttractor', 'LorenzAttractor', (['sigma', 'rho', 'beta', 'C', 'dt'], {'method': '"""midpoint"""'}), False, 'from ceem.systems import LorenzAttractor\n'), (39, 'torch.zeros_like', 'torch.zeros_like', (['x'], {}), False, 'import torch\n'), (28, 'torch.randn', 'torch.randn', (['B', '(1)', '(3)'], {}), False, 'import torch\n'), (37, 'torch.get_default_dtype', 'torch.get_default_dtype', ([], {}), False, 'import torch\n'), (41, 'torch.ones', 'torch.ones', (['(2)'], {}), False, 'import torch\n'), (43, 'torch.ones', 'torch.ones', (['(3)'], {}), False, 'import torch\n'), (58, 'torch.zeros_like', 'torch.zeros_like', (['x'], {}), False, 'import torch\n'), (32, 'torch.cat', 'torch.cat', (['xs'], {'dim': '(1)'}), False, 'import torch\n'), (30, 'torch.tensor', 'torch.tensor', (['([0.0] * B)'], {}), False, 'import torch\n'), (37, 'torch.arange', 'torch.arange', (['T'], {}), False, 'import torch\n'), (37, 'torch.arange', 'torch.arange', (['T'], {}), False, 'import torch\n')]