# Derived from NVIDIA LocateAnything-3B (modeling_vit.py / modeling_locateanything.py, Copyright (c) 2026 NVIDIA
# CORPORATION) under the NVIDIA License (see ../LICENSE); MoonViT is MIT-licensed by Moonshot AI. Modified.
import sys, math, torch, torch.nn as nn, torch.nn.functional as F
from safetensors import safe_open
GH = int(sys.argv[1]) if len(sys.argv) > 1 else 32     # patch grid (H/14, W/14)
GW = int(sys.argv[2]) if len(sys.argv) > 2 else 32
D, HEADS, HD, LAYERS, FF, P = 1152, 16, 72, 27, 4304, 14
MP = "./model"

def load():
    w = {}
    import json
    idx = json.load(open(f"{MP}/model.safetensors.index.json"))["weight_map"]
    files = sorted(set(idx.values()))
    for f in files:
        with safe_open(f"{MP}/{f}", "pt") as sf:
            for k in sf.keys():
                if k.startswith("vision_model.") or k.startswith("mlp1."):
                    w[k] = sf.get_tensor(k).float()
    return w

def gelu_tanh(x):
    return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * x * x * x)))

class Block(nn.Module):
    def __init__(s):
        super().__init__()
        s.norm0 = nn.LayerNorm(D); s.norm1 = nn.LayerNorm(D)
        s.wqkv = nn.Linear(D, 3 * D); s.wo = nn.Linear(D, D)
        s.fc0 = nn.Linear(D, FF); s.fc1 = nn.Linear(FF, D)
    def forward(s, x, cos, sin):                      # x (1,N,D); cos/sin (1,N,1,36)
        N = x.shape[1]
        qkv = s.wqkv(s.norm0(x)).view(1, N, 3, HEADS, HD)
        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]          # (1,N,H,HD)
        def rot(t):
            t = t.view(1, N, HEADS, HD // 2, 2)
            r, i = t[..., 0], t[..., 1]
            return torch.stack([r * cos - i * sin, r * sin + i * cos], dim=-1).view(1, N, HEADS, HD)
        q, k = rot(q), rot(k)
        q = q.permute(0, 2, 1, 3); k = k.permute(0, 2, 1, 3); v = v.permute(0, 2, 1, 3)
        a = torch.softmax(torch.matmul(q, k.transpose(-1, -2)) * (1.0 / math.sqrt(HD)), dim=-1)
        o = torch.matmul(a, v).permute(0, 2, 1, 3).reshape(1, N, D)
        x = x + s.wo(o)
        return x + s.fc1(gelu_tanh(s.fc0(s.norm1(x))))

class ViTFixed(nn.Module):
    def __init__(s, w):
        super().__init__()
        N = GH * GW
        s.proj = nn.Conv2d(3, D, P, P)
        s.proj.weight.data = w["vision_model.patch_embed.proj.weight"]; s.proj.bias.data = w["vision_model.patch_embed.proj.bias"]
        pe = w["vision_model.patch_embed.pos_emb.weight"]                # (64,64,D)
        if (GH, GW) != tuple(pe.shape[:2]):
            pe = F.interpolate(pe.permute(2, 0, 1).unsqueeze(0), size=(GH, GW), mode="bicubic").squeeze(0).permute(1, 2, 0)
        s.register_buffer("pos", pe.reshape(1, N, D).contiguous())
        # 2D rope: angle[:,2i]=x*f_i, angle[:,2i+1]=y*f_i  (x = column, y = row)
        f = 1.0 / (10000 ** (torch.arange(0, HD, 4)[: HD // 4].float() / HD))
        ys, xs = torch.meshgrid(torch.arange(GH).float(), torch.arange(GW).float(), indexing="ij")
        xs, ys = xs.reshape(-1), ys.reshape(-1)
        ang = torch.stack([torch.outer(xs, f), torch.outer(ys, f)], dim=-1).reshape(N, HD // 2)
        s.register_buffer("cos", ang.cos().view(1, N, 1, HD // 2)); s.register_buffer("sin", ang.sin().view(1, N, 1, HD // 2))
        s.blocks = nn.ModuleList([Block() for _ in range(LAYERS)])
        for i, b in enumerate(s.blocks):
            p = f"vision_model.encoder.blocks.{i}."
            b.norm0.weight.data = w[p+"norm0.weight"]; b.norm0.bias.data = w[p+"norm0.bias"]
            b.norm1.weight.data = w[p+"norm1.weight"]; b.norm1.bias.data = w[p+"norm1.bias"]
            b.wqkv.weight.data = w[p+"wqkv.weight"]; b.wqkv.bias.data = w[p+"wqkv.bias"]
            b.wo.weight.data = w[p+"wo.weight"]; b.wo.bias.data = w[p+"wo.bias"]
            b.fc0.weight.data = w[p+"mlp.fc0.weight"]; b.fc0.bias.data = w[p+"mlp.fc0.bias"]
            b.fc1.weight.data = w[p+"mlp.fc1.weight"]; b.fc1.bias.data = w[p+"mlp.fc1.bias"]
        s.fln = nn.LayerNorm(D)
        s.fln.weight.data = w["vision_model.encoder.final_layernorm.weight"]; s.fln.bias.data = w["vision_model.encoder.final_layernorm.bias"]
        s.m0 = nn.LayerNorm(4 * D); s.m1 = nn.Linear(4 * D, 2048); s.m3 = nn.Linear(2048, 2048)
        s.m0.weight.data = w["mlp1.0.weight"]; s.m0.bias.data = w["mlp1.0.bias"]
        s.m1.weight.data = w["mlp1.1.weight"]; s.m1.bias.data = w["mlp1.1.bias"]
        s.m3.weight.data = w["mlp1.3.weight"]; s.m3.bias.data = w["mlp1.3.bias"]
    def forward(s, pixel):                             # (1,3,H,W), normalized with mean=std=0.5
        x = s.proj(pixel).flatten(2).transpose(1, 2) + s.pos          # (1,N,D), raster order
        for b in s.blocks: x = b(x, s.cos, s.sin)
        x = s.fln(x)
        x = x.view(GH // 2, 2, GW // 2, 2, D).permute(0, 2, 1, 3, 4).reshape(1, (GH // 2) * (GW // 2), 4 * D)
        x = s.m3(F.gelu(s.m1(s.m0(x))))
        return x[0]                                     # (N/4, 2048)

if __name__ == "__main__":
    m = ViTFixed(load()).eval()
    ref = torch.load("ref_vit.pt")
    with torch.no_grad(): out = m(ref["pixel"])
    r = ref["out"]
    cos = F.cosine_similarity(out.flatten(), r.flatten(), dim=0).item()
    print("shape", tuple(out.shape), "cos", round(cos, 6), "max|diff|", (out - r).abs().max().item(), "ref std", r.std().item())
    if len(sys.argv) > 3 and sys.argv[3] == "export":
        torch.onnx.export(m, ref["pixel"], f"out/vit_{GH*14}x{GW*14}.onnx", input_names=["pixel"], output_names=["embeds"], opset_version=17, do_constant_folding=True)
        print("exported")