"""Real first-block inputs: attention layout, Metal 4 FC2 tiling, lossless decode."""

from __future__ import annotations

import argparse
import gc
import itertools
import json
from functools import lru_cache
from pathlib import Path

import mlx.core as mx

from experiments.r3.common import paired
from experiments.v2_mpp import gemm
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.ops import rms_norm, rotate
from h3_mlx.packed import PackedBF16
from h3_mlx.runtime import configure_memory, environment


@lru_cache(None)
def decoder(lanes):
    return mx.fast.metal_kernel(
        name=f"h3_unpack_bf16_palette12_vec{lanes}",
        input_names=["low", "nibbles", "escape_indices", "escape_values", "base"],
        output_names=["out"],
        source=r"""
        uint first = thread_position_in_grid.x * LANES;
        #pragma unroll
        for (uint lane = 0; lane < LANES; ++lane) {
            uint i = first + lane;
            if (i >= low_shape[0]) break;
            uint code = (nibbles[i / 2] >> ((i % 2) * 4)) & 15;
            uint count = escape_indices_shape[0];
            if (code == 0 && count != 0) {
                uint left = 0, right = count;
                while (left < right) {
                    uint mid = (left + right) / 2;
                    if (escape_indices[mid] < i) left = mid + 1;
                    else right = mid;
                }
                out[i] = escape_values[left];
            } else {
                out[i] = ushort(low[i]) | ushort(((code & 7) + base[0]) << 8)
                       | ushort((code & 8) << 12);
            }
        }
        """,
    )


def unpack(packed, lanes, threads):
    raw = decoder(lanes)(
        inputs=packed.state(),
        template=[("LANES", lanes)],
        grid=((packed.low.size + lanes - 1) // lanes, 1, 1),
        threadgroup=(threads, 1, 1),
        output_shapes=[(packed.low.size,)],
        output_dtypes=[mx.uint16],
    )[0]
    return raw.view(mx.bfloat16).reshape(packed.shape)


class Capture:
    def __init__(self, linear):
        self.linear = linear

    def __call__(self, x):
        self.input = x
        mx.eval(x)
        return self.linear(x)


def main(a):
    configure_memory(wired_gb=14)
    dest = Path(a.output)
    dest.parent.mkdir(parents=True, exist_ok=True)
    report = {
        "environment": environment(),
        "scope": "Actual saved first-block activations, one GPU job, counterbalanced warmed calls including layout/decode costs. Operator evidence only.",
        "trials": [],
    }

    def record(row):
        report["trials"].append(row)
        dest.write_text(json.dumps(report, indent=2) + "\n")
        print(json.dumps(row), flush=True)

    with FootprintMonitor() as memory:
        data = mx.load(a.trace)
        model = DiT(a.model)
        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")
        )
        indices = 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")
        captures = {}
        for name in ("qkv", "out", "fc1", "fc2"):
            capture = Capture(getattr(block, name))
            captures[name] = capture
            setattr(block, name, capture)
        y = block(x, params=params, indices=indices, rope=rope, mlp_chunk=0)
        mx.eval(y)
        report["activation_ranges"] = {
            name: {
                "shape": list(v.input.shape),
                "max_abs": mx.max(mx.abs(v.input.astype(mx.float32))).item(),
            }
            for name, v in captures.items()
        }
        if "attention" in a.groups:
            qkv = (
                captures["qkv"]
                .linear(captures["qkv"].input)
                .reshape(1, x.shape[0], 3, 56, 128)
            )
            q, k = (
                rotate(
                    rms_norm(qkv[:, :, j], weight, block.config.qk_norm_eps), *rope
                ).transpose(0, 2, 1, 3)
                for j, weight in ((0, block.qn), (1, block.kn))
            )
            v = qkv[:, :, 2].transpose(0, 2, 1, 3)
            mx.eval(q, k, v)

            def attention(layout="strided"):
                args = [
                    mx.contiguous(t) if axis in layout else t
                    for axis, t in zip("qkv", (q, k, v), strict=True)
                ]
                return mx.fast.scaled_dot_product_attention(*args, scale=128**-0.5)

            ref = attention()
            mx.eval(ref)
            for layout in ("q", "k", "v", "kv", "qk", "qkv"):

                def fn(layout=layout):
                    return attention(layout)

                record(
                    {
                        "group": "attention",
                        "variant": layout,
                        **paired(attention, fn, a.repeats),
                        **error(ref, fn()),
                    }
                )
            del qkv, q, k, v, ref
        if "mpp" in a.groups:
            count = 0
            for name in ("fc2", "fc1"):
                cap = captures[name]
                layer, xx = cap.linear, cap.input
                w = layer.weight
                b = (
                    layer.bias
                    if layer.bias is not None
                    else mx.zeros(w.shape[0], w.dtype)
                )
                ref = layer(xx)
                mx.eval(ref)
                tiles = itertools.product(
                    (
                        (64, 128, 4),
                        (64, 128, 8),
                        (128, 128, 4),
                        (128, 128, 8),
                        (128, 256, 8),
                        (64, 256, 8),
                    ),
                    (128, 256, 512),
                    (0, 8),
                )
                for (bm, bn, sg), bk, walk in tiles:
                    if w.shape[1] % bk or w.shape[0] % bn:
                        continue
                    count += 1
                    row = {
                        "group": "mpp",
                        "operator": name,
                        "tile": [bm, bn, bk],
                        "simdgroups": sg,
                        "walk": walk,
                    }
                    try:

                        def fn(bm=bm, bn=bn, bk=bk, sg=sg, walk=walk, xx=xx, w=w, b=b):
                            return gemm(xx, w, b, bm, bn, bk, sg, walk)

                        row.update(
                            paired(lambda layer=layer, xx=xx: layer(xx), fn, a.repeats),
                            **error(ref, fn()),
                        )
                    except Exception as exc:
                        row["error"] = f"{type(exc).__name__}: {exc}"
                    record(row)
                    if count >= a.max_tiles:
                        break
                if count >= a.max_tiles:
                    break
        if "unpack" in a.groups:
            for name in ("qkv", "fc2"):
                cap = captures[name]
                layer, xx = cap.linear, cap.input
                packed = PackedBF16.pack(layer.weight)
                ref = packed.unpack()
                mx.eval(ref)
                for lanes, threads in itertools.product((2, 4, 8), (128, 256)):

                    def fn(lanes=lanes, threads=threads, packed=packed):
                        return unpack(packed, lanes, threads)

                    candidate = fn()
                    mx.eval(candidate)
                    bit_equal = mx.array_equal(
                        ref.view(mx.uint16), candidate.view(mx.uint16)
                    ).item()
                    row = {
                        "group": "unpack",
                        "operator": name,
                        "lanes": lanes,
                        "threads": threads,
                        "bit_equal": bit_equal,
                        "decode": paired(packed.unpack, fn, a.repeats),
                    }

                    def project(w, xx=xx, layer=layer):
                        return (
                            xx @ w.T
                            if layer.bias is None
                            else mx.addmm(layer.bias, xx, w.T)
                        )

                    row["with_gemm"] = paired(
                        lambda project=project, packed=packed: project(packed.unpack()),
                        lambda project=project, fn=fn: project(fn()),
                        a.repeats,
                    )
                    record(row)
                del packed, ref
        model.clear()
        gc.collect()
    report["memory"] = memory.result()
    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(
        "--trace", default="outputs/suite/product/native-3.trace.safetensors"
    )
    p.add_argument("--groups", nargs="+", default=["attention", "mpp", "unpack"])
    p.add_argument("--max-tiles", type=int, default=72)
    p.add_argument("--repeats", type=int, default=6)
    p.add_argument("--output", default="reports/optimization-r3/screen.json")
    main(p.parse_args())