"""
Avalia e compara os modelos treinados.

Métricas:
  1. Perplexidade no val set
  2. Geração de texto a partir de prompts DevOps reais
  3. Top-k tokens mais prováveis para completar frases típicas

Uso:
  python eval.py
  python eval.py --prompt "apiVersion: apps/v1\nkind: Deployment"
"""

import json
import math
import argparse
import numpy as np
import torch
from pathlib import Path
from tokenizers import ByteLevelBPETokenizer

from model import GPT, GPTConfig
from train import BinaryDataset, DEVICE, DTYPE, EVAL_ITERS, SEQ_LEN, evaluate

DATA_DIR = Path(__file__).parent / "data"
CKPT_DIR = Path(__file__).parent / "checkpoints"

DEVOPS_PROMPTS = [
    "apiVersion: apps/v1\nkind: Deployment\nmetadata:\n  name:",
    "resource \"aws_eks_cluster\" \"main\" {\n  name =",
    "helm install my-release stable/nginx-ingress --namespace",
    "kubectl apply -f deployment.yaml --namespace",
    "kind: HelmRelease\napiVersion: helm.toolkit.fluxcd.io/v2beta1\nmetadata:\n  name:",
    "spec:\n  replicas: 3\n  selector:\n    matchLabels:\n      app:",
    "terraform {\n  required_providers {\n    kubernetes = {",
    "argocd app create --app-name",
]


def load_model(variant: str) -> tuple[GPT, dict]:
    ckpt_path = CKPT_DIR / f"{variant}_latest.pt"
    if not ckpt_path.exists():
        raise FileNotFoundError(f"Checkpoint não encontrado: {ckpt_path}")
    ck     = torch.load(ckpt_path, map_location=DEVICE, weights_only=False)
    cfg    = GPTConfig(**ck["config"])
    model  = GPT(cfg).to(DEVICE)
    # torch.compile adiciona prefixo "_orig_mod." — remover ao carregar
    state  = ck["model"]
    state  = {k.replace("_orig_mod.", ""): v for k, v in state.items()}
    model.load_state_dict(state)
    model.eval()
    return model, ck


def load_tokenizer() -> ByteLevelBPETokenizer:
    tok_dir = DATA_DIR / "tokenizer"
    return ByteLevelBPETokenizer(
        str(tok_dir / "vocab.json"),
        str(tok_dir / "merges.txt"),
    )


def run_eval(args):
    meta    = json.loads((DATA_DIR / "meta.json").read_text())
    val_ds  = BinaryDataset(DATA_DIR / "val.bin", SEQ_LEN)
    tok     = load_tokenizer()

    models = {}
    for v in ("baseline", "attnres"):
        try:
            m, ck = load_model(v)
            models[v] = (m, ck)
            print(f"Carregado: {v}  (step {ck.get('step', '?')})")
        except FileNotFoundError as e:
            print(f"SKIP: {e}")

    if not models:
        print("Nenhum checkpoint encontrado. Rode train.py primeiro.")
        return

    # ── 1. Perplexidade ──────────────────────────────────────────────────────
    print("\n" + "=" * 55)
    print("PERPLEXIDADE NO VAL SET")
    print("=" * 55)
    for v, (m, _) in models.items():
        loss = evaluate(m, val_ds)
        ppl  = math.exp(loss)
        params = m.n_params()
        print(f"  {v:10s}  loss {loss:.4f}  ppl {ppl:8.2f}  params {params:,}")

    # ── 2. Geração de texto ──────────────────────────────────────────────────
    prompts = [args.prompt] if args.prompt else DEVOPS_PROMPTS

    print("\n" + "=" * 55)
    print("GERAÇÃO DE TEXTO")
    print("=" * 55)

    for prompt in prompts[:4]:
        print(f"\nPrompt: {prompt!r}")
        enc_ids = tok.encode(prompt).ids
        x = torch.tensor([enc_ids], dtype=torch.long, device=DEVICE)

        for v, (m, _) in models.items():
            out_ids = m.generate(x, max_new=60, temperature=0.7, top_k=30)
            text    = tok.decode(out_ids[0].tolist())
            cont    = text[len(prompt):]
            print(f"  [{v:10s}] {cont[:120]!r}")

    # ── 3. Curvas de convergência (se ambos disponíveis) ─────────────────────
    if len(models) == 2:
        print("\n" + "=" * 55)
        print("CURVAS DE CONVERGÊNCIA (val_ppl por step)")
        print("=" * 55)
        print(f"{'Step':>6}  {'Baseline':>10}  {'AttnRes':>10}  {'Δ (A-B)':>10}")
        print("-" * 45)
        ck_b = models["baseline"][1]
        ck_a = models["attnres"][1]
        h_b  = ck_b.get("history", {})
        h_a  = ck_a.get("history", {})
        for s, lb, la in zip(h_b.get("val_steps", []),
                             h_b.get("val_loss",  []),
                             h_a.get("val_loss",  [])):
            pb, pa = math.exp(lb), math.exp(la)
            print(f"{s:6d}  {pb:10.2f}  {pa:10.2f}  {pa-pb:+10.2f}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--prompt", type=str, default=None)
    args = parser.parse_args()
    run_eval(args)