"""Phase 2 — tests du transformer texte from scratch.

Tout tourne sur CPU en quelques secondes avec les tailles minuscules de
`configs/model.yaml > test`.
"""

from __future__ import annotations

import pytest
import torch

from src.model.attention import MultiHeadSelfAttention, build_causal_mask
from src.model.block import TransformerBlock
from src.model.ffn import FeedForward
from src.model.rope import apply_rope, build_rope_cache, rotate_half
from src.model.text_model import CortexTextLM, TextConfig

torch.manual_seed(0)

TINY = TextConfig(
    vocab_size=64, n_layer=2, n_head=2, n_embd=32, ctx_len=16, dropout=0.0, tie_embeddings=True
)


# ---------------------------------------------------------------------------
# RoPE
# ---------------------------------------------------------------------------


def test_rope_cache_shapes():
    cos, sin = build_rope_cache(seq_len=8, head_dim=16)
    assert cos.shape == (8, 16) and sin.shape == (8, 16)


def test_rope_cache_rejects_odd_dim():
    with pytest.raises(ValueError):
        build_rope_cache(seq_len=4, head_dim=15)


def test_rotate_half_is_involution_up_to_sign_layout():
    x = torch.arange(8, dtype=torch.float32).view(1, 8)
    y = rotate_half(x)
    assert torch.equal(y, torch.tensor([[-4.0, -5, -6, -7, 0, 1, 2, 3]]))


def test_rope_preserves_norm():
    """RoPE est une rotation : la norme de chaque paire est conservée."""
    q = torch.randn(1, 2, 5, 16)
    k = torch.randn(1, 2, 5, 16)
    cos, sin = build_rope_cache(5, 16)
    q2, k2 = apply_rope(q, k, cos, sin)
    assert torch.allclose(q.norm(dim=-1), q2.norm(dim=-1), atol=1e-5)
    assert torch.allclose(k.norm(dim=-1), k2.norm(dim=-1), atol=1e-5)


def test_rope_position_zero_is_identity():
    q = torch.randn(1, 1, 3, 8)
    k = torch.randn(1, 1, 3, 8)
    cos, sin = build_rope_cache(3, 8)
    q2, _ = apply_rope(q, k, cos, sin)
    assert torch.allclose(q[:, :, 0], q2[:, :, 0], atol=1e-6)


def test_rope_makes_attention_depend_on_relative_position():
    """Le score q·k doit dépendre de l'écart de position, pas de l'absolu."""
    d = 16
    cos, sin = build_rope_cache(8, d)
    base = torch.randn(1, 1, 1, d)
    q = base.repeat(1, 1, 8, 1)
    k = base.repeat(1, 1, 8, 1)
    q_rot, k_rot = apply_rope(q, k, cos, sin)
    scores = (q_rot @ k_rot.transpose(-2, -1))[0, 0]
    # score[i, j] ne dépend que de (i - j)
    assert torch.allclose(scores[0, 1], scores[1, 2], atol=1e-5)
    assert torch.allclose(scores[2, 5], scores[3, 6], atol=1e-5)


# ---------------------------------------------------------------------------
# Masque causal + attention
# ---------------------------------------------------------------------------


def test_causal_mask_shape_and_content():
    mask = build_causal_mask(4)
    assert mask.shape == (4, 4)
    assert mask[0].tolist() == [True, False, False, False]
    assert mask[3].tolist() == [True, True, True, True]


def test_attention_output_shape():
    attn = MultiHeadSelfAttention(32, 2).eval()
    x = torch.randn(2, 7, 32)
    assert attn(x).shape == (2, 7, 32)


def test_attention_is_strictly_causal():
    """Le token t ne dépend pas des tokens > t : on modifie le futur, la sortie passée ne bouge pas."""
    attn = MultiHeadSelfAttention(32, 2, dropout=0.0).eval()
    x = torch.randn(1, 6, 32)
    with torch.no_grad():
        y1 = attn(x)
        x2 = x.clone()
        x2[:, 4:] += 100.0  # on ne modifie QUE le futur
        y2 = attn(x2)
    assert torch.allclose(y1[:, :4], y2[:, :4], atol=1e-5), "fuite d'information du futur"
    assert not torch.allclose(y1[:, 4:], y2[:, 4:])


def test_attention_matches_manual_inference_for_second_token():
    """Vérifie la formule d'attention à la main sur un cas simple."""
    attn = MultiHeadSelfAttention(8, 1, use_rope=False, dropout=0.0).eval()
    x = torch.randn(1, 3, 8)
    with torch.no_grad():
        out = attn(x)
        q = attn.q_proj(x)
        k = attn.k_proj(x)
        v = attn.v_proj(x)
        scores = (q @ k.transpose(-2, -1)) / (8 ** 0.5)
        scores = scores.masked_fill(~build_causal_mask(3), float("-inf"))
        w = torch.softmax(scores, dim=-1)
        manual = w @ v
        manual = attn.out_proj(manual)
    assert torch.allclose(out, manual, atol=1e-5)


# ---------------------------------------------------------------------------
# FFN + bloc
# ---------------------------------------------------------------------------


def test_ffn_shape_and_param_count():
    ffn = FeedForward(32, mult=4, bias=False)
    x = torch.randn(2, 5, 32)
    assert ffn(x).shape == (2, 5, 32)
    # 2 matrices : 32*128 + 128*32
    assert sum(p.numel() for p in ffn.parameters()) == 32 * 128 + 128 * 32


def test_block_preserves_shape():
    block = TransformerBlock(32, 2).eval()
    x = torch.randn(2, 5, 32)
    assert block(x).shape == (2, 5, 32)


def test_block_is_prenorm():
    block = TransformerBlock(32, 2)
    assert hasattr(block, "ln_1") and hasattr(block, "ln_2")


# ---------------------------------------------------------------------------
# Modèle complet
# ---------------------------------------------------------------------------


def test_model_forward_shapes():
    model = CortexTextLM(TINY).eval()
    x = torch.randint(0, 64, (2, 16))
    logits, loss = model(x)
    assert logits.shape == (2, 16, 64)
    assert loss is None


def test_model_loss_is_scalar():
    model = CortexTextLM(TINY).eval()
    x = torch.randint(0, 64, (2, 16))
    _, loss = model(x, targets=x)
    assert loss.dim() == 0
    # À l'initialisation, la perte doit être proche de log(vocab) (distribution quasi uniforme).
    assert abs(loss.item() - torch.log(torch.tensor(64.0)).item()) < 1.0


def test_model_rejects_too_long_sequence():
    model = CortexTextLM(TINY).eval()
    x = torch.randint(0, 64, (1, 17))
    with pytest.raises(ValueError):
        model(x)


def test_weight_tying():
    model = CortexTextLM(TINY)
    assert model.lm_head.weight is model.token_embedding.weight


def test_untied_model_has_more_params():
    tied = CortexTextLM(TINY)
    untied = CortexTextLM(
        TextConfig(**{**TINY.__dict__, "tie_embeddings": False})
    )
    assert untied.count_parameters() > tied.count_parameters()


def test_model_is_causal():
    model = CortexTextLM(TINY).eval()
    x = torch.randint(0, 64, (1, 10))
    with torch.no_grad():
        a, _ = model(x)
        x2 = x.clone()
        x2[:, 5:] = torch.randint(0, 64, (1, 5))
        b, _ = model(x2)
    assert torch.allclose(a[:, :5], b[:, :5], atol=1e-5), "fuite du futur dans les logits"


# ---------------------------------------------------------------------------
# Overfit 1 batch : la loss doit descendre à ~0
# ---------------------------------------------------------------------------


def test_overfit_single_batch():
    torch.manual_seed(0)
    model = CortexTextLM(TINY)
    x = torch.randint(0, 64, (4, 16))
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3)

    first_loss = None
    for step in range(200):
        logits, loss = model(x, targets=x)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if step == 0:
            first_loss = loss.item()

    final_loss = loss.item()
    assert first_loss is not None and first_loss > 1.0
    assert final_loss < 0.05, f"n'a pas overfit : loss finale {final_loss}"
    assert final_loss < first_loss / 20