"""Measure quantization drift against BF16; never used for calibration."""
import gc
import fcntl
import hashlib
import json
import math
from pathlib import Path

import mlx.core as mx
from mlx_lm import load

from release import ROOT, REV, VARIANTS, write_json

gpu_lock = (ROOT / "artifacts/.gpu.lock").open("a")
fcntl.flock(gpu_lock, fcntl.LOCK_EX)

corpus_path = ROOT / "evaluation.json"
corpus = json.loads(corpus_path.read_text())
corpus_hash = hashlib.sha256(corpus_path.read_bytes()).hexdigest()
reference_file = ROOT / "artifacts/bf16-reference-logprobs.safetensors"

if not reference_file.exists():
    model, tokenizer = load(str(ROOT / "artifacts/bf16"))
    reference = {}
    for i, item in enumerate(corpus):
        ids = tokenizer.encode(item["text"])
        logits = model(mx.array([ids[:-1]])).astype(mx.float32)
        logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
        mx.eval(logprobs)
        reference[str(i)] = logprobs
        del logits
    mx.save_safetensors(str(reference_file), reference, metadata={"corpus_sha256": corpus_hash, "source_revision": REV})
    del model, tokenizer, reference, logprobs
    gc.collect()
    mx.clear_cache()

reference, metadata = mx.load(str(reference_file), return_metadata=True)
assert metadata["corpus_sha256"] == corpus_hash
assert metadata["source_revision"] == REV
summary = {}
for name in VARIANTS:
    if not (ROOT / "artifacts" / name / "conversion.json").exists():
        continue
    report_path = ROOT / "reports" / f"{name}.fidelity.json"
    if report_path.exists():
        summary[name] = json.loads(report_path.read_text())
        continue
    rows = []
    if name != "bf16":
        model, tokenizer = load(str(ROOT / "artifacts" / name))
    for i, item in enumerate(corpus):
        teacher = reference[str(i)]
        if name == "bf16":
            student = teacher
        else:
            ids = tokenizer.encode(item["text"])
            logits = model(mx.array([ids[:-1]])).astype(mx.float32)
            student = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
        kl = float(mx.mean(mx.sum(mx.exp(teacher) * (teacher - student), axis=-1)).item())
        agreement = float(mx.mean(mx.argmax(teacher, axis=-1) == mx.argmax(student, axis=-1)).item())
        assert math.isfinite(kl) and kl >= -1e-6
        rows.append({"category": item["category"], "tokens": teacher.shape[1],
                     "kl_bf16_to_quant_nats": max(0.0, kl), "top1_agreement": agreement})
    count = sum(r["tokens"] for r in rows)
    report = {"variant": name, "corpus_sha256": corpus_hash, "scored_tokens": count,
              "mean_kl_bf16_to_quant_nats": sum(r["tokens"] * r["kl_bf16_to_quant_nats"] for r in rows) / count,
              "top1_agreement": sum(r["tokens"] * r["top1_agreement"] for r in rows) / count,
              "samples": rows,
              "scope": "Teacher-forced output-distribution fidelity on a small fixed synthetic corpus; not a capability benchmark"}
    write_json(report_path, report)
    summary[name] = report
    print(name, report["mean_kl_bf16_to_quant_nats"], report["top1_agreement"], flush=True)
    if name != "bf16":
        del model, tokenizer, logits
    del student
    gc.collect()
    mx.clear_cache()
write_json(ROOT / "reports/fidelity-summary.json", summary)