"""J2/J3 — full calibrated quantization of Nemotron-3-Embed-8B (fresh process per format).

Pipeline (one format per invocation; NEVER derive one format from another's model object):
  1. Build the calibration set deterministically in-job (calibration/build_calibration.py).
  2. Load pinned BF16, quantize with --config over the full token-bucketed calibration.
  3. Module inventory (fail on forbidden modules), fixture fidelity vs in-process BF16
     reference (computed BEFORE quantization from the same weights).
  4. export_hf_checkpoint + tokenizer; write manifest (hashes, versions, config, marker).
  5. Leave everything in artifacts_out/<name>/ for the launcher's upload step.

Usage:
  python quantization/quantize_full.py --config NVFP4_MLP_ONLY_CFG --name nvfp4
  python quantization/quantize_full.py --config FP8_DEFAULT_CFG   --name fp8
"""

from __future__ import annotations

import argparse
import hashlib
import json
import subprocess
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
from fixtures.adapter import BASE_MODEL, BASE_REVISION, format_input, token_hash  # noqa: E402

FIXTURES = [
    ("query", "what is the capital of France"),
    ("query", "postgres index bloat mitigation"),
    ("passage", "Paris is the capital and largest city of France."),
    ("passage", "REINDEX CONCURRENTLY rebuilds an index without blocking writes."),
    ("passage", "The bank raised rates. The river bank flooded in spring."),
]
NEVER_QUANTIZE_SUBSTRINGS = ("embed_tokens", "lm_head", "norm", "layernorm")
T0 = time.monotonic()


def log(msg):
    print(f"[j2 +{time.monotonic() - T0:7.1f}s] {msg}", flush=True)


def embed_fixtures(model, tokenizer, device):
    import torch
    import torch.nn.functional as F

    out = []
    for input_type, text in FIXTURES:
        fi = format_input(text, input_type)
        enc = tokenizer(fi.text, return_tensors="pt").to(device)
        ids = enc["input_ids"][0].tolist()
        with torch.no_grad():
            o = model(**enc)
        mask = enc["attention_mask"].unsqueeze(-1).float()
        pooled = (o.last_hidden_state * mask).sum(1) / mask.sum(1)
        emb = F.normalize(pooled, p=2, dim=-1)[0].float().cpu().numpy()
        out.append({"token_hash": token_hash(ids), "emb": emb})
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", required=True)
    ap.add_argument("--name", required=True)
    ap.add_argument("--calib-batch", type=int, default=4)
    args = ap.parse_args()

    import numpy as np
    import torch
    from transformers import AutoModel, AutoTokenizer

    art = ROOT / "artifacts_out" / args.name
    art.mkdir(parents=True, exist_ok=True)
    device = "cuda"

    # 1. calibration set (deterministic; hard-fails on missing buckets)
    calib_dir = ROOT / "calibration"
    if not (calib_dir / "samples.jsonl").exists():
        log("building calibration set")
        subprocess.run(
            [sys.executable, str(calib_dir / "build_calibration.py")], check=True
        )
    samples = [json.loads(l) for l in open(calib_dir / "samples.jsonl")]
    audit = json.loads((calib_dir / "audit.json").read_text())
    log(f"calibration: {len(samples)} samples, audit={audit['status']}")

    # 2. load + BF16 reference
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, revision=BASE_REVISION)
    model = AutoModel.from_pretrained(
        BASE_MODEL, revision=BASE_REVISION, torch_dtype=torch.bfloat16,
        trust_remote_code=True,
    ).to(device).eval()
    log("computing BF16 fixture reference (pre-quantization)")
    ref = embed_fixtures(model, tokenizer, device)
    np.savez(art / "bf16_reference_embeddings.npz", **{r["token_hash"]: r["emb"] for r in ref})

    # 3. quantize over full calibration
    import modelopt.torch.quantization as mtq

    cfg = getattr(mtq, args.config)

    def forward_loop(m):
        bs = args.calib_batch
        for i in range(0, len(samples), bs):
            batch = samples[i : i + bs]
            texts = [
                format_input(s["text"], s["input_type"]).text for s in batch
            ]
            enc = tokenizer(
                texts, return_tensors="pt", padding=True, truncation=True,
                max_length=16384,
            ).to(device)
            with torch.no_grad():
                m(**enc)
            if (i // bs) % 25 == 0:
                log(f"calibration forward {i}/{len(samples)}")

    log(f"quantizing with {args.config}")
    model = mtq.quantize(model, cfg, forward_loop)

    inventory, violations = [], []
    for name, mod in model.named_modules():
        qattrs = [a for a in ("weight_quantizer", "input_quantizer") if hasattr(mod, a)]
        if not qattrs:
            continue
        enabled = any(getattr(getattr(mod, a), "is_enabled", False) for a in qattrs)
        inventory.append({"module": name, "class": type(mod).__name__, "enabled": enabled})
        if enabled and any(s in name.lower() for s in NEVER_QUANTIZE_SUBSTRINGS):
            violations.append(name)
    log(f"inventory: {len(inventory)} instrumented, violations={violations}")
    if violations:
        sys.exit(2)

    # fixture fidelity (fake-quant path, calibrated)
    fq = embed_fixtures(model, tokenizer, device)
    # fake-quant vectors enable exact runtime-vs-simulation parity later (not just
    # quant-vs-BF16): vLLM kernel output should match THESE at ~0.999
    np.savez(art / "fakequant_reference_embeddings.npz", **{r["token_hash"]: r["emb"] for r in fq})
    cosines = {}
    for r, q in zip(ref, fq):
        assert r["token_hash"] == q["token_hash"]
        cosines[r["token_hash"]] = round(float(np.dot(r["emb"], q["emb"])), 6)
    log(f"calibrated fixture cosines: {sorted(cosines.values())}")

    # 4. export + manifest
    from modelopt.torch.export import export_hf_checkpoint

    export_dir = art / "checkpoint"
    log("exporting")
    export_hf_checkpoint(model, export_dir=str(export_dir))
    tokenizer.save_pretrained(str(export_dir))

    # Post-export fix: ModelOpt writes exclude_modules patterns without the "model."
    # prefix vLLM uses for this arch ("layers.0.self_attn*" vs
    # "model.layers.0.self_attn.qkv_proj"), so none of vLLM's three matching
    # mechanisms fire and excluded (BF16) layers get treated as quantized →
    # meta-tensor crash in process_weights_after_loading. Collapse to
    # leading-wildcard patterns, which all fnmatch correctly. (runs/stage_d logs)
    qcfg_path = export_dir / "hf_quant_config.json"
    if qcfg_path.exists():
        qcfg = json.loads(qcfg_path.read_text())
        excl = qcfg.get("quantization", {}).get("exclude_modules") or []
        if excl:
            collapsed = sorted(
                {
                    ("*self_attn*" if "self_attn" in e else "*" + e.strip("*") + "*")
                    for e in excl
                }
                # lm_head doesn't exist in the encoder at quantize time, but vLLM's
                # Ministral3 runtime creates one (tied embeddings) and would treat it
                # as quantized — the exact single-meta-tensor crash from stage D.
                | {"*lm_head*"}
            )
            qcfg["quantization"]["exclude_modules"] = collapsed
            qcfg_path.write_text(json.dumps(qcfg, indent=4))
            log(f"patched exclude_modules -> {collapsed}")

    shard_hashes = {}
    for shard in sorted(export_dir.glob("*.safetensors")):
        h = hashlib.sha256()
        with open(shard, "rb") as f:
            for chunk in iter(lambda: f.read(1 << 22), b""):
                h.update(chunk)
        shard_hashes[shard.name] = h.hexdigest()

    import modelopt
    import transformers

    manifest = {
        "name": args.name,
        "quant_config": args.config,
        "base_model": BASE_MODEL,
        "base_revision": BASE_REVISION,
        "n_calibration_samples": len(samples),
        "calibration_audit": audit,
        "module_inventory_len": len(inventory),
        "fixture_cosines_vs_bf16": cosines,
        "shard_sha256": shard_hashes,
        "versions": {
            "modelopt": modelopt.__version__,
            "transformers": transformers.__version__,
            "torch": torch.__version__,
        },
        "completion_marker": True,
    }
    (art / "manifest.json").write_text(json.dumps(manifest, indent=2))
    (art / "module_inventory.json").write_text(json.dumps(inventory, indent=1))
    import shutil

    shutil.copy(calib_dir / "manifest.json", art / "calibration_manifest.json")
    shutil.copy(calib_dir / "audit.json", art / "calibration_audit.json")
    log(f"J2 complete: {args.name}, {len(shard_hashes)} shards")


if __name__ == "__main__":
    main()