"""Real MTPLX module construction/strict lazy loading, no model inference.

Only MTP's two small pre-FC norms are evaluated. Run under test-supervisor.
"""
import sys
print("binding reference: importing pinned runtime", file=sys.stderr, flush=True)
import hashlib
import json
import os
from contextlib import redirect_stdout
from pathlib import Path
import mlx.core as mx
import numpy as np
from mlx.utils import tree_flatten
from mlx_lm.utils import load_model
from mtplx.models import qwen4_exp as qwen


def main(root, mode=None):
    assert mx.__version__ == "0.32.2"
    name_in_bank = lambda name: "mtp." + name[len("language_model.mtp."):] if name.startswith("language_model.mtp.") else name
    for mask, mtp in ((0, False), (7, False), (0, True), (7, True)):
        print(f"binding reference: mask={mask} mtp={mtp}", file=sys.stderr, flush=True)
        os.environ.update(MTPLX_FUSED_GATE_UP=str(mask & 1), MTPLX_FUSED_GDN_INPROJ=str((mask >> 1) & 1), MTPLX_FUSED_QSA_QKV=str((mask >> 2) & 1))
        with redirect_stdout(sys.stderr):
            model, _ = load_model(Path(root), lazy=True, strict=True, get_model_classes=lambda **_: (qwen.Model, qwen.ModelArgs))
            if mtp:
                assert model.attach_mtp(Path(root))
        if mode == "--parameter-order":
            names = [name_in_bank(name) for name, _ in tree_flatten(model.parameters())]
            print(json.dumps(dict(mask=mask, mtp=mtp, names=names),separators=(",", ":")),flush=True)
            del model, names
            continue
        if mode == "--post-weight-load":
            from mtplx.memory_plan import ngram_table_resident_policy
            assert not ngram_table_resident_policy(), "this test must not make the full table resident"
            with redirect_stdout(sys.stderr):
                model.post_weight_load(Path(root))
            tables = []
            for i, layer in enumerate(model.layers):
                if getattr(layer, "ple", None) is None:
                    continue
                table = layer.ple.ple_embedding.ngram_embedding
                sidecar = table._sidecar
                out = table(mx.array([0, 1, 255, 1], dtype=mx.int64))
                mx.eval(out)
                tables.append(dict(layer=i, resident=getattr(table, "_lazy_parts", None) is not None,
                    hot_mb=sidecar._hot_cap_rows * sidecar._hot_row_bytes // 2**20,
                    prefetch=sidecar._pool is not None, shape=list(out.shape),
                    sha256=hashlib.sha256(np.asarray(out.astype(mx.float32)).astype('<f4').tobytes()).hexdigest()))
            print(json.dumps(dict(mask=mask, mtp=mtp, tables=tables, ar_ready=model.set_ar_pipeline_mode(True)),separators=(",", ":")),flush=True)
            del model
            continue
        parameters = tree_flatten(model.parameters())
        rows = sorted([[name_in_bank(name), list(a.shape), str(a.dtype).removeprefix("mlx.core.")] for name, a in parameters])
        quantized = {}
        for name, module in model.named_modules():
            name = name_in_bank(name)
            if hasattr(module, "bits") and hasattr(module, "weight") and getattr(module, "scales", None) is not None:
                quantized[name] = [module.bits, module.group_size, list(module.weight.shape)]
            elif hasattr(module, "bits") and hasattr(module, "gu_weight"):
                quantized[name + ".gu"] = [module.bits, module.group_size, list(module.gu_weight.shape)]
        norms = {}
        if mtp:
            for name in ("pre_fc_norm_embedding", "pre_fc_norm_hidden"):
                a = getattr(model.mtp, name).weight
                mx.eval(a)
                norms["mtp." + name + ".weight"] = hashlib.sha256(np.asarray(a.astype(mx.float32)).astype('<f4').tobytes()).hexdigest()
        print(json.dumps(dict(mask=mask, mtp=mtp, parameter_count=len(rows),
            metadata_sha256=hashlib.sha256(json.dumps(rows,separators=(",", ":")).encode()).hexdigest(),
            quantized=quantized, norms=norms,
            ple=[i for i, layer in enumerate(model.layers) if getattr(layer,"ple",None) is not None],
            linear=[i for i, layer in enumerate(model.layers) if layer.is_linear]),separators=(",", ":")),flush=True)
        del model, parameters, quantized, rows


if __name__ == "__main__":
    assert len(sys.argv) in (2, 3)
    if len(sys.argv) == 3:
        assert sys.argv[2] in ("--post-weight-load", "--parameter-order")
    main(sys.argv[1], sys.argv[2] if len(sys.argv) == 3 else None)