"""Tests for ScrollCVAE model."""
from __future__ import annotations

import torch
import pytest

from ai_mouse.scroll.models import ScrollCVAE


class TestScrollCVAEForward:
    @pytest.fixture
    def model(self):
        return ScrollCVAE(seq_len=32, latent_dim=16, hidden=64, cond_dim=7)

    def test_output_shapes(self, model):
        batch = 4
        seq = torch.randn(batch, 32, 2)
        cond = torch.randn(batch, 7)
        recon, mu, logvar = model(seq, cond)
        assert recon.shape == (batch, 32, 2)
        assert mu.shape == (batch, 16)
        assert logvar.shape == (batch, 16)

    def test_decode_shape(self, model):
        z = torch.randn(4, 16)
        cond = torch.randn(4, 7)
        out = model.decode(z, cond)
        assert out.shape == (4, 32, 2)

    def test_decode_deterministic(self, model):
        model.eval()
        z = torch.randn(1, 16)
        cond = torch.randn(1, 7)
        with torch.no_grad():
            out1 = model.decode(z, cond)
            out2 = model.decode(z, cond)
        torch.testing.assert_close(out1, out2)