"""Génération autoregressive pour Cortex-0 — échantillonnage écrit à la main.

Implémente :
* température (recherche du logits / T) ;
* top-k (garder les k logits les plus grands) ;
* top-p / nucleus (garder le plus petit ensemble de probabilité >= p).

Aucune fonction `model.generate` n'est utilisée : la boucle est explicite.
"""

from __future__ import annotations

import torch

from src.model.text_model import CortexTextLM
from src.tokenizer import CortexTokenizer
from src.tokenizer.special_tokens import BOS_ID, EOS_ID


def _apply_top_k(logits: torch.Tensor, k: int) -> torch.Tensor:
    if k is None or k <= 0:
        return logits
    k = min(k, logits.size(-1))
    kth_value = torch.topk(logits, k, dim=-1).values[..., -1, None]
    return logits.masked_fill(logits < kth_value, float("-inf"))


def _apply_top_p(logits: torch.Tensor, p: float) -> torch.Tensor:
    if p is None or p >= 1.0:
        return logits
    sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
    probs = torch.softmax(sorted_logits, dim=-1)
    cumulative = torch.cumsum(probs, dim=-1)
    # On retire les tokens au-delà du seuil (en gardant le premier qui dépasse p).
    remove = cumulative - probs > p
    sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
    # Remise dans l'ordre d'origine.
    return logits.scatter(-1, sorted_idx, sorted_logits)


@torch.no_grad()
def generate(
    model: CortexTextLM,
    input_ids: torch.Tensor,
    max_new_tokens: int = 50,
    temperature: float = 1.0,
    top_k: int | None = None,
    top_p: float | None = None,
    greedy: bool = False,
    eos_id: int | None = EOS_ID,
) -> torch.Tensor:
    """Génère ``max_new_tokens`` tokens en continuation de ``input_ids`` (B, T)."""
    model.eval()
    ctx = model.config.ctx_len
    for _ in range(max_new_tokens):
        idx_cond = input_ids[:, -ctx:]
        logits, _ = model(idx_cond)
        logits = logits[:, -1, :]  # (B, vocab)

        if greedy:
            next_id = torch.argmax(logits, dim=-1, keepdim=True)
        else:
            logits = logits / max(temperature, 1e-6)
            logits = _apply_top_k(logits, top_k)
            logits = _apply_top_p(logits, top_p)
            probs = torch.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)

        input_ids = torch.cat([input_ids, next_id], dim=1)
        if eos_id is not None and (next_id == eos_id).all():
            break
    return input_ids


def generate_text(
    model: CortexTextLM,
    tokenizer: CortexTokenizer,
    prompt: str,
    max_new_tokens: int = 50,
    **kwargs,
) -> str:
    """Enveloppe pratique : encode le prompt, génère, décode la suite."""
    ids = [BOS_ID, *tokenizer.encode(prompt)]
    inp = torch.tensor([ids], dtype=torch.long)
    out = generate(model, inp, max_new_tokens=max_new_tokens, **kwargs)
    return tokenizer.decode(out[0].tolist(), skip_special_tokens=True)