"""Compare the original tensor reader with lazy native MLX safetensors I/O."""

from __future__ import annotations

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

import mlx.core as mx
import numpy as np

from h3_mlx.runtime import configure_memory, environment
from h3_mlx.weights import WeightStore


def main(a):
    configure_memory(wired_gb=4)
    store = WeightStore(Path(a.model) / "transformer", reader="memmap")
    report = {
        "environment": environment(),
        "scope": "One original tensor at a time. OS page cache is not flushed; includes read/copy and Metal synchronization.",
        "trials": [],
    }
    for key in (
        "blocks.0.attn.qkv_proj.weight",
        "blocks.25.mlp.fc1.weight",
        "blocks.49.mlp.fc2.weight",
    ):
        native = mx.load(str(store.index[key]))
        for mode in ("memmap_copy", "mlx_lazy_load"):
            samples = []
            for _ in range(a.repeats):
                gc.collect()
                mx.clear_cache()
                mx.synchronize()
                tick = time.perf_counter()
                if mode == "memmap_copy":
                    value = store.get(key)
                else:
                    # Reload dictionary so every sample has an unevaluated Load node.
                    value = mx.load(str(store.index[key]))[key]
                    mx.eval(value)
                mx.synchronize()
                samples.append(time.perf_counter() - tick)
                bits = hashlib.sha256(
                    np.asarray(value.view(mx.uint16)).tobytes()
                ).hexdigest()
                del value
            report["trials"].append(
                {
                    "key": key,
                    "mode": mode,
                    "median_s": statistics.median(samples),
                    "samples_s": samples,
                    "uint16_sha256": bits,
                }
            )
            print(json.dumps(report["trials"][-1]), flush=True)
        del native
    dest = Path(a.output)
    dest.parent.mkdir(parents=True, exist_ok=True)
    dest.write_text(json.dumps(report, indent=2) + "\n")


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--model", required=True)
    p.add_argument("--repeats", type=int, default=3)
    p.add_argument("--output", default="reports/optimization-v2/loading.json")
    main(p.parse_args())