"""Phase 3 — évaluation : perplexité + génération d'échantillons.

    python -m src.evaluate --ckpt checkpoints/text_50m/best.pt \\
        --data-dir data/processed/wikipedia_fr --eval-batches 50 --sample
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import numpy as np
import torch

from src.data.prepare_text import load_bin
from src.generate import generate_text
from src.model.text_model import CortexTextLM, TextConfig
from src.tokenizer import load_tokenizer
from src.train.train_text import TrainConfig, estimate_loss, pick_autocast, pick_device


def load_checkpoint(path: str | Path, device: torch.device) -> tuple[CortexTextLM, dict]:
    ckpt = torch.load(path, map_location=device, weights_only=False)
    model = CortexTextLM(TextConfig(**ckpt["text_config"])).to(device)
    model.load_state_dict(ckpt["model"])
    model.eval()
    return model, ckpt


@torch.no_grad()
def perplexity(model: CortexTextLM, data: np.memmap, cfg: TrainConfig,
               device: torch.device, autocast_ctx) -> float:
    metrics = estimate_loss(model, {"val": data}, cfg, device, autocast_ctx)
    return math.exp(min(metrics["val"], 20)), metrics["val"]


def main() -> None:
    parser = argparse.ArgumentParser(description="Évaluer un checkpoint texte Cortex-0.")
    parser.add_argument("--ckpt", type=Path, required=True)
    parser.add_argument("--data-dir", type=Path, required=True)
    parser.add_argument("--split", default="val")
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--eval-batches", type=int, default=50)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--sample", action="store_true")
    parser.add_argument("--prompt", default="La France est")
    parser.add_argument("--max-new-tokens", type=int, default=60)
    args = parser.parse_args()

    device = pick_device(args.device)
    model, ckpt = load_checkpoint(args.ckpt, device)
    data = load_bin(Path(args.data_dir) / f"{args.split}.bin")
    cfg = TrainConfig(batch_size=args.batch_size, eval_batches=args.eval_batches,
                      ctx_len=model.config.ctx_len, device=args.device)
    autocast_ctx, _ = pick_autocast(device, "float32")
    ppl, loss = perplexity(model, data, cfg, device, autocast_ctx)

    result = {"checkpoint": str(args.ckpt), "step": ckpt.get("step"),
              "split": args.split, "val_loss": loss, "perplexity": ppl,
              "n_params": model.count_parameters()}
    if args.sample:
        tokenizer = load_tokenizer()
        result["samples"] = [
            generate_text(model, tokenizer, args.prompt, max_new_tokens=args.max_new_tokens,
                          temperature=t, top_k=40)
            for t in (0.7, 1.0)
        ]
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()