"""Tests for TrajectoryFlowModel architecture."""
from __future__ import annotations

import torch
import pytest

from ai_mouse.models import TrajectoryFlowModel


class TestTrajectoryFlowModel:
    """Test the Conditional Flow Matching model."""

    @pytest.fixture
    def model(self):
        return TrajectoryFlowModel(
            seq_len=64, d_model=128, nhead=4, num_layers=4,
            dim_feedforward=256, dropout=0.1, cond_dim=3,
        )

    def test_output_shape(self, model):
        """(4, 64, 3) input → (4, 64, 3) output."""
        batch = 4
        x_t = torch.randn(batch, 64, 3)
        t = torch.rand(batch)
        cond = torch.randn(batch, 3)
        out = model(x_t, t, cond)
        assert out.shape == (batch, 64, 3)

    def test_single_sample(self, model):
        """(1, 64, 3) works."""
        x_t = torch.randn(1, 64, 3)
        t = torch.rand(1)
        cond = torch.randn(1, 3)
        out = model(x_t, t, cond)
        assert out.shape == (1, 64, 3)

    def test_deterministic(self, model):
        """Eval mode, same input → same output."""
        model.eval()
        x_t = torch.randn(2, 64, 3)
        t = torch.tensor([0.3, 0.7])
        cond = torch.randn(2, 3)
        with torch.no_grad():
            out1 = model(x_t, t, cond)
            out2 = model(x_t, t, cond)
        torch.testing.assert_close(out1, out2)

    def test_different_timesteps(self, model):
        """t=0.1 vs t=0.9 gives different output."""
        model.eval()
        x_t = torch.randn(1, 64, 3)
        cond = torch.randn(1, 3)
        with torch.no_grad():
            out_early = model(x_t, torch.tensor([0.1]), cond)
            out_late = model(x_t, torch.tensor([0.9]), cond)
        assert not torch.allclose(out_early, out_late, atol=1e-5)

    def test_gradient_flows(self, model):
        """Backward works, grad on x_t exists."""
        model.train()
        x_t = torch.randn(2, 64, 3, requires_grad=True)
        t = torch.rand(2)
        cond = torch.randn(2, 3)
        out = model(x_t, t, cond)
        loss = out.sum()
        loss.backward()
        assert x_t.grad is not None
        assert x_t.grad.shape == (2, 64, 3)
        assert x_t.grad.abs().sum() > 0