"""Probe projection components selected by MTPLX's Qwen implementation.

This is an operator diagnostic, not a full-model or serving benchmark. Only the
selected projections are materialized; no weights are downloaded.
"""

import argparse
import hashlib
import json
import statistics
import time
from pathlib import Path

import mlx.core as mx
import numpy as np
from mtplx.models import qwen4_exp


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("model", type=Path)
    parser.add_argument("--capture", type=Path, help="Capture only the 32-row router GPU operation")
    parser.add_argument("--check-split-k", action="store_true",
                        help="Diagnose router rounding with explicitly batched K partitions")
    args = parser.parse_args()
    assert mx.__version__ == "0.32.2", "Use the pinned MTPLX MLX environment"
    index = json.loads((args.model / "model.safetensors.index.json").read_text())["weight_map"]
    config = json.loads((args.model / "config.json").read_text())
    quantization = config["quantization"]
    model_args = qwen4_exp.TextArgs.from_dict(config["text_config"])
    source = Path(qwen4_exp.__file__).resolve()
    print(json.dumps({"reference": "MTPLX Qwen projection components",
                      "source": str(source),
                      "source_sha256": hashlib.sha256(source.read_bytes()).hexdigest(),
                      "mlx_version": mx.__version__}), flush=True)
    values = []
    state = 0x12345678
    for _ in range(2560):
        state = (state * 1664525 + 1013904223) & 0xFFFFFFFF
        values.append(((state >> 24) - 128) / 32)
    for suffix in ("linear_attn.in_proj_qkv", "mlp.gate",
                   "linear_attn.in_proj_b", "mlp.shared_expert_gate"):
        name = "language_model.model.layers.0." + suffix
        weights = {}
        for part in ("weight", "scales", "biases"):
            key = name + "." + part
            weights[part] = mx.load(str(args.model / index[key]))[key]
        mx.eval(*weights.values())
        layout = quantization[name]
        module = (qwen4_exp.GatedDeltaNet(model_args) if suffix.startswith("linear_attn.")
                  else qwen4_exp.SparseMoeBlock(model_args))
        attribute = suffix.split(".")[-1]
        projection = getattr(module, attribute).to_quantized(**layout)
        projection.load_weights(list(weights.items()), strict=True)
        setattr(module, attribute, projection)
        for rows in (1, 2, 3, 4, 31, 32, 63, 64, 2047, 2048):
            x = mx.array([[values] * rows], dtype=mx.bfloat16)
            mx.eval(x)
            capture = args.capture and suffix == "mlp.gate" and rows == 32
            if capture:
                mx.metal.start_capture(str(args.capture))
            try:
                result = getattr(module, attribute)(x)
                mx.eval(result)
            finally:
                if capture:
                    mx.metal.stop_capture()
            raw = np.asarray(result.astype(mx.float32)).astype("<f4").tobytes()
            digest = hashlib.sha256(raw).hexdigest()
            if suffix == "linear_attn.in_proj_qkv" and rows == 1:
                assert digest == "8952fadfdb1ef4450fd23b0060d62ea9b118367733014985a94c2d815a2be203"
            print(json.dumps({"mlx_version": mx.__version__, "projection": name,
                              "quantization": layout, "rows": rows,
                              "first_row_sha256": hashlib.sha256(raw[:result.shape[-1] * 4]).hexdigest(),
                              "sha256": digest}), flush=True)
            if rows in (2, 3, 4, 32, 64, 2048):
                # Operator call + eval wall time, not a serving benchmark or
                # an extra GPU counter probe. Discard the first warmup sample.
                timings = []
                for trial in range(33):
                    started = time.perf_counter_ns()
                    measured = getattr(module, attribute)(x)
                    mx.eval(measured)
                    elapsed = (time.perf_counter_ns() - started) / 1000
                    if trial:
                        timings.append(elapsed)
                print(json.dumps({"diagnostic": "projection_call_eval_wall",
                                  "projection": name, "rows": rows, "samples": len(timings),
                                  "median_us": statistics.median(timings)}), flush=True)
            if args.check_split_k and suffix == "mlp.gate" and rows in (32, 64):
                # Counterfactual component operations, not a replacement oracle:
                # a real weight batch bypasses the non-batched split-K dispatch.
                for splits in (1, 2, 4, 5, 8, 10, 20, 40):
                    width = x.shape[-1] // splits
                    batches = max(2, splits)
                    partitioned = {}
                    for key, value in weights.items():
                        block = value.reshape(value.shape[0], splits, -1).transpose(1, 0, 2)
                        partitioned[key] = mx.repeat(block, batches, axis=0) if splits == 1 else block
                    inputs = x.reshape(rows, splits, width).transpose(1, 0, 2)
                    if splits == 1:
                        inputs = mx.repeat(inputs, batches, axis=0)
                    elif rows == 32:
                        # Small K raises MTPLX's QMV threshold to 33 rows. Pad
                        # to retain BF16 QMM dequantization in this control.
                        inputs = mx.concatenate([inputs, inputs], axis=1)
                    parts = mx.quantized_matmul(
                        inputs, partitioned["weight"], partitioned["scales"],
                        partitioned["biases"], transpose=True, **layout,
                    )
                    combined = parts[0] if splits == 1 else parts.sum(axis=0)
                    combined = combined[:rows]
                    mx.eval(combined)
                    diagnostic_hash = hashlib.sha256(
                        np.asarray(combined.astype(mx.float32)).astype("<f4").tobytes()
                    ).hexdigest()
                    first_actual = np.asarray(combined.astype(mx.float32))[0]
                    first_expected = np.asarray(result.astype(mx.float32))[0, 0]
                    differences = np.flatnonzero(first_actual != first_expected)
                    print(json.dumps({"diagnostic": "explicit_batched_partitions",
                                      "rows": rows, "splits": splits,
                                      "partition_rows": inputs.shape[1],
                                      "sha256": diagnostic_hash,
                                      "different_columns": int(len(differences)),
                                      "first_differences": [[int(i), float(first_actual[i]), float(first_expected[i])] for i in differences[:4]],
                                      "matches_projection": diagnostic_hash == digest}), flush=True)
                    if splits == (20 if rows == 32 else 10):
                        assert diagnostic_hash == digest, "Split-K QMM decomposition differs from MTPLX projection"
                    if splits == 1:
                        expected_unsplit = {
                            32: "3da6d25816a695eda420aea0a76ff22ee4c59689857d0ad02ed16112d66625b6",
                            64: "8e0702490c8b94ea1c167fdbbca27107d7ecf6ec0c0b4f446ea3f8927d34c0f4",
                        }
                        assert diagnostic_hash == expected_unsplit[rows], "Unsplit control changed"


if __name__ == "__main__":
    main()