"""oQ5e (imatrix-enhanced) quant of the NEW merged Nemotron, KEEPING the MTP head.

Run with the omlx tool venv python (has omlx + mlx-lm):

    /Users/david/.local/share/uv/tools/omlx/bin/python quant_oq5e_mtp.py

Env overrides: SRC, OUT, OQ_LEVEL (default 5), OUT_SUFFIX (default oq5e).
The imatrix cache is shared across levels (keyed on source signature), so an
oQ4e follow-up pass reuses the calibration for free:

    OQ_LEVEL=4 OUT_SUFFIX=oq4e /Users/david/.local/share/uv/tools/omlx/bin/python quant_oq5e_mtp.py

See KENOSISTRON2_QUANT_PLAN.md for the full runbook.
"""
import os
import sys

sys.path.insert(0, "/Users/david/AI")
import nemotron_h_mtp_model as P  # standalone copy of the serving patch
from mlx_lm.models import nemotron_h as nh

P.set_mtp_active(True)
ok = P.apply()
print(
    "patch applied:", ok,
    "| Model.sanitize is ours:",
    getattr(nh.Model.__dict__.get("sanitize"), "_omlx_nh_mtp", False),
)
assert ok, "nemotron_h MTP patch failed to apply — mtp.* would be stripped"

from omlx.oq import quantize_oq_streaming

LEVEL = float(os.environ.get("OQ_LEVEL", "5"))
LEVEL = int(LEVEL) if LEVEL == int(LEVEL) else LEVEL
SUFFIX = os.environ.get("OUT_SUFFIX", "oq5e")
SRC = os.environ.get("SRC", "/Users/david/AI/NVIDIA-Nemotron-3-Super-120B-merged2")
OUT = os.environ.get("OUT", f"/Users/david/AI/kenosistron2-{SUFFIX}-mtp")
IMATRIX = os.environ.get("IMATRIX", "/Users/david/AI/kenosistron2_imatrix.oqe")


def cb(phase, pct):
    print(f"[{SUFFIX}-mtp] {phase}: {pct:.1f}%", flush=True)


quantize_oq_streaming(
    model_path=SRC,
    output_path=OUT,
    oq_level=LEVEL,
    group_size=64,
    preserve_mtp=True,
    dtype="bfloat16",
    enhanced=True,                  # oQe: imatrix-weighted quantization (v0.5.0)
    imatrix_cache_path=IMATRIX,     # reused across levels if source unchanged
    imatrix_reuse_cache=True,
    progress_callback=cb,
)
print("done. ->", OUT)

# --- serving fix: pin the tool parser ---------------------------------------
# The Nemotron-3-Super chat_template contains `<tool_call>` + `tool_call.name`
# but NOT the adjacent `<tool_call>\n<function=` marker, so mlx_lm auto-infers
# the json_tools parser — while the model actually emits the qwen3_coder
# `<function=name><parameter=k>v</parameter></function>` format. Mismatch => tool
# calls never parse (leak as text / empty). Pin the correct parser explicitly;
# tokenizer_config `tool_parser_type` overrides inference (mlx_lm tokenizer_utils).
import json as _json
_tc_path = os.path.join(OUT, "tokenizer_config.json")
try:
    _tc = _json.load(open(_tc_path))
    if _tc.get("tool_parser_type") != "qwen3_coder":
        _tc["tool_parser_type"] = "qwen3_coder"
        _json.dump(_tc, open(_tc_path, "w"), indent=2, ensure_ascii=False)
        print("pinned tool_parser_type=qwen3_coder in", _tc_path)
    else:
        print("tool_parser_type already qwen3_coder")
except FileNotFoundError:
    print("WARN: no tokenizer_config.json in output — set tool_parser_type manually")