# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "torchvision>=0.27.1",
#   "torch>=2.4",
#   "transformers>=4.50,<5",
#   "onnx>=1.16",
#   "onnxruntime>=1.18",
#   "pillow",
#   "numpy",
#   "huggingface-hub",
#   "onnxscript",
# ]
# ///
"""MoonVit.py — standalone build / eval / inference for moonshotai/MoonViT-SO-400M → ONNX.

MoonViT (the Kimi-VL vision tower) is a SigLIP-SO400M-style ViT (1152h / 27L / patch 14) with
NaViT-style packing: forward(pixel_values [L,3,14,14], grid_hws [N,2]) → merged tokens.
Its internals call `grid_hws.tolist()` (per-image Python loops for the interpolated pos-emb,
the 2D-RoPE freqs and the 2x2 patch merger), so the graph is DATA-DEPENDENT on the grid.

Export strategy: bake `grid_hws` as a constant for ONE image at a FIXED grid (default 28x28
patches = 392x392 px). The tracer unrolls the loops against the constant, producing a clean
fixed-shape graph. Need other resolutions → export more variants with --grid-h/--grid-w.

Usage (uv run resolves the inline deps):
  uv run MoonVit.py build     --output MoonVitOnnx/ [--grid-h 28 --grid-w 28]
  uv run MoonVit.py eval      --onnx-model-path MoonVitOnnx/
  uv run MoonVit.py inference --onnx-model-path MoonVitOnnx/ --image cat.png

build     downloads the model (trust_remote_code), exports moonvit.onnx + manifest.json.
eval      original PyTorch vs ONNX on identical random packed inputs → cosine / max|Δ|.
inference preprocesses an image (resize to the export grid, mean/std 0.5) and runs the ONNX.
"""
import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np

MODEL_ID = "moonshotai/MoonViT-SO-400M"
PATCH = 14
MEAN = 0.5
STD = 0.5

for _s in (sys.stdout, sys.stderr):
    try:
        _s.reconfigure(encoding="utf-8", errors="replace")
    except Exception:
        pass


# --------------------------------------------------------------------------- shared
def load_pytorch_model():
    """MoonViT in fp32 + sdpa (bf16 checkpoint → fp32 for export/parity; sdpa traces)."""
    import torch
    # Compat shim: the repo's remote code imports PytorchGELUTanh from
    # transformers.activations, which newer transformers removed. It is exactly
    # nn.GELU(approximate="tanh") — restore it before the remote module loads.
    import transformers.activations as _act
    if not hasattr(_act, "PytorchGELUTanh"):
        class PytorchGELUTanh(torch.nn.GELU):
            def __init__(self):
                super().__init__(approximate="tanh")
        _act.PytorchGELUTanh = PytorchGELUTanh
    from transformers import AutoModel
    model = AutoModel.from_pretrained(
        MODEL_ID, trust_remote_code=True, dtype=torch.float32,
        attn_implementation="sdpa",
    )
    return model.eval()


def make_inputs(grid_h, grid_w, seed=0):
    """Random packed patches for one image: pixel_values [L,3,14,14], grid_hws [[h,w]]."""
    import torch
    g = torch.Generator().manual_seed(seed)
    pixel_values = torch.randn(grid_h * grid_w, 3, PATCH, PATCH, generator=g,
                               dtype=torch.float32)
    grid_hws = torch.tensor([[grid_h, grid_w]], dtype=torch.int64)
    return pixel_values, grid_hws


def cosine(a, b):
    a = np.asarray(a, dtype=np.float64).ravel()
    b = np.asarray(b, dtype=np.float64).ravel()
    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-12))


def read_manifest(onnx_dir):
    p = Path(onnx_dir) / "manifest.json"
    if not p.exists():
        sys.exit(f"no manifest.json in {onnx_dir} — run `MoonVit.py build` first")
    return json.loads(p.read_text())


def _fix_mixed_index_dtypes(onnx_path):
    """Legacy-exporter quirk: a Slice/Gather can end up with MIXED int32/int64 index inputs
    (ORT load error: 'Type parameter (Tind) bound to different types'). Normalize every
    int32 index initializer/Constant feeding Slice starts/ends/axes/steps to int64."""
    import onnx
    from onnx import numpy_helper, TensorProto

    m = onnx.load(str(onnx_path))
    inits = {i.name: i for i in m.graph.initializer}
    consts = {n.output[0]: n for n in m.graph.node if n.op_type == "Constant"}
    index_inputs = set()
    for n in m.graph.node:
        if n.op_type == "Slice":
            index_inputs.update(n.input[1:])
        elif n.op_type in ("Gather", "GatherElements"):
            index_inputs.add(n.input[1])
    fixed = 0
    for name in index_inputs:
        if name in inits and inits[name].data_type == TensorProto.INT32:
            arr = numpy_helper.to_array(inits[name]).astype(np.int64)
            inits[name].CopyFrom(numpy_helper.from_array(arr, name))
            fixed += 1
        elif name in consts:
            for a in consts[name].attribute:
                if a.name == "value" and a.t.data_type == TensorProto.INT32:
                    arr = numpy_helper.to_array(a.t).astype(np.int64)
                    a.t.CopyFrom(numpy_helper.from_array(arr))
                    fixed += 1

    # COMPUTED index inputs (e.g. int32 cu_seqlens → Unsqueeze → Slice starts/ends): insert an
    # explicit Cast→int64 in front of the consumer. Cast(int64→int64) is a no-op, so this is
    # safe to apply to every non-constant index input.
    cast_cache, new_nodes, casts = {}, [], 0
    for n in list(m.graph.node):
        idx_slots = (range(1, len(n.input)) if n.op_type == "Slice"
                     else [1] if n.op_type in ("Gather", "GatherElements") else [])
        for slot in idx_slots:
            name = n.input[slot]
            if not name or name in inits or name in consts:
                continue
            if name not in cast_cache:
                out_name = f"{name}_as_i64"
                new_nodes.append(onnx.helper.make_node(
                    "Cast", [name], [out_name], name=f"Cast_i64_{casts}",
                    to=TensorProto.INT64))
                cast_cache[name] = out_name
                casts += 1
            n.input[slot] = cast_cache[name]
    if new_nodes:
        # prepend is invalid (defs must precede uses per producer order); insert each cast
        # right before its first consumer by rebuilding the node list topologically.
        consumers = {c.input[0]: c for c in new_nodes}       # original name -> cast node
        rebuilt = []
        emitted = set()
        for n in m.graph.node:
            rebuilt.append(n)
            for o in n.output:
                if o in consumers and o not in emitted:
                    rebuilt.append(consumers[o])
                    emitted.add(o)
        # casts whose source is a graph input
        for src, cnode in consumers.items():
            if src not in emitted and any(gi.name == src for gi in m.graph.input):
                rebuilt.insert(0, cnode)
                emitted.add(src)
        del m.graph.node[:]
        m.graph.node.extend(rebuilt)
    if fixed or casts:
        onnx.save(m, str(onnx_path))
        print(f"  index-dtype fixup: {fixed} const rewrites, {casts} Cast→int64 inserted")


def _patch_for_export(model, grid_hws, dynamo):
    """Neutralize MoonViT's data-dependent / complex-dtype ops for ONNX export, using the baked
    single-image grid. Every patch is mathematically identical for that grid (verified by the
    parity check in cmd_build). The rope patch is needed for BOTH exporters (ONNX has no complex
    dtype); the sdpa/merger/posemb patches are only needed by the dynamo exporter (the legacy
    tracer unrolls those loops against the baked grid on its own), but they're harmless for legacy.
    """
    import sys as _sys
    import torch
    import torch.nn.functional as F

    mm = _sys.modules[type(model).__module__]
    gh, gw = int(grid_hws[0, 0]), int(grid_hws[0, 1])

    # (1) 2D RoPE: complex64 freqs_cis → precomputed real cos/sin constants (ONNX has no complex).
    with torch.no_grad():
        freqs = model.encoder.rope_2d.get_freqs_cis(grid_hws)      # complex64 [L, hd/2]
    cos_c = freqs.real.contiguous().unsqueeze(-2)                  # [L, 1, hd/2]
    sin_c = freqs.imag.contiguous().unsqueeze(-2)

    def apply_rope_real(xq, xk, freqs_cis):
        def rot(x):
            xr = x.float().reshape(*x.shape[:-1], -1, 2)
            a, b = xr[..., 0], xr[..., 1]
            out = torch.stack([a * cos_c - b * sin_c, a * sin_c + b * cos_c], dim=-1)
            return out.reshape(*x.shape).type_as(x)
        return rot(xq), rot(xk)

    mm.apply_rope = apply_rope_real
    mm.Rope2DPosEmb.get_freqs_cis = lambda self, grid_hws=None, **k: torch.zeros(1)

    if not dynamo:
        return                                # legacy tracer handles the rest by unrolling

    # (2) sdpa: the in-place block-diagonal mask loop breaks torch.export. One baked image ⇒
    #     the mask is all-True ⇒ full attention (attn_mask=None), reshaped to 4D for onnxscript.
    def sdpa_full(q, k, v, q_cu_seqlens=None, k_cu_seqlens=None):
        seq = q.shape[0]
        q, k, v = (t.transpose(0, 1).unsqueeze(0) for t in (q, k, v))   # [1,H,L,hd]
        o = F.scaled_dot_product_attention(q, k, v, None, dropout_p=0.0)
        return o.squeeze(0).transpose(0, 1).reshape(seq, -1)
    mm.VL_VISION_ATTENTION_FUNCTIONS["sdpa"] = sdpa_full

    # (3) patch_merger + pos-emb interp iterate grid_hws.tolist() → unbacked-symint slices the
    #     onnxscript converter rejects. Bake concrete Python ints for the single image.
    def patch_merger_const(x, grid_hws, merge_kernel_size=model.merge_kernel_size):
        d = x.size(-1); kh, kw = merge_kernel_size; nh, nw = gh // kh, gw // kw
        rs = x.view(nh, kh, nw, kw, d).permute(0, 2, 1, 3, 4).contiguous().view(nh * nw, kh * kw, -1)
        return [rs]
    mm.patch_merger = patch_merger_const

    def posemb_const(self, x, grid_hws):
        if (gh, gw) == tuple(self.weight.shape[:-1]):
            pe = self.weight.flatten(end_dim=1)
        else:
            pe = (F.interpolate(self.weight.permute(2, 0, 1).unsqueeze(0), size=(gh, gw),
                                mode=self.interpolation_mode)
                  .squeeze(0).permute(1, 2, 0).flatten(end_dim=1))
        return x + pe
    mm.Learnable2DInterpPosEmb.forward = posemb_const


# --------------------------------------------------------------------------- build
def cmd_build(args):
    import torch

    out = Path(args.output)
    out.mkdir(parents=True, exist_ok=True)
    gh, gw = args.grid_h, args.grid_w
    print(f"=== build {MODEL_ID} → {out}  grid={gh}x{gw} ({gh*PATCH}x{gw*PATCH} px) ===")

    model = load_pytorch_model()

    class Wrapper(torch.nn.Module):
        """grid_hws baked as a constant so the data-dependent loops unroll at trace time."""
        def __init__(self, m, grid_hws):
            super().__init__()
            self.m = m
            self.register_buffer("grid_hws", grid_hws)

        def forward(self, pixel_values):
            out = self.m(pixel_values, self.grid_hws)
            # patch_merger returns a LIST of per-image tensors; concat (single image → its tokens)
            return torch.cat(out, dim=0) if isinstance(out, (list, tuple)) else out

    pixel_values, grid_hws = make_inputs(gh, gw)
    wrapper = Wrapper(model, grid_hws).eval()

    with torch.no_grad():
        ref = wrapper(pixel_values)          # golden reference computed BEFORE any patch
    print(f"  pytorch output: {tuple(ref.shape)}  (merged tokens, dim)")

    _patch_for_export(model, grid_hws, dynamo=args.dynamo)

    with torch.no_grad():                    # patched forward must still match the reference
        patched = wrapper(pixel_values)
    pc = cosine(patched.numpy(), ref.numpy())
    print(f"  export patches parity: cos={pc:.6f}  (dynamo={args.dynamo})")
    if pc < 0.9999:
        sys.exit(f"EXPORT PATCH MISMATCH (cos={pc:.6f}) — aborting")

    onnx_path = out / "moonvit.onnx"
    kw = dict(input_names=["pixel_values"], output_names=["image_features"], opset_version=19)
    with torch.no_grad():
        if args.dynamo:
            # NEW exporter: cleaner/smaller graph; needs the sdpa+merger+posemb patches above.
            torch.onnx.export(wrapper, (pixel_values,), str(onnx_path), dynamo=True, **kw)
        else:
            # LEGACY tracer: unrolls the .tolist() loops against the baked grid. `dynamo=False`
            # forces it on newer torch (which defaults to dynamo).
            try:
                torch.onnx.export(wrapper, (pixel_values,), str(onnx_path), dynamo=False,
                                  do_constant_folding=True, **kw)
            except TypeError:                  # older torch: no `dynamo` kwarg
                torch.onnx.export(wrapper, (pixel_values,), str(onnx_path),
                                  do_constant_folding=True, **kw)
    size_mb = onnx_path.stat().st_size / 1e6
    data = out / "moonvit.onnx.data"
    if data.exists():
        size_mb += data.stat().st_size / 1e6
    print(f"  exported {onnx_path.name}  ({size_mb:.0f} MB, {'dynamo' if args.dynamo else 'legacy'})")

    if not args.dynamo:
        _fix_mixed_index_dtypes(onnx_path)   # legacy-only: int32/int64 Slice index mismatch

    # immediate sanity: ORT vs the PyTorch reference on the export input
    import onnxruntime as ort
    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
    got = sess.run(None, {"pixel_values": pixel_values.numpy()})[0]
    c = cosine(got, ref.numpy())
    maxd = float(np.abs(got - ref.numpy()).max())
    print(f"  build sanity: ort vs pytorch  cos={c:.6f}  max|Δ|={maxd:.3e}")

    (out / "manifest.json").write_text(json.dumps({
        "model_id": MODEL_ID, "file": "moonvit.onnx",
        "grid_h": gh, "grid_w": gw, "patch_size": PATCH,
        "input": {"pixel_values": [gh * gw, 3, PATCH, PATCH]},
        "output": {"image_features": list(ref.shape)},
        "image_size_px": [gh * PATCH, gw * PATCH],
        "normalize": {"mean": MEAN, "std": STD},
        "exporter": "dynamo" if args.dynamo else "legacy",
        "note": "fixed-grid export: grid_hws baked as a constant (internals are data-dependent)",
    }, indent=2))
    print("  wrote manifest.json")
    if c < 0.999:
        sys.exit(f"BUILD SANITY FAILED (cos={c:.4f} < 0.999)")
    print("build OK")


# --------------------------------------------------------------------------- eval
def cmd_eval(args):
    import torch
    import onnxruntime as ort

    man = read_manifest(args.onnx_model_path)
    gh, gw = man["grid_h"], man["grid_w"]
    onnx_path = Path(args.onnx_model_path) / man["file"]
    print(f"=== eval original vs onnx  grid={gh}x{gw}  {onnx_path} ===")

    model = load_pytorch_model()
    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])

    worst = 1.0
    for seed in range(args.samples):
        pixel_values, grid_hws = make_inputs(gh, gw, seed=seed)
        with torch.no_grad():
            ref = model(pixel_values, grid_hws)
            ref = (torch.cat(ref, 0) if isinstance(ref, (list, tuple)) else ref).numpy()
        got = sess.run(None, {"pixel_values": pixel_values.numpy()})[0]
        c = cosine(got, ref)
        maxd = float(np.abs(got - ref).max())
        worst = min(worst, c)
        print(f"  sample {seed}: cos={c:.6f}  max|Δ|={maxd:.3e}  out={got.shape}")

    verdict = "PASS" if worst >= args.tol else "FAIL"
    print(f"=== {verdict}  (worst cosine {worst:.6f}, tol {args.tol}) ===")
    sys.exit(0 if verdict == "PASS" else 1)


# --------------------------------------------------------------------------- inference
def preprocess(image_path, gh, gw):
    """PIL image → packed patches [L,3,14,14], raster order (matches conv patchify)."""
    from PIL import Image
    img = Image.open(image_path).convert("RGB").resize((gw * PATCH, gh * PATCH), Image.BICUBIC)
    x = np.asarray(img, dtype=np.float32) / 255.0            # [H,W,3]
    x = (x - MEAN) / STD
    x = x.transpose(2, 0, 1)                                  # [3,H,W]
    # [3, gh, 14, gw, 14] → [gh, gw, 3, 14, 14] → [L,3,14,14]
    x = x.reshape(3, gh, PATCH, gw, PATCH).transpose(1, 3, 0, 2, 4)
    return np.ascontiguousarray(x.reshape(gh * gw, 3, PATCH, PATCH))


def cmd_inference(args):
    import onnxruntime as ort

    man = read_manifest(args.onnx_model_path)
    gh, gw = man["grid_h"], man["grid_w"]
    onnx_path = Path(args.onnx_model_path) / man["file"]
    print(f"=== inference  {onnx_path}  grid={gh}x{gw} ===")

    if args.image:
        pixel_values = preprocess(args.image, gh, gw)
        src = args.image
    else:
        pixel_values = make_inputs(gh, gw)[0].numpy()
        src = "random tensor (no --image given)"
    sess = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])

    t0 = time.time()
    feats = sess.run(None, {"pixel_values": pixel_values.astype(np.float32)})[0]
    dt = time.time() - t0

    print(f"  input : {src} → pixel_values {pixel_values.shape}")
    print(f"  output: image_features {feats.shape}  dtype={feats.dtype}  ({dt*1000:.0f} ms)")
    print(f"  stats : mean={feats.mean():+.4f}  std={feats.std():.4f}  "
          f"finite={bool(np.isfinite(feats).all())}")
    if args.save:
        np.save(args.save, feats)
        print(f"  saved features → {args.save}")


# --------------------------------------------------------------------------- cli
def main():
    ap = argparse.ArgumentParser(description="MoonViT-SO-400M → ONNX (build / eval / inference)")
    sub = ap.add_subparsers(dest="cmd", required=True)

    b = sub.add_parser("build", help="export the model to ONNX")
    b.add_argument("--output", default=r"MoonVitOnnx")
    b.add_argument("--grid-h", type=int, default=28, help="patch-grid height (image h = 14*grid_h)")
    b.add_argument("--grid-w", type=int, default=28, help="patch-grid width  (image w = 14*grid_w)")
    b.add_argument("--dynamo", action="store_true",
                   help="use the new dynamo exporter (cleaner/smaller graph; needs onnxscript). "
                        "Default: legacy tracer.")
    b.set_defaults(fn=cmd_build)

    e = sub.add_parser("eval", help="original PyTorch vs ONNX parity")
    e.add_argument("--onnx-model-path", required=True)
    e.add_argument("--samples", type=int, default=3)
    e.add_argument("--tol", type=float, default=0.999)
    e.set_defaults(fn=cmd_eval)

    i = sub.add_parser("inference", help="run the ONNX on an image")
    i.add_argument("--onnx-model-path", required=True)
    i.add_argument("--image", help="input image (any size; resized to the export grid)")
    i.add_argument("--save", help="save features to this .npy")
    i.set_defaults(fn=cmd_inference)

    args = ap.parse_args()
    args.fn(args)


if __name__ == "__main__":
    main()