"""Export SAM 3's vision graph once for both prompts: the Perception Encoder read of a photo feeds
the detector's neck (words) and the tracker's neck (clicks and boxes).

Outputs: features_288/144/72 (as vision.onnx, for the detector) and image_embeddings.0/1/2 (as
onnx-community's sam3-tracker encoder writes them, for its prompt decoder).
    python export_shared.py <weights dir> <out dir>
"""
import sys, time
from pathlib import Path

import torch
from torch import nn
from transformers import Sam3VideoModel

from sam3_graphs import IMAGE_SIZE


class Sam3SharedVision(nn.Module):
    def __init__(self, video):
        super().__init__()
        self.vision_encoder = video.detector_model.vision_encoder
        self.tracker_neck = video.tracker_neck
        self.conv_s0 = video.tracker_model.mask_decoder.conv_s0
        self.conv_s1 = video.tracker_model.mask_decoder.conv_s1
        self.no_memory_embedding = video.tracker_model.no_memory_embedding

    def forward(self, pixel_values: torch.Tensor):
        out = self.vision_encoder(pixel_values)
        levels = out.fpn_hidden_states
        hidden = out.last_hidden_state
        b, side = hidden.shape[0], IMAGE_SIZE // 14
        spatial = hidden.view(b, side, side, -1).permute(0, 3, 1, 2)
        tracked, _ = self.tracker_neck(spatial)
        e0 = self.conv_s0(tracked[0])
        e1 = self.conv_s1(tracked[1])
        e2 = tracked[2] + self.no_memory_embedding.reshape(1, -1, 1, 1)
        return levels[0], levels[1], levels[2], e0, e1, e2


if __name__ == "__main__":
    weights, out = Path(sys.argv[1]), Path(sys.argv[2])
    out.mkdir(parents=True, exist_ok=True)
    t = time.time()
    video = Sam3VideoModel.from_pretrained(weights, attn_implementation="eager", torch_dtype=torch.float32).eval()
    print(f"loaded in {time.time() - t:.0f} s", flush=True)
    module = Sam3SharedVision(video).eval()
    pixels = torch.zeros(1, 3, IMAGE_SIZE, IMAGE_SIZE)
    with torch.no_grad():
        shapes = [tuple(x.shape) for x in module(pixels)]
    print("outputs", shapes, flush=True)
    t = time.time()
    with torch.no_grad():
        torch.onnx.export(module, (pixels,), str(out / "vision.onnx"), input_names=["pixel_values"],
                          output_names=["features_288", "features_144", "features_72",
                                        "image_embeddings.0", "image_embeddings.1", "image_embeddings.2"],
                          opset_version=18, dynamo=True, external_data=True, optimize=True)
    print(f"exported in {time.time() - t:.0f} s", flush=True)