"""Export complete native VAE tiles to Core ML and inspect real device placement.

The source BF16 weights/input values are preserved through FP16 conversion.
Internal FP16 rounding differs from the native BF16 decoder; compare images.
Run export/benchmark in .venv-coreml, then use v2_decode.py for image evidence.
"""

from __future__ import annotations

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

import coremltools as ct
import numpy as np
import torch
import torch.nn.functional as F
from safetensors import safe_open

from experiments.v2_coreml import devices


class TileDecoder(torch.nn.Module):
    def __init__(self, path, batch):
        super().__init__()
        self.batch = batch
        for p in sorted(Path(path).glob("*.safetensors")):
            with safe_open(p, framework="pt", device="cpu") as f:
                for key in f.keys():
                    if key.startswith("decoder.") or key.startswith("post_quant_conv."):
                        self.register_buffer(
                            key.replace(".", "__"), f.get_tensor(key).bfloat16().half()
                        )
        # Native BF16 grid arithmetic, including BF16 products in 2D RoPE.
        coords = 2 * ((torch.arange(16, dtype=torch.bfloat16) + 0.5) / 16) - 1
        yy, xx = torch.meshgrid(coords, coords, indexing="ij")
        pos = torch.stack((yy, xx), dim=-1).reshape(-1, 2)
        pos = torch.cat((pos, torch.zeros(5, 2, dtype=torch.bfloat16)))
        inv = 1 / (100 ** torch.arange(0, 1, 0.125, dtype=torch.float32))
        scaled = (2 * math.pi * pos.float()).bfloat16().float()
        angles = (scaled[:, :, None] * inv[None, None, :]).reshape(261, 16)
        angles = torch.cat((angles, angles), dim=-1)
        self.register_buffer("cosine", angles.cos().bfloat16().half()[None, :, None, :])
        self.register_buffer("sine", angles.sin().bfloat16().half()[None, :, None, :])

    def w(self, key):
        return getattr(self, key.replace(".", "__"))

    def linear(self, x, key):
        return F.linear(x, self.w(key + ".weight"), self.w(key + ".bias"))

    def norm(self, x, key=None):
        xf = x.float()
        y = xf * torch.rsqrt(torch.mean(xf * xf, dim=-1, keepdim=True) + 1e-5)
        if key is not None:
            y = y * self.w(key).float()
        return y.half()

    def rope(self, x):
        rotary, tail = x[..., :32], x[..., 32:]
        a, b = rotary.chunk(2, dim=-1)
        return torch.cat(
            (rotary * self.cosine + torch.cat((-b, a), dim=-1) * self.sine, tail),
            dim=-1,
        )

    def forward(self, raw):
        x = raw.permute(0, 2, 3, 1)
        x = F.linear(
            x,
            self.w("post_quant_conv.weight").reshape(24, 24),
            self.w("post_quant_conv.bias"),
        )
        x = self.linear(x, "decoder.x_embedder").reshape(self.batch, 256, 1024)
        x = torch.cat(
            (
                x,
                self.w("decoder.register_tokens").expand(self.batch, 4, 1024),
                torch.zeros(self.batch, 1, 1024, dtype=torch.float16),
            ),
            dim=1,
        )
        for i in range(24):
            p = f"decoder.transformer_blocks.{i}"
            h = self.norm(x, p + ".norm1.weight")
            qkv = self.linear(h, p + ".attn.to_qkv").reshape(self.batch, 261, 16, 3, 64)
            q, k, v = (qkv[:, :, :, j] for j in range(3))
            q, k = self.rope(self.norm(q)), self.rope(self.norm(k))
            attn = (
                F.scaled_dot_product_attention(
                    q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
                )
                .transpose(1, 2)
                .reshape(self.batch, 261, 1024)
            )
            x = x + self.linear(attn, p + ".attn.to_out") * self.w(p + ".scale1")
            h = self.norm(x, p + ".norm2.weight")
            gate, up = self.linear(h, p + ".ff.w1").chunk(2, dim=-1)
            x = x + self.linear(F.silu(gate) * up, p + ".ff.w2") * self.w(p + ".scale2")
        x = F.layer_norm(
            x.float(),
            (1024,),
            self.w("decoder.norm_out.weight").float(),
            self.w("decoder.norm_out.bias").float(),
            1e-5,
        ).half()
        x = self.linear(x, "decoder.proj_out")[:, :256]
        return (
            x.reshape(self.batch, 16, 16, 3, 16, 16)
            .permute(0, 3, 1, 4, 2, 5)
            .reshape(self.batch, 3, 256, 256)
        )


def main(a):
    torch.set_num_threads(a.threads)
    package = Path(a.package)
    example = (
        np.random.default_rng(731).normal(size=(a.batch, 24, 16, 16)).astype(np.float16)
    )
    if not package.exists():
        model = TileDecoder(a.vae, a.batch).eval()
        with torch.inference_mode():
            traced = torch.jit.trace(model, torch.from_numpy(example))
        converted = ct.convert(
            traced,
            convert_to="mlprogram",
            inputs=[ct.TensorType(name="raw", shape=example.shape, dtype=np.float16)],
            outputs=[ct.TensorType(name="pixels", dtype=np.float16)],
            minimum_deployment_target=ct.target.macOS15,
            compute_precision=ct.precision.FLOAT16,
            skip_model_load=True,
        )
        files = sorted(Path(a.vae).glob("*.safetensors"))
        identity = [
            (str(p.resolve()), p.stat().st_size, p.stat().st_mtime_ns) for p in files
        ]
        converted.user_defined_metadata.update(
            {
                "h3.decoder.format": "h3-coreml-tile-v1",
                "h3.decoder.source": hashlib.sha256(
                    json.dumps(identity).encode()
                ).hexdigest(),
                "h3.decoder.batch": str(a.batch),
                "h3.decoder.precision": "Source BF16 weights and inputs, FP16 internal arithmetic, BF16 tile assembly",
            }
        )
        package.parent.mkdir(parents=True, exist_ok=True)
        converted.save(str(package))
        del model, traced, converted
        gc.collect()
    report = {
        "coremltools": ct.__version__,
        "torch": torch.__version__,
        "package": str(package),
        "batch": a.batch,
        "scope": "Complete 24-layer tile decoder. FP16 internals, source weights rounded to BF16 before conversion. NumPy I/O included; tile assembly and normalization excluded.",
        "trials": [],
    }
    dest = Path(a.output)
    dest.parent.mkdir(parents=True, exist_ok=True)
    for unit in (getattr(ct.ComputeUnit, name) for name in a.units):
        row = {"compute_units": unit.name}
        try:
            tick = time.perf_counter()
            model = ct.models.MLModel(
                str(package),
                compute_units=unit,
                optimization_hints={
                    "specializationStrategy": ct.SpecializationStrategy.FastPrediction
                },
            )
            row["load_compile_s"] = time.perf_counter() - tick
            plan = ct.models.compute_plan.MLComputePlan.load_from_path(
                model.get_compiled_model_path(), compute_units=unit
            )
            row["placement"] = devices(plan)
            tick = time.perf_counter()
            y = model.predict({"raw": example})["pixels"]
            row["first_predict_ms"] = 1000 * (time.perf_counter() - tick)
            samples = []
            for _ in range(a.repeats):
                tick = time.perf_counter()
                y = model.predict({"raw": example})["pixels"]
                samples.append(1000 * (time.perf_counter() - tick))
            row.update(
                samples_ms=samples,
                median_ms=statistics.median(samples),
                finite=bool(np.isfinite(y).all()),
            )
            del model, plan, y
            gc.collect()
        except Exception as exc:
            row["error"] = f"{type(exc).__name__}: {exc}"
        report["trials"].append(row)
        dest.write_text(json.dumps(report, indent=2) + "\n")
        print(
            json.dumps({k: v for k, v in row.items() if k != "placement"}), flush=True
        )


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--vae", required=True)
    p.add_argument("--batch", type=int, default=8)
    p.add_argument("--threads", type=int, default=10)
    p.add_argument(
        "--units",
        nargs="+",
        choices=("CPU_ONLY", "CPU_AND_GPU", "CPU_AND_NE", "ALL"),
        default=["CPU_ONLY", "CPU_AND_GPU", "CPU_AND_NE", "ALL"],
    )
    p.add_argument("--repeats", type=int, default=3)
    p.add_argument("--package", default="cache/coreml-v2/vae-batch8.mlpackage")
    p.add_argument("--output", default="reports/optimization-v2/coreml-vae.json")
    main(p.parse_args())