#!/usr/bin/env python3
"""Fast baseline eval for prod9 champion checkpoint.

Runs BPB, PPL, and generates sample text for human inspection.
No training loop — eval only.
"""
import os
# WSL CUDA library path — must be set BEFORE any htm_rust import (Rust cudarc)
os.environ["LD_LIBRARY_PATH"] = "/usr/lib/wsl/lib"

import sys, math, torch, gc

os.environ.setdefault("HYDRA_SEQ_LEN", "1024")
os.environ.setdefault("HYDRA_N_LAYER", "8")
os.environ.setdefault("HYDRA_D_MODEL", "256")
os.environ.setdefault("HYDRA_EXPAND", "2")
os.environ.setdefault("HYDRA_HEADDIM", "32")
os.environ.setdefault("HYDRA_D_STATE", "64")
os.environ.setdefault("HYDRA_ENGRAM_N_COLUMNS", "1024")
os.environ.setdefault("HYDRA_ENGRAM_TOPK", "64")
os.environ.setdefault("HYDRA_BATCH_SIZE", "1")
os.environ.setdefault("HYDRA_TOTAL_BATCH", "131072")
os.environ.setdefault("HYDRA_SAMPLED_SOFTMAX", "1024")
os.environ.setdefault("HYDRA_MTP_K", "1")
os.environ.setdefault("HYDRA_GDN_LAYERS", "")
os.environ.setdefault("HYDRA_HYENA_LAYERS", "")
os.environ.setdefault("HYDRA_USE_MDLM", "0")
os.environ.setdefault("HYDRA_SKIP_FACTUAL_EVAL", "1")
os.environ.setdefault("HYDRA_EVAL_BATCH", "1")
os.environ.setdefault("HYDRA_DISABLE_FUSED_SDR_TRITON", "1")
os.environ.setdefault("HYDRA_FUSED_SDR_PROJECT", "0")
os.environ.setdefault("HYDRA_HTM_FUSED", "1")
os.environ.setdefault("HYDRA_HTM_BATCHED_FUSED", "1")

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import prepare_nemotron as _p_nemo
_p_nemo.ensure_tokenizer()

from prepare import Tokenizer
from hydra.model import PostSemClawModel
from hydra.config import PostSemClawConfig
from prepare_nemotron import evaluate_bpb

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"[eval] device={device}", flush=True)

tokenizer = Tokenizer.from_directory()
vocab_size = tokenizer.get_vocab_size()
print(f"[eval] vocab_size={vocab_size}", flush=True)

cfg = PostSemClawConfig(
    vocab_size=tokenizer.get_vocab_size(),
    n_layer=int(os.environ["HYDRA_N_LAYER"]),
    d_model=int(os.environ["HYDRA_D_MODEL"]),
    d_state=int(os.environ["HYDRA_D_STATE"]),
    headdim=int(os.environ["HYDRA_HEADDIM"]),
    expand=int(os.environ["HYDRA_EXPAND"]),
    engram_n_columns=int(os.environ["HYDRA_ENGRAM_N_COLUMNS"]),
    sequence_len=int(os.environ["HYDRA_SEQ_LEN"]),
)
model = PostSemClawModel(cfg).to(device)
print(f"[eval] params={sum(p.numel() for p in model.parameters())/1e6:.2f}M", flush=True)

CKPT = os.environ.get("HYDRA_RESUME_CKPT", "checkpoints/prod9_champion/bootstrap/prod9_b16_step21500_trainbest_bpb0p8726_latest.pt")
sd = torch.load(CKPT, map_location=device, weights_only=False)
model.load_state_dict(sd["model_state_dict"])
step = sd.get("step", "?")
epoch = sd.get("epoch", "?")
print(f"[eval] loaded ckpt step={step} epoch={epoch} from {CKPT}", flush=True)

gc.collect()
if device.type == "cuda":
    torch.cuda.empty_cache()

model.eval()
with torch.no_grad():
    print("[eval] running evaluate_bpb ...", flush=True)
    bpb = evaluate_bpb(model, tokenizer, int(os.environ["HYDRA_EVAL_BATCH"]))
    ppl = 2.0 ** bpb
    print(f"[EVAL_RESULT] bpb={bpb:.6f} ppl={ppl:.4f}", flush=True)

# Sample generation
print("[eval] generating sample completions ...", flush=True)
model.train()  # needed for engram/htm learning during generation
gc.collect()

# Pick a neutral English prompt
prompt_texts = [
    "In 1492, Christopher Columbus sailed across the Atlantic Ocean and",
    "The theory of relativity, developed by Albert Einstein, explains that",
    "A well-balanced diet should include plenty of vegetables, fruits, and",
    "The history of artificial intelligence began in the 1950s when",
    "To bake a simple loaf of bread, you need flour, water, yeast, and",
]

max_new_tokens = 64
for prompt in prompt_texts[:3]:
    input_ids = torch.tensor([tokenizer.encode(prompt)], device=device)
    with torch.no_grad():
        for _ in range(max_new_tokens):
            logits = model(input_ids, input_ids, reduction="none")  # dummy y
            next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
            input_ids = torch.cat([input_ids, next_token], dim=1)
    completion = tokenizer.decode(input_ids[0].tolist())
    print(f"\n--- PROMPT: {prompt}")
    print(f"--- COMPLETION: {completion}")
    print("---")
    gc.collect()

print("[eval] baseline eval complete", flush=True)