import numpy as np, torch
from spectral_world_models.v9_causal_generalization import make_trajectory_cf_batch,trajectory_cf_loss,INTERP,EXTRAP,STRUCTURAL
from spectral_world_models.models import build_model

def test_v9_regimes_are_distinct():
    assert INTERP.acceleration < EXTRAP.acceleration
    assert STRUCTURAL.restitution < INTERP.restitution

def test_trajectory_cf_batch_and_backward():
    b=make_trajectory_cf_batch(n=2,horizon=3,seed=1,device='cpu')
    assert b['left'].shape==(2,3,1,32,32)
    m=build_model('swm_structured_cf_v9')
    loss,parts=trajectory_cf_loss(m,b)
    assert torch.isfinite(loss)
    loss.backward()
    assert any(p.grad is not None for p in m.parameters())
    assert 'traj_cf_dir_loss' in parts