"""Apple GPU/CPU alternatives through PyTorch, including MPS Inductor fusion."""

from __future__ import annotations

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

import torch
import torch.nn.functional as F
from safetensors import safe_open


def load(folder, key):
    for p in sorted(Path(folder).glob("*.safetensors")):
        with safe_open(p, framework="pt", device="cpu") as f:
            if key in f.keys():
                return f.get_tensor(key)
    raise KeyError(key)


def measure(fn, device, repeats):
    def sync():
        if device == "mps":
            torch.mps.synchronize()

    for _ in range(2):
        fn()
        sync()
    samples = []
    for _ in range(repeats):
        sync()
        tick = time.perf_counter()
        fn()
        sync()
        samples.append(1000 * (time.perf_counter() - tick))
    return {"median_ms": statistics.median(samples), "samples_ms": samples}


def main(a):
    torch.set_num_threads(10)
    torch.manual_seed(731)
    root = Path(a.model) / "transformer"
    out = Path(a.output)
    out.parent.mkdir(parents=True, exist_ok=True)
    report = {
        "platform": platform.platform(),
        "torch": torch.__version__,
        "threads": 10,
        "repeats": a.repeats,
        "warmup": 2,
        "scope": "Actual H3 block-0 weights, synthetic activations; model loading and dtype conversion excluded.",
        "trials": [],
    }
    with torch.inference_mode():
        for name in ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2"):
            w = load(root, "blocks.0." + name + ".weight")
            x = torch.randn(a.tokens, w.shape[1]).bfloat16()
            # A single FP32 CPU reference distinguishes accumulation precision.
            ref = F.linear(x.float(), w.float())
            for device, dtype in (
                ("cpu", torch.bfloat16),
                ("cpu", torch.float32),
                ("mps", torch.bfloat16),
                ("mps", torch.float16),
                ("mps", torch.float32),
            ):
                row = {
                    "operator": name,
                    "device": device,
                    "dtype": str(dtype),
                    "shape": [a.tokens, w.shape[1], w.shape[0]],
                }
                try:
                    xx, ww = x.to(device, dtype), w.to(device, dtype)
                    row.update(
                        measure(
                            lambda xx=xx, ww=ww: F.linear(xx, ww), device, a.repeats
                        )
                    )
                    y = F.linear(xx, ww).cpu().float()
                    row["relative_rmse"] = torch.sqrt(
                        torch.mean((ref - y) ** 2) / torch.mean(ref**2)
                    ).item()
                    del xx, ww, y
                    if device == "mps":
                        torch.mps.empty_cache()
                except Exception as exc:
                    row["error"] = f"{type(exc).__name__}: {exc}"
                report["trials"].append(row)
                out.write_text(json.dumps(report, indent=2) + "\n")
                print(json.dumps(row), flush=True)
            del x, w, ref
            gc.collect()
        w1 = load(root, "blocks.0.mlp.fc1.weight").to("mps")
        w2 = load(root, "blocks.0.mlp.fc2.weight").to("mps")
        x = torch.randn(a.tokens, 5376, device="mps").bfloat16()

        def mlp(x):
            gate, up = F.linear(x, w1).chunk(2, dim=-1)
            return F.linear(F.silu(gate) * up, w2)

        reference = mlp(x)
        for compiled in (False, True):
            row = {
                "operator": "mlp",
                "compiled": compiled,
                "device": "mps",
                "dtype": "torch.bfloat16",
            }
            try:
                fn = torch.compile(mlp) if compiled else mlp
                row.update(measure(lambda fn=fn: fn(x), "mps", a.repeats))
                y = fn(x)
                row["equal_fraction"] = (reference == y).float().mean().item()
            except Exception as exc:
                row["error"] = f"{type(exc).__name__}: {exc}"
            report["trials"].append(row)
            out.write_text(json.dumps(report, indent=2) + "\n")
            print(json.dumps(row), flush=True)


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