"""Phase 4 — tests de l'encodeur image ViT from scratch.

CPU, rapides : shapes, patch embedding, position apprise, projecteur, et
**overfit** sur un petit jeu d'images jouets (classification + reconstruction).
"""

from __future__ import annotations

import numpy as np
import pytest
import torch

from src.data.toy_images import class_names, generate_dataset, make_image, to_arrays
from src.model.projectors import MLPProjector, build_projector
from src.model.vit import ImageClassifier, ImageDecoder, PatchEmbedding, VisionEncoder, ViTConfig

torch.manual_seed(0)

TINY = ViTConfig(image_size=16, patch_size=8, n_layer=2, n_head=2, n_embd=32)


# ---------------------------------------------------------------------------
# Générateur jouet
# ---------------------------------------------------------------------------


def test_toy_dataset_shapes_and_balance():
    items = generate_dataset(90, grid=16, seed=0)
    images, labels = to_arrays(items)
    assert images.shape == (90, 3, 16, 16)
    assert images.min() >= 0.0 and images.max() <= 1.0
    counts = np.bincount(labels, minlength=9)
    assert counts.min() == counts.max() == 10  # équilibré
    assert len(class_names()) == 9


def test_toy_dataset_is_deterministic():
    a = to_arrays(generate_dataset(10, grid=16, seed=42))[0]
    b = to_arrays(generate_dataset(10, grid=16, seed=42))[0]
    assert np.array_equal(a, b)


def test_caption_is_french():
    item = generate_dataset(1, grid=16, seed=1)[0]
    assert "un cercle" in item.caption or "un carre" in item.caption or "un triangle" in item.caption


def test_make_image_shape():
    img = make_image("cercle", "rouge", grid=24)
    assert img.shape == (3, 24, 24)
    assert img.dtype == np.uint8


# ---------------------------------------------------------------------------
# Patch embedding
# ---------------------------------------------------------------------------


def test_patch_embedding_shape():
    pe = PatchEmbedding(16, 8, 3, 32)
    assert pe.grid == 2 and pe.n_patches == 4
    out = pe(torch.randn(2, 3, 16, 16))
    assert out.shape == (2, 4, 32)


def test_patch_embedding_rejects_bad_size():
    with pytest.raises(ValueError):
        PatchEmbedding(15, 8, 3, 32)


# ---------------------------------------------------------------------------
# Encodeur
# ---------------------------------------------------------------------------


def test_encoder_output_shape_with_cls():
    enc = VisionEncoder(TINY).eval()
    out = enc(torch.randn(2, 3, 16, 16))
    assert out.shape == (2, 1 + 4, 32)
    assert enc.pool(out).shape == (2, 32)


def test_encoder_is_bidirectional_not_causal():
    """Régression : un ViT doit être bidirectionnel (sinon le CLS ne voit pas les patchs)."""
    enc = VisionEncoder(TINY).eval()
    assert all(block.attn.causal is False for block in enc.blocks)
    x = torch.randn(1, 3, 16, 16)
    with torch.no_grad():
        a = enc(x)[:, 0]  # CLS
        x2 = x.clone()
        x2[0, :, :8, :8] += 1.0  # modifie un patch au hasard
        b = enc(x2)[:, 0]
    assert not torch.allclose(a, b), "le CLS ne dépend pas des patchs -> attention causale !"


def test_encoder_mean_pool_has_no_cls():
    cfg = ViTConfig(**{**TINY.__dict__, "pool": "mean"})
    enc = VisionEncoder(cfg).eval()
    out = enc(torch.randn(2, 3, 16, 16))
    assert out.shape == (2, 4, 32)
    assert enc.pool(out).shape == (2, 32)


def test_positional_embedding_is_learned_parameter():
    enc = VisionEncoder(TINY)
    assert enc.pos_embedding.requires_grad
    assert enc.pos_embedding.shape == (1, 5, 32)


def test_encoder_is_shape_robust_to_batch():
    enc = VisionEncoder(TINY).eval()
    for b in (1, 3, 5):
        assert enc(torch.randn(b, 3, 16, 16)).shape[0] == b


# ---------------------------------------------------------------------------
# Projecteur
# ---------------------------------------------------------------------------


def test_projector_maps_to_text_dim():
    proj = build_projector(32, 64)
    out = proj(torch.randn(4, 32))
    assert out.shape == (4, 64)


def test_projector_with_sequence_input():
    proj = MLPProjector(32, 64)
    out = proj(torch.randn(2, 5, 32))
    assert out.shape == (2, 5, 64)


# ---------------------------------------------------------------------------
# Overfit
# ---------------------------------------------------------------------------


def test_overfit_toy_classification():
    torch.manual_seed(0)
    images, labels = to_arrays(generate_dataset(18, grid=16, seed=0))
    x = torch.from_numpy(images)
    y = torch.from_numpy(labels)
    model = ImageClassifier(TINY, n_classes=9)
    opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
    for _ in range(120):
        logits = model(x)
        loss = torch.nn.functional.cross_entropy(logits, y)
        opt.zero_grad()
        loss.backward()
        opt.step()
    acc = (model(x).argmax(-1) == y).float().mean().item()
    assert acc > 0.9, f"n'a pas overfit : acc {acc}"
    assert loss.item() < 0.5


def test_overfit_toy_reconstruction():
    torch.manual_seed(0)
    images, _ = to_arrays(generate_dataset(8, grid=16, seed=0))
    x = torch.from_numpy(images)
    cfg = ViTConfig(**{**TINY.__dict__, "pool": "mean"})
    model = ImageDecoder(cfg, hidden=64)
    opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
    for _ in range(120):
        recon = model(x)
        loss = torch.nn.functional.mse_loss(recon, x)
        opt.zero_grad()
        loss.backward()
        opt.step()
    assert loss.item() < 0.05, f"reconstruction insuffisante : mse {loss.item()}"