"""Tests for LTI injection in ACT loops."""
import torch
from arbitor.components import LTIInjection, ByteHead
from arbitor.decoders import VideoHead


def test_lti_basic_properties():
    lti = LTIInjection(64)
    h = torch.randn(2, 10, 64)
    e = torch.randn(2, 10, 64)
    t = torch.randn(2, 10, 64)
    out = lti(h, e, t)
    assert out.shape == h.shape
    assert torch.isfinite(out).all()


def test_lti_spectral_radius():
    lti = LTIInjection(64)
    A = lti.get_A()
    assert (A > 0).all()
    assert (A < 1).all()


def test_lti_learnable_params():
    lti = LTIInjection(128)
    assert lti.log_A.shape == (128,)
    assert lti.log_dt.shape == (1,)
    assert lti.B.shape == (128,)
    assert sum(p.numel() for p in lti.parameters()) == 128 + 1 + 128


def test_lti_state_decay():
    lti = LTIInjection(8)
    h = torch.ones(1, 1, 8) * 100.0
    e = torch.zeros(1, 1, 8)
    t = torch.zeros(1, 1, 8)
    out = lti(h, e, t)
    assert (out.abs() < 50).all()


def test_lti_initial_state_small():
    lti = LTIInjection(8)
    h = torch.zeros(1, 1, 8)
    e = torch.ones(1, 1, 8) * 5.0
    t = torch.zeros(1, 1, 8)
    out = lti(h, e, t)
    assert (out > 0).all()
    assert (out < 5).all()


def test_bytehead_lti_integration():
    bh = ByteHead()
    x = torch.randn(2, 10, 8192)
    logits = bh(x)
    assert logits.shape[-1] == 288
    assert bh.lti is not None
    assert isinstance(bh.lti, LTIInjection)


def test_bytehead_no_act():
    bh_single = ByteHead(act_max_iters=1)
    assert bh_single.lti is None
    x = torch.randn(1, 5, 8192)
    logits = bh_single(x)
    assert logits.shape[-1] == 288