"""Screen a 512 MiB free-allocator cache on the actual fully tuned first block."""

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

import mlx.core as mx

from h3_mlx.benchmark import error
from h3_mlx.config import SIGMAS
from h3_mlx.dit import DiT
from h3_mlx.layout import image_positions, rope_cache
from h3_mlx.memory import FootprintMonitor
from h3_mlx.runtime import configure_memory, environment


def main(a):
    configure_memory(wired_gb=14)
    data = mx.load("outputs/suite/product/native-3.trace.safetensors")
    model = DiT(a.model, kernel="metal-exact", tune_matmul=True)
    model.prepare_modulation(SIGMAS)
    text = model.refine(data["conditioning"])
    x = mx.concatenate(
        (
            text,
            model._linear("video_patch_proj")(data["initial_rows"]).astype(mx.bfloat16),
        ),
        axis=0,
    )
    rope = rope_cache(
        image_positions(len(text), 128, 128), model._constant("rope.inv_freq")
    )
    ids = mx.concatenate((data["text_tags"], mx.zeros(4096, mx.int32)))
    params = tuple(model.modulation["block.0"][0, :, j] for j in range(6))
    block = model._block("blocks.0")
    fn = mx.compile(
        lambda x, p, i, r: block(
            x, params=p, indices=i, rope=r, kernel="metal-exact", mlp_chunk=0
        )
    )
    ref = fn(x, params, ids, rope)
    mx.eval(ref)
    report = {
        "environment": environment(),
        "scope": "Real complete first block, not full-image performance. Cache is free allocations only; same compiled graph and original weights.",
        "runs": [],
    }
    for repeat in range(6):
        for cap in (1.0, 0.5) if repeat % 2 == 0 else (0.5, 1.0):
            mx.set_cache_limit(int(cap * 1024**3))
            for _ in range(2):
                mx.eval(fn(x, params, ids, rope))
            mx.reset_peak_memory()
            with FootprintMonitor(interval=0.001) as memory:
                mx.synchronize()
                tick = time.perf_counter()
                result = fn(x, params, ids, rope)
                mx.eval(result)
                mx.synchronize()
                elapsed = time.perf_counter() - tick
            row = {
                "cache_gb": cap,
                "repeat": repeat,
                "milliseconds": elapsed * 1000,
                "physical": memory.result(),
                "mlx_peak_bytes": mx.get_peak_memory(),
                "mlx_cache_bytes": mx.get_cache_memory(),
                **error(ref, result),
            }
            report["runs"].append(row)
            print(json.dumps(row), flush=True)
    report["summary"] = {
        str(cap): {
            "median_ms": statistics.median(
                r["milliseconds"] for r in report["runs"] if r["cache_gb"] == cap
            ),
            "max_physical_bytes": max(
                r["physical"]["sampled_peak_physical_bytes"]
                for r in report["runs"]
                if r["cache_gb"] == cap
            ),
        }
        for cap in (1.0, 0.5)
    }
    Path(a.output).write_text(json.dumps(report, indent=2) + "\n")
    model.clear()


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--model", required=True)
    p.add_argument("--output", default="reports/optimization-r3/cache-screen.json")
    main(p.parse_args())