"""P3b — convert nvidia/Nemotron-3-Embed-8B-BF16 to a quantized MLX checkpoint.

The NVIDIA HF checkpoint's weight names (embed_tokens/layers.N/norm, no "model."
prefix) and config keys match the proven bidirectional MLX implementation bundled
with the mlx-community 1B port exactly (see runs/mlx_contract_1b_bf16.json: the
implementation is 0.9999-exact vs Transformers). Conversion is therefore:
mmap-load shards -> load into NemotronEmbedModel -> nn.quantize -> save.

One recipe per process invocation (24GB unified memory: 16GB bf16 mmap + the
quantized materialization must not coexist across recipes).

Usage:
  source ~/.venvs/nq-mlx/bin/activate
  python mlx/convert_8b.py --bits 4 --group-size 32 \
      --src ~/.cache/nq-models/nv-8b-bf16 --out ~/.cache/nq-models/mlx-8b-4bit-gs32
  python mlx/convert_8b.py --bits 4 --group-size 32 --mixed \
      --out ~/.cache/nq-models/mlx-8b-mixed-gs32   # embeddings kept 8-bit
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import shutil
import sys
from pathlib import Path

MLX_1B_DIR = Path.home() / ".cache/nq-models/mlx-1b-4bit"  # source of the impl


def load_impl():
    spec = importlib.util.spec_from_file_location(
        "nemotron3_embed_mlx", MLX_1B_DIR / "nemotron3_embed_mlx.py"
    )
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", default=str(Path.home() / ".cache/nq-models/nv-8b-bf16"))
    ap.add_argument("--out", required=True)
    ap.add_argument("--bits", type=int, required=True)
    ap.add_argument("--group-size", type=int, required=True)
    ap.add_argument(
        "--mixed",
        action="store_true",
        help="keep embed_tokens at 8-bit (body at --bits)",
    )
    args = ap.parse_args()

    import mlx.core as mx
    import mlx.nn as nn
    from mlx_lm.models.ministral3 import ModelArgs

    src, out = Path(args.src).expanduser(), Path(args.out).expanduser()
    out.mkdir(parents=True, exist_ok=True)
    mod = load_impl()

    cfg = json.loads((src / "config.json").read_text())
    model = mod.NemotronEmbedModel(ModelArgs.from_dict(cfg))

    weights = {}
    for shard in sorted(src.glob("model-*.safetensors")):
        weights.update(mx.load(str(shard)))  # mmap — not materialized yet
    print(
        f"loaded {len(weights)} tensors from {len(list(src.glob('model-*.safetensors')))} shards"
    )
    model.load_weights(list(weights.items()))

    if args.mixed:

        def class_predicate(path, module):
            if not hasattr(module, "to_quantized"):
                return False
            if "embed_tokens" in path:
                return {"group_size": args.group_size, "bits": 8}
            return True

        nn.quantize(
            model,
            group_size=args.group_size,
            bits=args.bits,
            class_predicate=class_predicate,
        )
        quant_cfg = {
            "group_size": args.group_size,
            "bits": args.bits,
            "mode": "affine",
            "overrides": {"embed_tokens": {"bits": 8}},
        }
    else:
        nn.quantize(model, group_size=args.group_size, bits=args.bits)
        quant_cfg = {"group_size": args.group_size, "bits": args.bits, "mode": "affine"}

    mx.eval(model.parameters())

    from mlx.utils import tree_flatten

    flat = dict(tree_flatten(model.parameters()))
    mx.save_safetensors(
        str(out / "model.safetensors"), flat, metadata={"format": "mlx"}
    )

    cfg_out = dict(cfg)
    cfg_out["quantization"] = quant_cfg
    cfg_out["quantization_config"] = quant_cfg
    (out / "config.json").write_text(json.dumps(cfg_out, indent=2))
    for f in ("tokenizer.json", "tokenizer_config.json"):
        shutil.copy(src / f, out / f)
    shutil.copy(MLX_1B_DIR / "nemotron3_embed_mlx.py", out / "nemotron3_embed_mlx.py")

    size_gb = (out / "model.safetensors").stat().st_size / 1e9
    print(
        json.dumps(
            {
                "out": str(out),
                "quantization": quant_cfg,
                "n_tensors": len(flat),
                "size_gb": round(size_gb, 2),
            }
        )
    )


if __name__ == "__main__":
    main()