#!/usr/bin/env python3
"""Convert the Laya checkpoint to a Core ML model and check it on tests/vectors.json.

    # The iOS file: fp16, one fixed length of 512 tokens, for the Neural Engine.
    uv run --no-project --python 3.12 \
      --with coremltools==9.0 --with torch==2.14.1 --with transformers==5.17.0 \
      --with laya==0.3.22 --with numpy==2.3.5 \
      python scripts/export_laya_coreml.py --precision fp16 --fixed-length 512

    # Check the existing package only (needs coremltools and numpy, not laya or torch):
    uv run --no-project --python 3.12 --with coremltools==9.0 --with numpy==2.3.5 \
      python scripts/export_laya_coreml.py --precision fp16 --fixed-length 512 --skip-export

    # The earlier fp32 package with a flexible length (the default).
    uv run --no-project --python 3.12 \
      --with coremltools==9.0 --with torch==2.14.1 --with transformers==5.17.0 \
      --with laya==0.3.22 --with numpy==2.5.3 python scripts/export_laya_coreml.py

The fixed-length conversion needs numpy 2.3: coremltools 9.0 calls int() on a
one-element array, which numpy 2.4 and later reject. With numpy 2.5.3 the
flexible conversion works. The numpy version changes the folded constants, so
pin it to reproduce a hash.

Both are MLPrograms with an iOS 16 deployment target. The graph is
`DecisionModel.forward` for one `choice` row (qtype 0, every marker valid)
without the act head.

Fixed length N (`--fixed-length N`, output `coreml/laya-<precision>-<N>.mlpackage`
and a zip of it):

    input_ids       int32 [1, N]   the Laya sequence, right-padded with the pad token id
    attention_mask  int32 [1, N]   1 for real tokens, 0 for padding
    marker_pos      int32 [1, 5]   option marker positions
    logits          float [1, 5]   raw option scores, before temperature

  The attention masks are additive float masks (-1e4, finite in fp16) built from
  `attention_mask` with plain ops, for the global layers, the sliding-window
  layers, and the two head layers. A pair (query i, key j) is masked when i or j
  is padding, so padded positions never attend and are never attended. The
  rotary tables and the window band are constants for length N. Three rewrites
  keep the program on the Neural Engine (found by OpenJevSwift pull request 85):
  the head layers are written with rank-4 tensors and scaled_dot_product_attention
  instead of nn.MultiheadAttention's rank-5 unpacking, the type embedding is a
  constant row, and the marker rows are picked with a one-hot matmul instead of
  a gather.

Flexible length (no `--fixed-length`, output `coreml/laya.<precision>.mlpackage`):

    input_ids   int32 [1, tokens]   the Laya sequence, 2 to 512 tokens, not padded
    marker_pos  int32 [1, 5]        option marker positions
    logits      float [1, 5]        raw option scores, before temperature

  Because the input is never padded, the key padding mask is all ones.

Before the conversion, the fixed-length wrapper runs in PyTorch fp32 on every
case of tests/vectors.json. After the conversion (or with `--skip-export`, on
the existing package) the script runs every case through coremltools on macOS,
once per compute unit. It prints the max probability error, the argmax changes,
the cases whose top probability changes side of 0.40, the Core ML compute plan
(which device each op is planned on), the latency over 30 runs, and the sha256
of the package. A conversion also writes a zip file of the package.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
import sys
import tempfile
import time
import zipfile
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "tests"))

from check import load_vectors  # noqa: E402

MODEL = "convaiinnovations/laya"
REVISION = "55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851"
OUT_DIR = ROOT / "coreml"
PAD_ID = 50283
OPTIONS = 5
MAX_TOKENS = 512
NEG = -1e4  # finite in fp16, so a fully masked row gives finite values instead of NaN
THRESHOLD = 0.40  # an app may add a second route below this top probability
UNITS = ("CPU_ONLY", "CPU_AND_GPU", "CPU_AND_NE")


def build_flexible(agent):
    import torch

    model = agent.model.eval().float()
    torch.backends.mha.set_fastpath_enabled(False)
    window = model.encoder.config.sliding_window

    class Router(torch.nn.Module):
        def __init__(self, m):
            super().__init__()
            self.m = m

        def forward(self, input_ids, marker_pos):
            m = self.m
            ids = input_ids.long()
            n = ids.shape[1]
            pos = torch.arange(n)
            near = ((pos[:, None] - pos[None, :]).abs() <= window).float()[None, None]
            full = torch.zeros(1, 1, 1, n)
            local = (near - 1.0) * 1e4
            h = m.encoder(
                input_ids=ids,
                attention_mask={"full_attention": full, "sliding_attention": local},
            ).last_hidden_state
            h = h + m.type_emb.weight[0]
            pad = torch.zeros(1, n, dtype=torch.bool)
            for layer in m.head.layers:
                h = layer(h, src_key_padding_mask=pad)
            picked = torch.index_select(h[0], 0, marker_pos[0])
            return m.scorer(picked).squeeze(-1)[None]

    return Router(model).eval()


def head_layer(layer, h, mask):
    """nn.TransformerEncoderLayer's norm-first eval forward with rank-4 tensors and an additive mask."""
    import torch.nn.functional as F

    attn = layer.self_attn
    b, s, d = h.shape
    heads = attn.num_heads
    q, k, v = F.linear(layer.norm1(h), attn.in_proj_weight, attn.in_proj_bias).split(d, dim=-1)
    q, k, v = (t.reshape(b, s, heads, d // heads).transpose(1, 2) for t in (q, k, v))
    a = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
    h = h + attn.out_proj(a.transpose(1, 2).reshape(b, s, d))
    return h + layer.linear2(layer.activation(layer.linear1(layer.norm2(h))))


def build_fixed(agent, length: int):
    import torch
    import torch.nn.functional as F

    model = agent.model.eval().float()
    torch.backends.mha.set_fastpath_enabled(False)
    for layer in model.head.layers:
        assert layer.norm_first and layer.activation is F.relu

    class Router(torch.nn.Module):
        def __init__(self, m):
            super().__init__()
            self.m = m
            encoder = m.encoder
            pos = torch.arange(length)
            far = (pos[None, :] - pos[:, None]).abs() > encoder.config.sliding_window
            self.register_buffer("band", (far.float() * NEG)[None, None], persistent=False)
            self.register_buffer("choice_row", m.type_emb.weight[0].detach().clone(), persistent=False)
            dummy = torch.zeros(1, length, encoder.config.hidden_size)
            with torch.no_grad():
                for kind in ("full_attention", "sliding_attention"):
                    cos, sin = encoder.rotary_emb(dummy, pos[None], kind)
                    self.register_buffer(f"cos_{kind}", cos, persistent=False)
                    self.register_buffer(f"sin_{kind}", sin, persistent=False)

        def forward(self, input_ids, attention_mask, marker_pos):
            m = self.m
            keep = attention_mask.to(torch.float32)
            pair = keep[:, :, None] * keep[:, None, :]  # 1 only when query and key are both real
            full = ((1.0 - pair) * NEG)[:, None]  # [1, 1, N, N]
            masks = {"full_attention": full, "sliding_attention": full + self.band}
            rotary = {kind: (getattr(self, f"cos_{kind}"), getattr(self, f"sin_{kind}")) for kind in masks}
            h = m.encoder.embeddings(input_ids=input_ids.long())
            for layer in m.encoder.layers:
                h = layer(h, attention_mask=masks[layer.attention_type],
                          position_embeddings=rotary[layer.attention_type])
            h = m.encoder.final_norm(h)
            h = h + self.choice_row
            for layer in m.head.layers:
                h = head_layer(layer, h, full)
            positions = torch.arange(length, dtype=marker_pos.dtype)
            onehot = (marker_pos[0][:, None] == positions[None, :]).to(h.dtype)  # [5, N]
            picked = onehot @ h[0]  # [5, d]
            return m.scorer(picked).squeeze(-1)[None].float()

    return Router(model).eval()


def check_wrapper_fixed(agent, wrapper, length: int, cases: list[dict], t: float) -> float:
    """The fixed wrapper in PyTorch fp32 against tests/vectors.json, before any conversion."""
    import torch

    worst = 0.0
    with torch.inference_mode():
        for case in cases:
            feed = pad_inputs(case["input_ids"], case["markers"], length, agent.tok.pad_token_id)
            logits = wrapper(*(torch.from_numpy(feed[k]) for k in ("input_ids", "attention_mask", "marker_pos")))
            worst = max(worst, float(np.abs(decode(logits.numpy(), t) - case["reference"]).max()))
    return worst


def convert(agent, out: Path, precision: str, length: int | None, fp32_ops: list[str]) -> dict:
    import coremltools as ct
    import torch

    if fp32_ops:
        compute_precision = ct.transform.FP16ComputePrecision(op_selector=lambda op: op.op_type not in fp32_ops)
    else:
        compute_precision = ct.precision.FLOAT32 if precision == "fp32" else ct.precision.FLOAT16
    common = dict(
        outputs=[ct.TensorType(name="logits", dtype=np.float32)],
        convert_to="mlprogram",
        compute_precision=compute_precision,
        minimum_deployment_target=ct.target.iOS16,
        skip_model_load=True,
    )
    if length is None:
        wrapper = build_flexible(agent)
        example = (
            torch.full((1, 131), 1000, dtype=torch.int32),
            torch.tensor([[14, 40, 64, 80, 99]], dtype=torch.int32),
        )
        tokens = torch.export.Dim("tokens", min=2, max=MAX_TOKENS)
        with torch.no_grad():
            program = torch.export.export(
                wrapper, example, dynamic_shapes={"input_ids": {1: tokens}, "marker_pos": None}
            ).run_decompositions()
        mlmodel = ct.convert(program, inputs=[
            ct.TensorType(name="input_ids", shape=ct.Shape((1, ct.RangeDim(2, MAX_TOKENS, default=MAX_TOKENS))),
                          dtype=np.int32),
            ct.TensorType(name="marker_pos", shape=(1, OPTIONS), dtype=np.int32),
        ], **common)
        about = "flexible length 2 to 512, not padded"
    else:
        wrapper = build_fixed(agent, length)
        example = (
            torch.full((1, length), agent.tok.pad_token_id, dtype=torch.int32),
            torch.zeros((1, length), dtype=torch.int32),
            torch.tensor([[14, 40, 64, 80, 99]], dtype=torch.int32),
        )
        example[1][0, :131] = 1
        with torch.no_grad():
            traced = torch.jit.trace(wrapper, example, check_trace=False)
        mlmodel = ct.convert(traced, inputs=[
            ct.TensorType(name="input_ids", shape=(1, length), dtype=np.int32),
            ct.TensorType(name="attention_mask", shape=(1, length), dtype=np.int32),
            ct.TensorType(name="marker_pos", shape=(1, OPTIONS), dtype=np.int32),
        ], **common)
        about = f"fixed length {length}, right-padded with token {agent.tok.pad_token_id}"
    mlmodel.short_description = (f"Laya router ({MODEL}@{REVISION}), {precision}, {about}: "
                                 "option logits before temperature")
    mlmodel.author = "Unofficial conversion of convaiinnovations/laya"
    mlmodel.license = "Apache-2.0"
    mlmodel.version = REVISION[:12]
    ops: dict[str, int] = {}
    for fn in mlmodel._mil_program.functions.values():
        for op in fn.operations:
            ops[op.op_type] = ops.get(op.op_type, 0) + 1
    out.parent.mkdir(parents=True, exist_ok=True)
    if out.exists():
        shutil.rmtree(out)
    mlmodel.save(str(out))
    return dict(sorted(ops.items(), key=lambda kv: -kv[1]))


def temperature() -> float:
    """The calibrated temperature of a choice question with 3 to 5 options, from rl_agent_config.json."""
    config = json.loads((ROOT / "rl_agent_config.json").read_text(encoding="utf-8"))
    return min(max(config["temperature_by_options"]["choice:3-5"], 0.5), 5.0)


def decode(logits, t: float) -> np.ndarray:
    z = np.asarray(logits, np.float64).reshape(-1)[:OPTIONS] / t
    p = np.exp(z - z.max())
    return p / p.sum()


def pad_inputs(ids: list[int], markers: list[int], length: int | None, pad_id: int) -> dict[str, np.ndarray]:
    if length is None:
        return {"input_ids": np.asarray([ids], np.int32), "marker_pos": np.asarray([markers], np.int32)}
    assert len(ids) <= length, f"{len(ids)} tokens do not fit in {length}"
    input_ids = np.full((1, length), pad_id, np.int32)
    input_ids[0, :len(ids)] = ids
    mask = np.zeros((1, length), np.int32)
    mask[0, :len(ids)] = 1
    return {"input_ids": input_ids, "attention_mask": mask, "marker_pos": np.asarray([markers], np.int32)}


def vector_cases() -> list[dict]:
    return load_vectors()[1]


def check(compiled: Path, units: str, t: float, length: int | None, pad_id: int,
          cases: list[dict]) -> dict:
    import coremltools as ct

    start = time.perf_counter()
    model = ct.models.CompiledMLModel(str(compiled), compute_units=getattr(ct.ComputeUnit, units))
    load_s = time.perf_counter() - start
    worst, flips, sides = 0.0, [], []
    for case in cases:
        p = decode(model.predict(pad_inputs(case["input_ids"], case["markers"], length, pad_id))["logits"], t)
        ref = case["reference"]
        worst = max(worst, float(np.abs(p - ref).max()))
        if int(p.argmax()) != int(ref.argmax()):
            flips.append(case["id"])
        if (p.max() < THRESHOLD) != (ref.max() < THRESHOLD):
            sides.append({"id": case["id"], "top": round(float(p.max()), 4), "reference": float(ref.max())})
    timing = latency(model, cases, length, pad_id)
    return {"units": units, "load_seconds": round(load_s, 2), "cases": len(cases),
            "max_abs_error": worst, "argmax_changes": flips, "threshold_side_changes": sides, **timing}


def latency(model, cases: list[dict], length: int | None, pad_id: int, runs: int = 30) -> dict:
    case = next(c for c in cases if c["id"] == "conv_lamp")
    feed = pad_inputs(case["input_ids"], case["markers"], length, pad_id)
    model.predict(feed)
    times = []
    for _ in range(runs):
        start = time.perf_counter()
        model.predict(feed)
        times.append((time.perf_counter() - start) * 1000)
    times.sort()
    return {"latency_input": f"{case['id']} ({len(case['input_ids'])} tokens)",
            "median_ms": round(times[len(times) // 2], 1), "p90_ms": round(times[int(runs * 0.9) - 1], 1)}


def compute_plan(compiled: Path, units: str) -> dict:
    """How many ops Core ML plans on each device for these compute units."""
    import coremltools as ct
    from coremltools.models.compute_plan import MLComputePlan

    plan = MLComputePlan.load_from_path(str(compiled), compute_units=getattr(ct.ComputeUnit, units))
    counts: dict[str, int] = {}
    on_cpu: dict[str, int] = {}
    for function in plan.model_structure.program.functions.values():
        for op in function.block.operations:
            usage = plan.get_compute_device_usage_for_mlprogram_operation(op)
            if usage is None:
                continue  # constants have no device
            device = type(usage.preferred_compute_device).__name__.replace("ML", "").replace("ComputeDevice", "")
            counts[device] = counts.get(device, 0) + 1
            if device == "CPU":
                on_cpu[op.operator_name] = on_cpu.get(op.operator_name, 0) + 1
    return {"units": units, "ops_by_device": counts, "cpu_ops": on_cpu}


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def sha256_tree(path: Path) -> str:
    digest = hashlib.sha256()
    for file in sorted(p for p in path.rglob("*") if p.is_file()):
        digest.update(str(file.relative_to(path)).encode())
        with file.open("rb") as f:
            for chunk in iter(lambda: f.read(1 << 20), b""):
                digest.update(chunk)
    return digest.hexdigest()


def write_zip(package: Path) -> Path:
    """One file for the app to download. Weights are already dense, so store without compression."""
    target = package.with_name(package.name + ".zip")
    with zipfile.ZipFile(target, "w", compression=zipfile.ZIP_STORED) as z:
        for file in sorted(p for p in package.rglob("*") if p.is_file()):
            z.write(file, file.relative_to(package.parent))
    return target


def main() -> None:
    import coremltools as ct

    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--precision", choices=("fp32", "fp16"), default="fp32")
    parser.add_argument("--fixed-length", type=int, default=None,
                        help="one fixed input length, padded, with an attention_mask input")
    parser.add_argument("--fp32-ops", default="", help="fp16 only: comma-separated MIL op types kept in fp32")
    parser.add_argument("--out", type=Path, default=None)
    parser.add_argument("--units", default=",".join(UNITS), help="compute units to check")
    parser.add_argument("--skip-export", action="store_true", help="check an existing package only")
    args = parser.parse_args()
    if args.fixed_length is not None and not 2 <= args.fixed_length <= MAX_TOKENS:
        parser.error(f"--fixed-length must be 2 to {MAX_TOKENS}")
    fp32_ops = [op for op in args.fp32_ops.split(",") if op]
    if fp32_ops and args.precision != "fp16":
        parser.error("--fp32-ops needs --precision fp16")
    out = args.out or OUT_DIR / (f"laya-{args.precision}-{args.fixed_length}.mlpackage" if args.fixed_length
                                 else f"laya.{args.precision}.mlpackage")
    units = [u for u in args.units.split(",") if u]
    if args.fixed_length is None:
        units = [u for u in units if u != "CPU_AND_NE"]  # a flexible length does not reach the ANE

    pad_id = PAD_ID
    t = temperature()
    cases = vector_cases()
    report: dict = {"model": MODEL, "revision": REVISION, "package": str(out), "precision": args.precision,
                    "fixed_length": args.fixed_length, "fp32_ops": fp32_ops, "pad_token_id": pad_id,
                    "temperature": t, "coremltools": ct.__version__}
    if not args.skip_export:
        import laya

        agent = laya.load(MODEL, device="cpu", revision=REVISION)
        assert int(agent.tok.pad_token_id) == PAD_ID
        if args.fixed_length is not None:
            report["wrapper_fp32_max_abs_error"] = check_wrapper_fixed(
                agent, build_fixed(agent, args.fixed_length), args.fixed_length, cases, t)
            print(f"PyTorch fixed wrapper vs vectors: {report['wrapper_fp32_max_abs_error']:.2e}", flush=True)
        start = time.perf_counter()
        report["ops"] = convert(agent, out, args.precision, args.fixed_length, fp32_ops)
        print(f"converted {out} in {time.perf_counter() - start:.0f} s", flush=True)
        del agent

    scratch = Path(tempfile.mkdtemp(prefix="laya-coreml-"))
    try:
        compiled = Path(ct.utils.compile_model(str(out), str(scratch / "laya.mlmodelc")))
        report["plans"] = [compute_plan(compiled, u) for u in units]
        report["checks"] = []
        for u in units:
            result = check(compiled, u, t, args.fixed_length, pad_id, cases)
            report["checks"].append(result)
            print(json.dumps(result), flush=True)
    finally:
        shutil.rmtree(scratch, ignore_errors=True)
    report["sha256_tree"] = sha256_tree(out)
    report["weights_sha256"] = sha256_tree(out / "Data/com.apple.CoreML/weights")
    report["package_bytes"] = sum(p.stat().st_size for p in out.rglob("*") if p.is_file())
    archive = out.with_name(out.name + ".zip")
    if args.fixed_length is not None and not args.skip_export:
        archive = write_zip(out)
    if archive.exists():
        report["zip"] = {"path": str(archive), "bytes": archive.stat().st_size, "sha256": sha256_file(archive)}
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()