#!/usr/bin/env python3
"""Load each delivered checkpoint on CPU and verify finite logits."""

from __future__ import annotations

import gc
import json
import os
import platform
import sys
import traceback
from pathlib import Path


THREADS = "4"
for variable in (
    "OMP_NUM_THREADS",
    "MKL_NUM_THREADS",
    "OPENBLAS_NUM_THREADS",
    "NUMEXPR_NUM_THREADS",
    "VECLIB_MAXIMUM_THREADS",
    "BLIS_NUM_THREADS",
):
    os.environ.setdefault(variable, THREADS)
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")

import torch  # noqa: E402
import safetensors  # noqa: E402
import transformers  # noqa: E402
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer  # noqa: E402


ROOT = Path(__file__).resolve().parents[1]
PROMPT = 'def rmsd(a, b):\n    """Root-mean-square deviation."""\n'
TARGETS = {
    "pretrained/codegpt_multilingual_5epoch": ROOT / "models/pretrained/codegpt_multilingual_5epoch",
    "pretrained/gpt2_124m_code_5epoch": ROOT / "models/pretrained/gpt2_124m_code_5epoch",
    "pretrained/qwen25_coder_7b_cpt": ROOT / "models/pretrained/qwen25_coder_7b_cpt",
    "pretrained/stage1_cpt": ROOT / "models/pretrained/stage1_cpt",
    "sft/sft_f3_refined_instruct": ROOT / "models/sft/sft_f3_refined_instruct",
}
for checkpoint in sorted((ROOT / "models/sft").iterdir()):
    if checkpoint.is_dir() and (checkpoint / "config.json").is_file():
        TARGETS.setdefault(f"sft/{checkpoint.name}", checkpoint)


def check(name: str, path: Path) -> dict:
    result = {"model": name, "path": str(path.relative_to(ROOT)), "ok": False}
    try:
        config = AutoConfig.from_pretrained(path)
        result.update(
            architecture=(config.architectures or [type(config).__name__])[0],
            hidden_size=getattr(config, "hidden_size", None),
            layers=getattr(config, "num_hidden_layers", None),
            vocab_size=getattr(config, "vocab_size", None),
        )
        tokenizer = AutoTokenizer.from_pretrained(path)
        model = AutoModelForCausalLM.from_pretrained(
            path,
            dtype="auto",
            low_cpu_mem_usage=True,
        )
        model.eval()
        inputs = tokenizer(PROMPT, return_tensors="pt")
        with torch.no_grad():
            logits = model(**inputs).logits
        last = logits[0, -1].float()
        result.update(
            tokenizer=type(tokenizer).__name__,
            parameters_million=round(sum(p.numel() for p in model.parameters()) / 1e6, 1),
            logits_shape=list(logits.shape),
            finite=bool(torch.isfinite(logits).all()),
            top1_zscore=round(((last.max() - last.mean()) / last.std()).item(), 2),
        )
        result["ok"] = result["finite"] and result["top1_zscore"] > 3.0
        del model, tokenizer, inputs, logits, last
        gc.collect()
    except Exception as error:  # noqa: BLE001
        result["error"] = f"{type(error).__name__}: {error}"
        result["traceback"] = traceback.format_exc(limit=3)
    return result


def main() -> int:
    torch.set_num_threads(int(THREADS))
    results = [check(name, path) for name, path in TARGETS.items()]
    report = {
        "validation_environment": {
            "python": platform.python_version(),
            "torch": torch.__version__,
            "transformers": transformers.__version__,
            "safetensors": safetensors.__version__,
            "device": "cpu",
            "threads": int(THREADS),
        },
        "results": results,
    }
    output = ROOT / "evidence/smoke_test_results.json"
    output.write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    for result in results:
        status = "PASS" if result["ok"] else "FAIL"
        detail = f"{result.get('parameters_million', '?')}M, z={result.get('top1_zscore', '?')}"
        print(f"{status}: {result['model']} ({detail})")
    passed = sum(result["ok"] for result in results)
    print(f"{passed}/{len(results)} checkpoints passed; results: {output}")
    return 0 if passed == len(results) else 1


if __name__ == "__main__":
    sys.exit(main())