# tests/test_fractus_core.py
"""Tests du noyau fractus natif (torch optionnel : SKIP sans torch).

Le module entier est skippé dans le venv mingw (sans torch) via
pytest.importorskip — le substrat fractus_vorax reste utilisable sans torch.
"""
import pytest

torch = pytest.importorskip("torch", reason="noyau torch optionnel (substrat sans torch)")

from fractus_vorax.model.fractus_core import FractusCore, FractusCoreConfig


def _tiny():
    return FractusCoreConfig(vocab_size=128, d_model=64, n_layers=2, n_experts=4, top_k=2, rank=16, max_seq_len=64)


def test_forward_logits_shape():
    torch.manual_seed(0)
    core = FractusCore(_tiny())
    ids = torch.randint(0, 128, (1, 10))
    logits = core(ids)
    assert logits.shape == (1, 10, 128)


def test_generate_greedy_deterministic():
    torch.manual_seed(0)
    core = FractusCore(_tiny())
    ids = torch.tensor([[1, 2, 3]])
    a = core.generate(ids, max_new_tokens=8)
    b = core.generate(ids, max_new_tokens=8)
    assert torch.equal(a, b)
    assert a.shape == (1, 11)


def test_generate_stops_at_eos():
    torch.manual_seed(0)
    core = FractusCore(_tiny())
    ids = torch.tensor([[1, 2, 3]])
    # déterminer le token que greedy émettrait au 1er pas, l'utiliser comme EOS
    with torch.no_grad():
        first = int(core(ids)[0, -1, :].argmax().item())
    out = core.generate(ids, max_new_tokens=8, eos_token_id=first)
    assert out.shape == (1, 4)          # prompt + 1 token, puis arrêt
    assert int(out[0, -1].item()) == first


def test_state_dict_matches_canonical_fractus_keys():
    """Les clés du state_dict reproduisent EXACTEMENT le motif du checkpoint réel (doc §5).

    Top-level (7) + par bloc : 16 clés hors experts + 2×E×4 clés experts.
    """
    cfg = _tiny()
    core = FractusCore(cfg)
    keys = set(core.state_dict().keys())

    expected = {
        "embed.tok_embed.weight",
        "embed.pos_embed.weight",
        "embed.norm.weight",
        "embed.norm.bias",
        "norm.weight",
        "norm.bias",
        "lm_head.weight",
    }
    for i in range(cfg.n_layers):
        expected |= {
            f"blocks.{i}.norm1.weight", f"blocks.{i}.norm1.bias",
            f"blocks.{i}.attn.w_qkv", f"blocks.{i}.attn.b_qkv",
            f"blocks.{i}.attn.w_out", f"blocks.{i}.attn.b_out",
            f"blocks.{i}.attn.level_logits", f"blocks.{i}.attn.level_offsets",
            f"blocks.{i}.norm_kur.weight", f"blocks.{i}.norm_kur.bias",
            f"blocks.{i}.kuramoto.omega", f"blocks.{i}.kuramoto.coupling_u",
            f"blocks.{i}.kuramoto.coupling_lambda",
            f"blocks.{i}.norm_moe.weight", f"blocks.{i}.norm_moe.bias",
            f"blocks.{i}.moe.expert_phases",
        }
        for e in range(cfg.n_experts):
            for w in ("w1", "w2"):
                expected |= {
                    f"blocks.{i}.moe.experts_{w}.{e}.U",
                    f"blocks.{i}.moe.experts_{w}.{e}.V",
                    f"blocks.{i}.moe.experts_{w}.{e}.scale",
                    f"blocks.{i}.moe.experts_{w}.{e}.bias",
                }

    assert keys == expected, (
        f"manquantes: {sorted(expected - keys)[:5]} / inattendues: {sorted(keys - expected)[:5]}"
    )
    # lm_head liée à l'embedding (doc §5 : poids identiques).
    assert core.lm_head.weight is core.embed.tok_embed.weight


def test_checkpoint_roundtrip(tmp_path):
    """Sauvegarde/charge au format fractus (clé 'model') — strict load."""
    torch.manual_seed(0)
    core = FractusCore(_tiny())
    torch.save({"model": core.state_dict()}, tmp_path / "tiny.pt")
    loaded = FractusCore.from_checkpoint(tmp_path / "tiny.pt")
    assert torch.equal(loaded.generate(torch.tensor([[1, 2, 3]]), 4),
                       core.generate(torch.tensor([[1, 2, 3]]), 4))


@pytest.mark.slow
def test_real_checkpoint_strict_load():
    """Marqué slow : charge le VRAI checkpoint fractus (lecture seule).
    Skip si absent. Ne s'exécute que dans .venv-torch."""
    import pathlib
    ckpt = pathlib.Path(r"C:\Users\PHIL\ZCodeProject\fractus\checkpoints\checkpoints\fractus_1b_latest.pt")
    if not ckpt.exists():
        pytest.skip("checkpoint fractus non présent")
    core = FractusCore.from_checkpoint(ckpt)
    ids = torch.zeros(1, 4, dtype=torch.long)
    with torch.no_grad():
        logits = core(ids)
    assert logits.ndim == 3