tests/test_memory.py

"""Phase 11 — tests mémoire : RAG (mode A) et apprentissage permanent (mode B)."""

from __future__ import annotations

import numpy as np
import pytest
import torch

from src.memory.permanent import PermanentLearner, ReplayBuffer, _encode_chunks
from src.memory.rag import (
    LessonMemory,
    VectorIndex,
    chunk_text,
    compute_idf,
    hashed_embed,
)
from src.model.text_model import CortexTextLM, TextConfig
from src.tokenizer import load_tokenizer


@pytest.fixture(scope="module")
def tokenizer():
    tok = load_tokenizer()
    if tok.vocab_size != 32000:
        pytest.skip("tokenizer final 32k requis")
    return tok


# ---------------------------------------------------------------------------
# Embeddings / index (mode A — briques)
# ---------------------------------------------------------------------------


def test_hashed_embed_is_normalized_and_deterministic():
    v = hashed_embed("le chat dort sur le tapis", dim=128)
    assert v.shape == (128,)
    assert abs(float(np.linalg.norm(v)) - 1.0) < 1e-5
    assert np.allclose(v, hashed_embed("le chat dort sur le tapis", dim=128))


def test_similar_texts_are_closer():
    a = hashed_embed("la pomme est un fruit rouge")
    b = hashed_embed("la pomme est un fruit rouge sucré")
    c = hashed_embed("le moteur de la voiture est bruyant")
    assert float(a @ b) > float(a @ c)


def test_chunking_respects_max_words():
    text = " ".join(["Phrase numéro un très courte."] * 20)
    passages = chunk_text(text, max_words=30)
    assert len(passages) >= 2
    assert all(len(p.text.split()) <= 40 for p in passages)


def test_vector_index_search_orders_by_similarity():
    idx = VectorIndex(dim=64)
    idx.add(np.stack([
        hashed_embed("le chien court dans le jardin", 64),
        hashed_embed("la voiture roule sur la route", 64),
        hashed_embed("le soleil brille dans le ciel", 64),
    ]))
    hits = idx.search(hashed_embed("un chien dans le jardin", 64), k=3)
    assert hits[0][0] == 0  # le passage sur le chien est le plus proche
    assert hits[0][1] >= hits[1][1] >= hits[2][1]


def test_idf_gives_rare_words_more_weight():
    texts = ["chat chien", "chat chien", "chat chien", "ornithorynque"]
    idf = compute_idf(texts)
    assert idf["ornithorynque"] > idf["chat"]


# ---------------------------------------------------------------------------
# Mode A — RAG complet sur un « cours »
# ---------------------------------------------------------------------------


def _course():
    return (
        "Le cycle de l'eau décrit le mouvement de l'eau sur Terre. "
        "L'eau s'évapore des océans sous l'effet du soleil. "
        "La vapeur monte, se refroidit et forme des nuages. "
        "La pluie retombe ensuite sur les continents. "
        "Enfin l'eau ruisselle vers les rivières puis la mer. "
        "La photosynthèse permet aux plantes de fabriquer du sucre. "
        "Les plantes absorbent le dioxyde de carbone et rejettent de l'oxygène. "
        "La chlorophylle donne leur couleur verte aux feuilles."
    )


def test_rag_retrieves_relevant_passage():
    mem = LessonMemory(dim=512, max_words=20)
    mem.add_text(_course(), source="cours_svt")
    ans = mem.answer("Comment se forment les nuages ?", k=2)
    assert ans["sources"] == ["cours_svt"]
    assert ans["confidence"] > 0.1
    # Le meilleur passage doit parler de vapeur / nuages.
    assert "nuage" in ans["answer"].lower() or "vapeur" in ans["answer"].lower()


def test_rag_answer_is_extractive_verbatim():
    mem = LessonMemory(max_words=20)
    mem.add_text(_course())
    ans = mem.answer("Que rejettent les plantes ?", k=1)
    assert ans["answer"] in [p.text for p in mem.passages]  # verbatim, pas d'invention


def test_rag_video_ingestion_keeps_sources():
    mem = LessonMemory(max_words=15)
    mem.add_video(
        transcript=[(0.0, "Bienvenue dans ce cours sur les volcans."),
                    (12.0, "Un volcan crache de la lave très chaude."),
                    (30.0, "La lave refroidit et devient de la roche.")],
        frame_captions=[(12.0, "une coulée de lave orange descend la pente")],
        source="volcans.mp4",
    )
    ans = mem.answer("De quoi est faite la lave refroidie ?", k=3)
    assert "volcans.mp4" in ans["sources"]
    assert any("roche" in r["text"].lower() for r in ans["retrieved"])


def test_rag_empty_memory_returns_empty():
    mem = LessonMemory()
    ans = mem.answer("une question sans cours")
    assert ans["answer"] == "" and ans["confidence"] == 0.0


# ---------------------------------------------------------------------------
# Mode B — apprentissage permanent (mécanique + honnêteté)
# ---------------------------------------------------------------------------


def _quick_model(tokenizer) -> CortexTextLM:
    torch.manual_seed(0)
    return CortexTextLM(TextConfig(vocab_size=tokenizer.vocab_size, n_layer=2, n_head=2,
                                   n_embd=64, ctx_len=48, tie_embeddings=True))


def test_replay_buffer_capacity_and_sampling():
    buf = ReplayBuffer(capacity=3)
    for i in range(5):
        buf.add([i, i + 1])
    assert len(buf) == 3                       # capacity respected
    assert buf.items[0] == [2, 3]              # oldest evicted
    rng = np.random.default_rng(0)
    assert len(buf.sample(2, rng)) == 2


def test_encode_chunks_wraps_bos_eos(tokenizer):
    chunks = _encode_chunks(tokenizer, "Bonjour le monde.", ctx_len=64, max_chunks=4)
    assert chunks and chunks[0][0] == __import__("src.tokenizer.special_tokens",
                                                 fromlist=["BOS_ID"]).BOS_ID
    assert chunks[0][-1] == __import__("src.tokenizer.special_tokens",
                                       fromlist=["EOS_ID"]).EOS_ID


def test_permanent_learning_memorizes_new_lesson(tokenizer):
    """Mode B : après ingestion, la perte sur la leçon a chuté (mémorisation)."""
    model = _quick_model(tokenizer)
    learner = PermanentLearner(model=model, tokenizer=tokenizer, lr=5e-3,
                               steps_per_lesson=40, ctx_len=48, device="cpu")
    lesson = "La photosynthèse transforme la lumière en sucre dans les feuilles vertes."
    metrics = learner.ingest(lesson, seed=0)
    assert metrics["steps"] == 40
    assert metrics["loss_after"] < metrics["loss_before"]  # elle mémorise
    assert metrics["replay_size"] >= 1


def test_permanent_learning_can_forget_without_replay(tokenizer):
    """Sans rejeu, ingérer une leçon fait *monter* la perte sur une autre.

    C'est la démonstration honnête de l'oubli catastrophique : le mode B est instable.
    """
    model = _quick_model(tokenizer)
    learner = PermanentLearner(model=model, tokenizer=tokenizer, lr=5e-3,
                               steps_per_lesson=40, replay_ratio=0.0, device="cpu")
    kept = "Le chat noir dort sur le tapis rouge du salon."
    new = "La lave brûlante coule le long de la pente du volcan actif."
    x = torch.tensor([learner.tokenizer.encode(kept)], dtype=torch.long)
    y = torch.tensor([learner.tokenizer.encode(kept)], dtype=torch.long)
    loss_before = float(model(x, y)[1].detach())
    learner.ingest(new, seed=0)
    loss_after = float(model(x, y)[1].detach())
    assert loss_after > loss_before  # oubli : la perte sur l'ancienne phrase augmente