"""Phase 11 — mode B : « apprentissage permanent » — EXPÉRIMENTAL ET INSTABLE.

⚠️ AVERTISSEMENT HONNÊTE. Ce module met à jour les **poids** du modèle pendant
l'« ingestion » d'un cours (quelques pas de gradient). Ce n'est :

* **pas** un vrai fine-tuning (peu de pas, peu de données) ;
* **pas** stable : l'oubli catastrophique est réel, et une mise à jour peut
  **dégrader** d'autres capacités du modèle déjà mesurées (accord image/audio, etc.) ;
* à utiliser qu'avec un **tampon de rejeu** (on ré-entraîne sur d'anciens exemples
  mémorisés) pour limiter l'oubli.

Ce que cela sait faire, honnêtement : mémoriser un court texte (une leçon) et le
restituer. Ce que cela ne fait pas : apprendre un domaine, raisonner, ou égaler un
entraînement complet. Mesures dans docs/architecture.md §10.

Aucun optimiseur externe : boucle SGD/AdamW PyTorch écrite à la main.
"""

from __future__ import annotations

from collections import deque
from dataclasses import dataclass, field

import torch

from src.tokenizer import CortexTokenizer
from src.tokenizer.special_tokens import BOS_ID, EOS_ID


@dataclass
class ReplayBuffer:
    """Tampon de rejeu : garde les derniers exemples vus pour limiter l'oubli."""

    capacity: int = 32
    items: deque = field(default_factory=lambda: deque(maxlen=32))

    def __post_init__(self) -> None:
        if self.items.maxlen != self.capacity:
            self.items = deque(self.items, maxlen=self.capacity)

    def add(self, token_ids: list[int]) -> None:
        self.items.append(list(token_ids))

    def sample(self, n: int, rng) -> list[list[int]]:
        if not self.items:
            return []
        n = min(n, len(self.items))
        idx = rng.choice(len(self.items), size=n, replace=False)
        return [self.items[int(i)] for i in idx]

    def __len__(self) -> int:
        return len(self.items)


def _encode_chunks(tokenizer: CortexTokenizer, text: str, ctx_len: int,
                   max_chunks: int) -> list[list[int]]:
    """Découpe un texte en fenêtres ``<bos> ... <eos>`` de ``ctx_len`` tokens max."""
    ids = [BOS_ID, *tokenizer.encode(text), EOS_ID]
    chunks = [ids[i:i + ctx_len] for i in range(0, len(ids), ctx_len)] if ids else []
    return [c for c in chunks if len(c) >= 2][:max_chunks]


@dataclass
class PermanentLearner:
    """Ingestion d'un cours par mise à jour des poids (mode B)."""

    model: torch.nn.Module
    tokenizer: CortexTokenizer
    lr: float = 3e-4
    steps_per_lesson: int = 20
    ctx_len: int = 128
    replay_ratio: float = 0.5  # part de rejeu vs nouveau texte
    replay: ReplayBuffer = field(default_factory=ReplayBuffer)
    device: str = "cpu"

    def __post_init__(self) -> None:
        self.model.to(self.device)
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.lr)

    def _loss_on(self, token_ids: list[int]) -> torch.Tensor:
        x = torch.tensor([token_ids[:-1]], dtype=torch.long, device=self.device)
        y = torch.tensor([token_ids[1:]], dtype=torch.long, device=self.device)
        _logits, loss = self.model(x, y)
        return loss

    def ingest(self, text: str, seed: int = 0) -> dict:
        """Met à jour les poids sur ``text`` (+ rejeu). Retourne les métriques."""
        rng = __import__("numpy").random.default_rng(seed)
        chunks = _encode_chunks(self.tokenizer, text, self.ctx_len, max_chunks=8)
        if not chunks:
            return {"loss_before": 0.0, "loss_after": 0.0, "steps": 0, "replay_size": len(self.replay)}

        before = float(sum(self._loss_on(c).item() for c in chunks) / len(chunks))

        self.model.train()
        for step in range(self.steps_per_lesson):
            new = chunks[step % len(chunks)]
            loss = self._loss_on(new)
            n_replay = max(1, int(round(self.replay_ratio * 4)))
            for old in self.replay.sample(n_replay, rng):
                loss = loss + self._loss_on(old)
            self.optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            self.optimizer.step()

        for c in chunks:
            self.replay.add(c)
        self.model.eval()
        after = float(sum(self._loss_on(c).item() for c in chunks) / len(chunks))
        return {"loss_before": round(before, 4), "loss_after": round(after, 4),
                "steps": self.steps_per_lesson, "replay_size": len(self.replay)}

    def recall(self, prompt: str, max_new_tokens: int = 40) -> str:
        """Teste la restitution après ingestion (génération greedy)."""
        from src.generate import generate_text

        self.model.eval()
        return generate_text(self.model, self.tokenizer, prompt,
                             max_new_tokens=max_new_tokens, greedy=True)