tests/test_model.py

"""
Smoke tests for all model components.

Run: python -m pytest tests/test_model.py -v
"""

import torch
import pytest
import unittest
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


@pytest.fixture
def device():
    # Use CPU for testing — MPS has bus errors with some operations on older PyTorch
    # Training can still use MPS; this just ensures tests are stable
    return torch.device("cpu")


class TestAttention:
    def test_multi_head_attention_bidirectional(self, device):
        from model.attention import MultiHeadAttention
        mha = MultiHeadAttention(768, 12, mode="bidirectional").to(device)
        x = torch.randn(2, 10, 768, device=device)
        out = mha(x)
        assert out.shape == (2, 10, 768)

    def test_multi_head_attention_causal(self, device):
        from model.attention import MultiHeadAttention
        mha = MultiHeadAttention(768, 12, mode="causal").to(device)
        x = torch.randn(2, 10, 768, device=device)
        out = mha(x)
        assert out.shape == (2, 10, 768)


class TestTransformerBlock:
    def test_forward(self, device):
        from model.transformer import TransformerBlock
        block = TransformerBlock(768, 12, mode="bidirectional").to(device)
        x = torch.randn(2, 10, 768, device=device)
        out = block(x)
        assert out.shape == (2, 10, 768)


class TestPatchEmbed:
    def test_patch_embedding(self, device):
        from model.patch_embed import PatchEmbedding
        pe = PatchEmbedding(384, 16, 3, 768).to(device)
        img = torch.randn(2, 3, 384, 384, device=device)
        out = pe(img)
        assert out.shape == (2, 576, 768), f"Expected (2, 576, 768), got {out.shape}"


class TestViT:
    def test_vjepa_encoder(self, device):
        from model.vit import VJEPAEncoder
        vit = VJEPAEncoder(384, 16, 768, 12, num_blocks=2).to(device)  # 2 blocks for speed
        img = torch.randn(2, 3, 384, 384, device=device)
        out = vit(img)
        assert out.shape == (2, 576, 768)


class TestTokenizer(unittest.TestCase):
    def test_special_tokens_extended(self):
        """32K tokenizer has 15 special tokens including new task types."""
        from model.tokenizer import BPETokenizer
        tok = BPETokenizer(vocab_size=32768)
        assert tok.pad_id == 0
        assert tok.bos_id == 1
        assert tok.eos_id == 2
        for name in ["<count>", "<ocr>", "<reason>", "<temporal>", "<multi_cam>", "<system>"]:
            assert name in tok.SPECIAL_TOKENS, f"Missing special token {name}"
        assert len(tok.SPECIAL_TOKENS) == 15

    def test_byte_fallback(self):
        """Unknown characters fall back to byte-level encoding."""
        from model.tokenizer import BPETokenizer
        tok = BPETokenizer(vocab_size=32768)
        tok.train(["hello world", "the quick brown fox"])
        ids = tok.encode("hello 你好")
        text = tok.decode(ids)
        assert "hello" in text

    def test_encode_decode_roundtrip_32k(self):
        """Encode/decode preserves text with 32K vocab."""
        from model.tokenizer import BPETokenizer
        tok = BPETokenizer(vocab_size=32768)
        tok.train(["the cat sat on the mat", "a dog ran in the park"] * 100)
        text = "the cat ran"
        ids = tok.encode(text)
        decoded = tok.decode(ids)
        assert "the cat ran" in decoded

    def test_save_load_32k(self):
        """Save and load preserves 32K tokenizer state."""
        from model.tokenizer import BPETokenizer
        tok = BPETokenizer(vocab_size=32768)
        tok.train(["hello world test data"] * 50)
        tok.save("/tmp/test_tok_32k.json")
        tok2 = BPETokenizer(vocab_size=32768)
        tok2.load("/tmp/test_tok_32k.json")
        assert tok.encode("hello world") == tok2.encode("hello world")


class TestYEncoder:
    def test_forward(self, device):
        from model.y_encoder import YEncoder
        enc = YEncoder(vocab_size=100, hidden_dim=128, embed_dim=256, num_heads=4, num_blocks=2).to(device)
        ids = torch.randint(0, 100, (2, 20), device=device)
        mask = torch.ones(2, 20, dtype=torch.bool, device=device)
        out = enc(ids, mask)
        assert out.shape == (2, 256)
        # Verify L2 normalized
        norms = torch.norm(out, dim=-1)
        assert torch.allclose(norms, torch.ones_like(norms), atol=1e-5)


class TestPredictor:
    def test_forward_with_query(self, device):
        from model.predictor import JEPAPredictor
        pred = JEPAPredictor(hidden_dim=128, embed_dim=256, num_heads=4, num_blocks=2, vocab_size=100).to(device)
        visual = torch.randn(2, 36, 128, device=device)  # small patch count for test
        query = torch.randint(0, 100, (2, 10), device=device)
        mask = torch.ones(2, 10, dtype=torch.bool, device=device)
        out = pred(visual, query, mask)
        assert out.shape == (2, 256)

    def test_forward_without_query(self, device):
        from model.predictor import JEPAPredictor
        pred = JEPAPredictor(hidden_dim=128, embed_dim=256, num_heads=4, num_blocks=2, vocab_size=100).to(device)
        visual = torch.randn(2, 36, 128, device=device)
        out = pred(visual)
        assert out.shape == (2, 256)


class TestInfoNCELoss:
    def test_matching_pairs_low_loss(self, device):
        from model.losses import InfoNCELoss
        loss_fn = InfoNCELoss(temperature=0.07)
        # Matching embeddings should have low loss
        embeds = torch.randn(8, 256, device=device)
        embeds = embeds / embeds.norm(dim=-1, keepdim=True)
        loss = loss_fn(embeds, embeds)
        assert loss.item() < 1.0  # self-matching should be very low loss

    def test_random_pairs_higher_loss(self, device):
        from model.losses import InfoNCELoss
        loss_fn = InfoNCELoss(temperature=0.07)
        pred = torch.randn(8, 256, device=device)
        target = torch.randn(8, 256, device=device)
        pred = pred / pred.norm(dim=-1, keepdim=True)
        target = target / target.norm(dim=-1, keepdim=True)
        loss = loss_fn(pred, target)
        # Random pairs should have higher loss than self-matching
        self_loss = loss_fn(pred, pred)
        assert loss.item() > self_loss.item()


class TestMoE:
    def test_moe_layer(self, device):
        from model.moe import MoELayer
        moe = MoELayer(embed_dim=128, num_experts=5, top_k=2).to(device)
        x = torch.randn(2, 10, 128, device=device)
        out = moe(x)
        assert out.shape == (2, 10, 128)

    def test_expert_routing(self, device):
        from model.moe import MoELayer
        moe = MoELayer(embed_dim=128, num_experts=5, top_k=2).to(device)
        x = torch.randn(2, 10, 128, device=device)
        _ = moe(x)
        data = moe.get_load_balancing_data()
        assert data is not None
        gate_probs, expert_indices = data
        assert gate_probs.shape[1] == 5  # 5 experts
        assert expert_indices.shape[1] == 2  # top-2


class TestMoEDecoder:
    def test_forward_training(self, device):
        from model.decoder import MoEDecoder
        dec = MoEDecoder(hidden_dim=128, embed_dim=256, vocab_size=100, num_heads=4,
                         num_blocks=2, num_experts=5, top_k=2).to(device)
        embed = torch.randn(2, 256, device=device)
        targets = torch.randint(1, 100, (2, 20), device=device)
        logits, loss = dec(embed, targets)
        assert logits.shape[0] == 2
        assert loss is not None and loss.item() > 0

    def test_generate(self, device):
        from model.decoder import MoEDecoder
        dec = MoEDecoder(hidden_dim=128, embed_dim=256, vocab_size=100, num_heads=4,
                         num_blocks=2, num_experts=5, top_k=2).to(device)
        embed = torch.randn(1, 256, device=device)
        generated = dec.generate(embed, max_new_tokens=10)
        assert generated.shape[0] == 1
        assert generated.shape[1] <= 10


class TestSelectiveDecoder:
    def test_first_frame_always_decodes(self):
        from model.selective_decode import SelectiveDecoder
        sd = SelectiveDecoder(similarity_threshold=0.95)
        embed = torch.randn(1536)
        assert sd.should_decode("cam1", embed, 0.0) is True

    def test_similar_frame_skipped(self):
        from model.selective_decode import SelectiveDecoder
        sd = SelectiveDecoder(similarity_threshold=0.95)
        embed = torch.randn(1536)
        embed = embed / embed.norm()
        sd.should_decode("cam1", embed, 0.0)
        # Same embedding = similarity 1.0 > 0.95 → skip
        assert sd.should_decode("cam1", embed, 2.0) is False

    def test_different_frame_decodes(self):
        from model.selective_decode import SelectiveDecoder
        sd = SelectiveDecoder(similarity_threshold=0.95)
        embed1 = torch.randn(1536)
        embed1 = embed1 / embed1.norm()
        sd.should_decode("cam1", embed1, 0.0)
        embed2 = torch.randn(1536)
        embed2 = embed2 / embed2.norm()
        # Different random embedding → low similarity → should decode
        assert sd.should_decode("cam1", embed2, 2.0) is True


class TestFullModel:
    def test_stage1_forward(self, device):
        """Test Stage 1 JEPA pretraining forward pass."""
        import yaml
        with open("configs/default.yaml") as f:
            config = yaml.safe_load(f)

        # Use small model for testing
        config["vision"]["num_blocks"] = 1
        config["predictor"]["num_blocks"] = 1
        config["y_encoder"]["num_blocks"] = 1
        config["decoder"]["num_blocks"] = 1

        from model.vlm import VLJEPAModel
        model = VLJEPAModel(config).to(device)

        B = 2
        images = torch.randn(B, 3, 384, 384, device=device)
        answer_ids = torch.randint(1, 100, (B, 20), device=device)
        answer_mask = torch.ones(B, 20, dtype=torch.bool, device=device)

        output = model.forward_stage1(images, None, None, answer_ids, answer_mask)
        assert "loss" in output
        assert output["loss"].requires_grad
        assert output["pred_embeds"].shape == (B, 1536)
        assert output["target_embeds"].shape == (B, 1536)

    def test_stage2_forward(self, device):
        """Test Stage 2 finetuning forward pass."""
        import yaml
        with open("configs/default.yaml") as f:
            config = yaml.safe_load(f)

        config["vision"]["num_blocks"] = 1
        config["predictor"]["num_blocks"] = 1
        config["y_encoder"]["num_blocks"] = 1
        config["decoder"]["num_blocks"] = 1

        from model.vlm import VLJEPAModel
        model = VLJEPAModel(config).to(device)

        B = 2
        images = torch.randn(B, 3, 384, 384, device=device)
        q_ids = torch.randint(1, 100, (B, 10), device=device)
        q_mask = torch.ones(B, 10, dtype=torch.bool, device=device)
        a_ids = torch.randint(1, 100, (B, 20), device=device)

        output = model.forward_stage2(images, q_ids, q_mask, a_ids)
        assert "loss" in output
        assert "decode_loss" in output
        assert "load_balance_loss" in output

    def test_parameter_count(self):
        import yaml
        with open("configs/default.yaml") as f:
            config = yaml.safe_load(f)

        from model.vlm import VLJEPAModel
        model = VLJEPAModel(config)
        params = model.count_parameters()
        print(f"\nFull model parameters: {params['total']:,}")
        assert params["total"] > 0
        assert params["x_encoder"] > 0
        assert params["predictor"] > 0
        assert params["decoder"] > 0


class TestScaledComponents(unittest.TestCase):
    def test_predictor_2048_embed(self):
        """Predictor with 1024d hidden, 2048 embed dim."""
        from model.predictor import JEPAPredictor
        pred = JEPAPredictor(
            hidden_dim=1024, embed_dim=2048, num_heads=16,
            num_blocks=12, vocab_size=32768, max_query_len=512
        )
        visual = torch.randn(2, 1024, 1024)  # 1024 patches, 1024d
        query_ids = torch.randint(0, 32768, (2, 32))
        out = pred(visual, query_ids)
        assert out.shape == (2, 2048), f"Expected (2, 2048), got {out.shape}"

    def test_y_encoder_2048_embed(self):
        """Y-Encoder with 1024d hidden, 2048 embed dim."""
        from model.y_encoder import YEncoder
        enc = YEncoder(
            vocab_size=32768, hidden_dim=1024, embed_dim=2048,
            num_heads=16, num_blocks=8, max_seq_len=512
        )
        ids = torch.randint(0, 32768, (2, 64))
        out = enc(ids)
        assert out.shape == (2, 2048), f"Expected (2, 2048), got {out.shape}"

    def test_decoder_8_experts_4096_seq(self):
        """MoE Decoder with 8 experts, 4096 max seq, 32K vocab."""
        from model.decoder import MoEDecoder
        dec = MoEDecoder(
            hidden_dim=1024, embed_dim=2048, vocab_size=32768,
            num_heads=16, num_blocks=4, num_experts=8, top_k=2,
            max_seq_len=4096
        )
        embedding = torch.randn(2, 2048)
        target_ids = torch.randint(0, 32768, (2, 64))
        logits, loss = dec(embedding, target_ids)
        assert logits.shape[0] == 2
        assert logits.shape[2] == 32768, f"Expected vocab 32768, got {logits.shape[2]}"
        assert loss is not None

    def test_moe_8_experts(self):
        """MoE layer with 8 experts, top-2 routing."""
        from model.moe import MoELayer
        moe = MoELayer(embed_dim=1024, num_experts=8, top_k=2, expansion=4)
        x = torch.randn(2, 32, 1024)
        out = moe(x)
        assert out.shape == (2, 32, 1024), f"Expected (2, 32, 1024), got {out.shape}"

    def test_load_balance_8_experts(self):
        """Load balancing loss works with 8 experts."""
        from model.losses import LoadBalancingLoss
        lb = LoadBalancingLoss(num_experts=8)
        gate_probs = torch.softmax(torch.randn(64, 8), dim=-1)
        indices = torch.randint(0, 8, (64, 2))
        loss = lb(gate_probs, indices)
        assert loss.shape == ()
        assert loss.item() > 0


class TestViTScaled(unittest.TestCase):
    def test_vit_1024d_patch14(self):
        """1.3B ViT: 448x448, patch 14, 24 layers, 1024d, 16 heads."""
        from model.vit import VJEPAEncoder
        encoder = VJEPAEncoder(
            img_size=448, patch_size=14, hidden_dim=1024,
            num_heads=16, num_blocks=24, dropout=0.1
        )
        x = torch.randn(2, 3, 448, 448)
        out = encoder(x)
        assert out.shape == (2, 1024, 1024), f"Expected (2, 1024, 1024), got {out.shape}"

    def test_vit_param_count_scaled(self):
        """Scaled ViT should have ~300M+ params."""
        from model.vit import VJEPAEncoder
        encoder = VJEPAEncoder(
            img_size=448, patch_size=14, hidden_dim=1024,
            num_heads=16, num_blocks=24
        )
        params = sum(p.numel() for p in encoder.parameters())
        assert params > 200_000_000, f"Expected >200M params, got {params:,}"


class TestTemporalAttention(unittest.TestCase):
    def test_single_frame(self):
        """Single frame passthrough."""
        from model.temporal_attention import TemporalAttention
        ta = TemporalAttention(embed_dim=2048, hidden_dim=1024, num_heads=16, num_blocks=4, max_frames=32)
        embeddings = torch.randn(2, 1, 2048)
        out = ta(embeddings)
        assert out.shape == (2, 1, 1024), f"Expected (2, 1, 1024), got {out.shape}"

    def test_multi_frame(self):
        """4 frames from same camera."""
        from model.temporal_attention import TemporalAttention
        ta = TemporalAttention(embed_dim=2048, hidden_dim=1024, num_heads=16, num_blocks=4, max_frames=32)
        embeddings = torch.randn(2, 4, 2048)
        out = ta(embeddings)
        assert out.shape == (2, 4, 1024), f"Expected (2, 4, 1024), got {out.shape}"

    def test_max_frames(self):
        """32 frames — maximum temporal window."""
        from model.temporal_attention import TemporalAttention
        ta = TemporalAttention(embed_dim=2048, hidden_dim=1024, num_heads=16, num_blocks=4, max_frames=32)
        embeddings = torch.randn(1, 32, 2048)
        out = ta(embeddings)
        assert out.shape == (1, 32, 1024), f"Expected (1, 32, 1024), got {out.shape}"

    def test_temporal_position_embeddings(self):
        """Different positions produce different outputs."""
        from model.temporal_attention import TemporalAttention
        ta = TemporalAttention(embed_dim=2048, hidden_dim=1024, num_heads=16, num_blocks=4, max_frames=32)
        single = torch.randn(1, 1, 2048)
        repeated = single.expand(1, 4, 2048).clone()
        out = ta(repeated)
        assert not torch.allclose(out[0, 0], out[0, 3], atol=1e-4), "Positions should produce different outputs"


class TestFullModel1_3B(unittest.TestCase):
    """Integration tests for VLJEPAModel with temporal attention at 1.3B-shaped config."""

    def _make_config(self):
        """Minimal 1.3B-shaped config (tiny dims for fast testing)."""
        return {
            "vision": {"img_size": 56, "patch_size": 14, "hidden_dim": 64, "num_heads": 4, "num_blocks": 2, "dropout": 0.0},
            "predictor": {"num_blocks": 2, "hidden_dim": 64, "num_heads": 4, "embed_dim": 128, "max_query_len": 32, "dropout": 0.0},
            "y_encoder": {"num_blocks": 2, "hidden_dim": 64, "num_heads": 4, "embed_dim": 128, "max_seq_len": 32, "lr_multiplier": 0.05, "dropout": 0.0},
            "decoder": {"num_blocks": 2, "hidden_dim": 64, "num_heads": 4, "num_experts": 8, "top_k": 2, "vocab_size": 256, "max_seq_len": 128, "dropout": 0.0},
            "temporal": {"num_blocks": 2, "hidden_dim": 64, "num_heads": 4, "max_frames": 8, "dropout": 0.0},
            "selective_decode": {"similarity_threshold": 0.95, "min_decode_interval": 1.0},
            "train_stage1": {"infonce_temperature": 0.07},
        }

    def test_multi_frame_embedding(self):
        """Multi-frame forward pass through temporal attention."""
        from model.vlm import VLJEPAModel
        config = self._make_config()
        model = VLJEPAModel(config)
        images = torch.randn(2, 4, 3, 56, 56)  # [B, 4 frames, C, H, W]
        query_ids = torch.randint(0, 256, (2, 8))
        embeddings = model.get_embedding_multi_frame(images, query_ids)
        assert embeddings.shape == (2, 4, 64), f"Expected (2, 4, 64), got {embeddings.shape}"

    def test_generate_multi_frame(self):
        """Generate text from multi-frame input."""
        from model.vlm import VLJEPAModel
        config = self._make_config()
        model = VLJEPAModel(config)
        images = torch.randn(1, 4, 3, 56, 56)
        query_ids = torch.randint(0, 256, (1, 8))
        output_ids = model.generate_multi_frame(images, query_ids, max_new_tokens=16)
        assert output_ids.shape[0] == 1
        assert output_ids.shape[1] > 0

    def test_backward_compat_no_temporal(self):
        """Model without temporal config still works (backward compat)."""
        from model.vlm import VLJEPAModel
        config = self._make_config()
        del config["temporal"]
        model = VLJEPAModel(config)
        assert model.temporal_attention is None
        # Single-frame still works
        images = torch.randn(2, 3, 56, 56)
        query_ids = torch.randint(0, 256, (2, 8))
        emb = model.get_embedding(images, query_ids)
        assert emb.shape == (2, 128)

    def test_count_params_includes_temporal(self):
        """Parameter count includes temporal attention."""
        from model.vlm import VLJEPAModel
        config = self._make_config()
        model = VLJEPAModel(config)
        counts = model.count_parameters()
        assert "temporal_attention" in counts
        assert counts["temporal_attention"] > 0