"""Reproducible conversion, release gates, benchmarking, and publication."""
import argparse
import fcntl
import hashlib
import json
import math
import shutil
import statistics
import struct
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
REV = "c1093c501cc6460427c0ac634db7223d0450c0ee"
SOURCE = "Nanbeige/Nanbeige4.3-33B"
VARIANTS = {"bf16": None, "4bit": (4, 64), "6bit": (6, 64), "8bit": (8, 64),
            "4bit-g32": (4, 32), "4bit-g128": (4, 128), "mixed-4-6-g64": (4, 64)}


def write_json(path, value):
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")
    tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n")
    tmp.replace(path)


def inventory(path):
    tensors = {}
    for f in sorted(path.glob("*.safetensors")):
        with f.open("rb") as stream:
            length = struct.unpack("<Q", stream.read(8))[0]
            header = json.loads(stream.read(length))
        for key, value in header.items():
            if key == "__metadata__":
                continue
            assert key not in tensors, key
            tensors[key] = {k: value[k] for k in ("dtype", "shape")}
    assert tensors, "No weights"
    index = json.loads((path / "model.safetensors.index.json").read_text())
    assert set(index["weight_map"]) == set(tensors)
    return tensors


def convert(name):
    from mlx_lm.convert import convert as mlx_convert
    path = ROOT / "artifacts" / name
    source = ROOT / ("source" if name == "bf16" else "artifacts/bf16")
    options = dict(hf_path=str(source), mlx_path=str(path), dtype="bfloat16")
    if VARIANTS[name]:
        bits, group = VARIANTS[name]
        options.update(quantize=True, q_bits=bits, q_group_size=group, q_mode="affine")
        if name.startswith("mixed"):
            options["quant_predicate"] = "mixed_4_6"
    mlx_convert(**options)
    # Preserve upstream tokenizer assets byte-for-byte, including the template.
    for pattern in ("tokenizer*", "special_tokens_map.json", "added_tokens.json", "generation_config.json"):
        for file in (ROOT / "source").glob(pattern):
            if file.is_file():
                shutil.copy2(file, path / file.name)
    write_json(path / "conversion.json", {"source": SOURCE, "revision": REV,
               "variant": name, "options": options})


def validate(name, model_path=None):
    import mlx.core as mx
    import mlx.nn as nn
    from mlx_lm import load, stream_generate
    from mlx_lm.sample_utils import make_sampler
    from transformers import AutoTokenizer
    path = Path(model_path) if model_path else ROOT / "artifacts" / name
    weights = inventory(path)
    config = json.loads((path / "config.json").read_text())
    original = json.loads((ROOT / "source/config.json").read_text())
    for key in ("rope_theta", "max_position_embeddings", "vocab_size", "num_hidden_layers", "num_attention_heads", "num_key_value_heads", "hidden_size"):
        assert config[key] == original[key], key
    if name == "bf16":
        assert weights == inventory(ROOT / "source"), "BF16 tensor inventory changed"
        assert all(v["dtype"] == "BF16" for v in weights.values())
    else:
        assert config.get("quantization"), "Missing quantization metadata"
        assert all(v["dtype"] in ("BF16", "U32") for v in weights.values())
    source_tok = AutoTokenizer.from_pretrained(ROOT / "source")
    model, tok = load(str(path))
    for text in ("Hello, world!", "请解释量化。", "def add(a, b):\n    return a+b", "<think></think>"):
        assert source_tok.encode(text) == tok.encode(text)
    tools = [{"type": "function", "function": {"name": "weather", "description": "Get weather", "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}]
    messages = [{"role": "user", "content": "Hello"}, {"role": "assistant", "content": "<think>Consider a greeting.</think>Hello!"}, {"role": "user", "content": "What is the weather in Boston?"}]
    for thinking in (True, False):
        for preserve in (True, False):
            for fmt in ("xml", "json"):
                kw = dict(tokenize=False, add_generation_prompt=True, enable_thinking=thinking, preserve_thinking=preserve, tools=tools, tool_call_format=fmt)
                assert source_tok.apply_chat_template(messages, **kw) == tok.apply_chat_template(messages, **kw)
    smoke = []
    for prompt in ("In one sentence, explain why the sky looks blue.", "请用一句中文解释为什么天空是蓝色的。"):
        rendered = tok.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True, enable_thinking=False)
        logits = model(mx.array([tok.encode(rendered)]))
        assert bool(mx.all(mx.isfinite(logits)).item()), "Non-finite logits"
        del logits
        text = "".join(r.text for r in stream_generate(model, tok, rendered, max_tokens=128, sampler=make_sampler(temp=0)))
        assert text.strip(), "Empty generation"
        smoke.append({"prompt": prompt, "response": text})
        print(text, flush=True)
    corpus_file = ROOT / "evaluation.json"
    corpus = json.loads(corpus_file.read_text())
    losses = []
    for item in corpus:
        ids = tok.encode(item["text"])
        x = mx.array([ids])
        logits = model(x[:, :-1]).astype(mx.float32)
        loss = float(nn.losses.cross_entropy(logits, x[:, 1:], reduction="mean").item())
        assert math.isfinite(loss)
        losses.append({"category": item["category"], "tokens": len(ids)-1, "nll": loss})
        del logits, x
        mx.clear_cache()
    nll = sum(x["tokens"] * x["nll"] for x in losses) / sum(x["tokens"] for x in losses)
    report = {"variant": name, "tensor_count": len(weights), "smoke": smoke,
              "corpus_sha256": hashlib.sha256(corpus_file.read_bytes()).hexdigest(),
              "losses": losses, "mean_nll": nll, "perplexity": math.exp(nll),
              "weight_bytes": sum(f.stat().st_size for f in path.glob("*.safetensors")),
              "template_cases_passed": 8, "status": "passed"}
    if name != "bf16":
        baseline = json.loads((ROOT / "reports/bf16.validation.json").read_text())
        report["nll_delta_bf16"] = nll - baseline["mean_nll"]
    write_json(ROOT / "reports" / f"{name}.validation.json", report)
    print(json.dumps(report, ensure_ascii=False), flush=True)


def benchmark(name):
    import mlx.core as mx
    from mlx_lm import load, stream_generate
    from mlx_lm.sample_utils import make_sampler
    model, tok = load(str(ROOT / "artifacts" / name))
    tok._eos_token_ids = set()
    results = []
    for length in (512, 4096, 16384):
        mx.random.seed(42)
        prompt = mx.random.randint(0, model.args.vocab_size, (length,)).tolist()
        trials = []
        for trial in range(4):
            mx.reset_peak_memory()
            for response in stream_generate(model, tok, prompt, max_tokens=256, sampler=make_sampler(temp=0), prefill_step_size=512):
                pass
            assert response.generation_tokens == 256
            row = {k: getattr(response, k) for k in ("prompt_tps", "generation_tps", "peak_memory")}
            print(name, length, trial, row, flush=True)
            if trial:
                trials.append(row)
            mx.clear_cache()
        results.append({"prompt_tokens": length, "generation_tokens": 256, "trials": trials,
                        "median": {k: statistics.median(t[k] for t in trials) for k in trials[0]}})
        write_json(ROOT / "reports" / f"{name}.benchmark.json", {"hardware": "Apple M5 Max, 128 GB", "kv_cache": "unquantized", "prefill_step_size": 512, "results": results})


def publish(name):
    from huggingface_hub import HfApi
    path = ROOT / "artifacts" / name
    report = json.loads((ROOT / "reports" / f"{name}.validation.json").read_text())
    assert report["status"] == "passed"
    repo = f"fcmeyer/Nanbeige4.3-33B-MLX-{name}"
    conversion = json.loads((path / "conversion.json").read_text())
    conversion["options"]["hf_path"] = "source" if name == "bf16" else "artifacts/bf16"
    conversion["options"]["mlx_path"] = f"artifacts/{name}"
    write_json(path / "conversion.json", conversion)
    for f in (ROOT / "reports").glob(f"{name}.*.json"):
        if not f.name.endswith(".publication.json"):
            shutil.copy2(f, path / f.name)
    if (ROOT / "reports/environment.json").exists():
        shutil.copy2(ROOT / "reports/environment.json", path / "environment.json")
    for f in ("requirements.lock.txt", "evaluation.json"):
        shutil.copy2(ROOT / f, path / f)
    (path / "scripts").mkdir(exist_ok=True)
    shutil.copy2(ROOT / "scripts/release.py", path / "scripts/release.py")
    if (ROOT / "reports" / f"{name}.fidelity.json").exists():
        shutil.copy2(ROOT / "scripts/fidelity.py", path / "scripts/fidelity.py")
    # Earlier initial releases placed this script at the repository root.
    # Its relative project paths require the scripts/ layout for reproduction.
    if (path / "release.py").exists():
        (path / "release.py").unlink()
    if (ROOT / "LICENSE").exists():
        shutil.copy2(ROOT / "LICENSE", path / "LICENSE")
    quant = "Unquantized BF16." if name == "bf16" else f"Affine quantization: {VARIANTS[name][0]} base bits, group size {VARIANTS[name][1]}." + (" MLX-LM mixed_4_6 recipe." if name.startswith("mixed") else "")
    performance = "Performance benchmarking is pending."
    benchmark_file = ROOT / "reports" / f"{name}.benchmark.json"
    if benchmark_file.exists():
        benchmark_report = json.loads(benchmark_file.read_text())
        performance = "Measured on Apple M5 Max with 128 GB unified memory, unquantized KV cache, 512-token prefill steps, 256 output tokens; one warm-up and three measured trials per length. Prompts are seeded synthetic token sequences. Values are medians.\n\n| Input tokens | Prefill tok/s | Decode tok/s | Peak GB |\n|---:|---:|---:|---:|\n"
        for row in benchmark_report["results"]:
            med = row["median"]
            performance += f"| {row['prompt_tokens']} | {med['prompt_tps']:.2f} | {med['generation_tps']:.2f} | {med['peak_memory']:.2f} |\n"
        performance += "\nPeak memory is MLX's reported peak allocation, excluding other applications and OS memory. See `environment.json` for OS/runtime details.\n"
    delta = f"NLL difference versus BF16: {report['nll_delta_bf16']:+.6f}." if "nll_delta_bf16" in report else ""
    fidelity_text = ""
    fidelity_path = ROOT / "reports" / f"{name}.fidelity.json"
    if fidelity_path.exists():
        fidelity = json.loads(fidelity_path.read_text())
        fidelity_text = f"Teacher-forced fidelity versus BF16 on the same {fidelity['scored_tokens']} tokens: mean KL(BF16 || quant) {fidelity['mean_kl_bf16_to_quant_nats']:.6f} nats/token; top-1 token agreement {100*fidelity['top1_agreement']:.2f}%. Lower KL means closer output distributions on these inputs. This does not establish downstream capability equivalence."
    card = f'''---
license: apache-2.0
library_name: mlx
pipeline_tag: text-generation
base_model: {SOURCE}
language:
- en
- zh
tags:
- mlx
- nanbeige
---
# Nanbeige4.3-33B MLX {name}

MLX conversion of [{SOURCE}](https://huggingface.co/{SOURCE}), by the Nanbeige Team.
Source revision: `{REV}`. {quant}

## Usage
```python
from mlx_lm import load, generate
model, tokenizer = load("{repo}")
prompt = tokenizer.apply_chat_template(
    [{{"role": "user", "content": "Hello!"}}],
    tokenize=False, add_generation_prompt=True, enable_thinking=False)
print(generate(model, tokenizer, prompt=prompt, max_tokens=256))
```

The upstream chat template is preserved. `enable_thinking` controls current reasoning;
`preserve_thinking=True` retains previous reasoning for agent workflows. Upstream recommends
`tool_call_format="xml"`; JSON is also supported. Template support does not certify any server's tool parser.

## Validation
Weight files: {report['weight_bytes']:,} bytes. Tensor count: {report['tensor_count']}.
English/Chinese generation, finite smoke logits, tokenizer equivalence, and eight thinking/tool template cases passed.
Small fixed synthetic diagnostic corpus ({sum(x['tokens'] for x in report['losses'])} scored tokens): mean token NLL {report['mean_nll']:.6f}, perplexity {report['perplexity']:.6f}.
{delta}
{fidelity_text}
This is a conversion regression check, not a standardized capability benchmark or a claim of benchmark equivalence.
See the included validation JSON for outputs and per-category measurements.
The source's 524,288-token context setting is retained; full 512K context has not been validated.
Performance results, when included, apply only to the stated hardware and tested lengths.

## Performance
{performance}

## Reproduction
Use the included `requirements.lock.txt`, `conversion.json`, `evaluation.json`, and `scripts/release.py`.
From the repository root:
```sh
uv venv --python 3.12 .venv
uv pip install --python .venv/bin/python -r requirements.lock.txt
.venv/bin/hf download {SOURCE} --revision {REV} --local-dir source
.venv/bin/python scripts/release.py convert bf16
.venv/bin/python scripts/release.py validate bf16
```
For quantized variants, then run `.venv/bin/python scripts/release.py convert {name}` and
`.venv/bin/python scripts/release.py validate {name}`. All quantized variants derive independently from BF16.

## Attribution
Original model: Nanbeige Team, Nanbeige4.3-33B (2026), Apache-2.0.
Conversion and measurements: fcmeyer. No claim is made that this is universally the best quantization.
'''
    (path / "README.md").write_text(card)
    api = HfApi()
    api.create_repo(repo_id=repo, repo_type="model", private=False, exist_ok=True)
    api.upload_folder(repo_id=repo, repo_type="model", folder_path=path,
                      ignore_patterns=[".cache/**"],
                      delete_patterns=["release.py"],
                      commit_message=f"Publish validated MLX {name} conversion and measurements")
    info = api.model_info(repo, files_metadata=True)
    remote = {f.rfilename: f.size for f in info.siblings}
    for f in path.iterdir():
        if f.is_file():
            assert remote.get(f.name) == f.stat().st_size, f.name
    write_json(ROOT / "reports" / f"{name}.publication.json", {"repo": repo, "revision": info.sha, "inventory_verified": True})
    print(f"PUBLISHED https://huggingface.co/{repo} revision={info.sha}", flush=True)


def verify_remote(name):
    """Reload published config/tokenizer with locally hash-verified remote weights."""
    from huggingface_hub import HfApi, snapshot_download
    from mlx_lm import load, stream_generate
    from mlx_lm.sample_utils import make_sampler
    report_path = ROOT / "reports" / f"{name}.publication.json"
    report = json.loads(report_path.read_text())
    repo, revision = report["repo"], report["revision"]
    path = ROOT / "artifacts" / (name + "-remote-check")
    snapshot_download(repo, revision=revision, local_dir=path, ignore_patterns=["*.safetensors"])
    info = HfApi().model_info(repo, revision=revision, files_metadata=True)
    for file in info.siblings:
        if not file.rfilename.endswith(".safetensors"):
            continue
        local = ROOT / "artifacts" / name / file.rfilename
        with local.open("rb") as stream:
            digest = hashlib.file_digest(stream, "sha256").hexdigest()
        assert file.lfs and digest == file.lfs.sha256, file.rfilename
        target = path / file.rfilename
        if not target.exists():
            target.symlink_to(local)
    model, tokenizer = load(str(path))
    prompt = tokenizer.apply_chat_template([{"role": "user", "content": "Say hello."}], tokenize=False, add_generation_prompt=True, enable_thinking=False)
    text = "".join(r.text for r in stream_generate(model, tokenizer, prompt, max_tokens=32, sampler=make_sampler(temp=0)))
    assert text.strip()
    report.update(remote_reload="passed", weight_sha256_verified=True, reload_response=text,
                  verification_method="Fresh published metadata/tokenizer; local weights SHA256 matched to remote LFS hashes, avoiding redundant transfer")
    write_json(report_path, report)
    print(json.dumps(report), flush=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("action", choices=["convert", "validate", "benchmark", "publish", "verify_remote"])
    parser.add_argument("variant", choices=VARIANTS)
    args = parser.parse_args()
    if args.action == "publish":
        publish(args.variant)
    else:
        (ROOT / "artifacts").mkdir(exist_ok=True)
        with (ROOT / "artifacts/.gpu.lock").open("a") as gpu_lock:
            fcntl.flock(gpu_lock, fcntl.LOCK_EX)
            globals()[args.action](args.variant)