"""Runs SAM 3 in PyTorch on the test images, whole and as the three split graphs, checks they
agree, and saves the inputs and the split graphs' outputs as references for the ONNX checks."""

import sys
import time
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from transformers import Sam3Model, Sam3Processor

from sam3_graphs import Sam3Detector, Sam3Text, Sam3Vision, position_encoding_72

ROOT = Path(__file__).parent
WEIGHTS = ROOT / "weights"
CASES = {
    "car.png": ["car"],
    "shivan_dragon.jpg": ["dragon", "fire"],
    "serra_angel.jpg": ["angel", "wings"],
    "llanowar_elves.jpg": ["elf", "tree"],
    "lightning_bolt.jpg": ["lightning"],
    "island.jpg": ["sky", "water"],
}


def main() -> None:
    torch.set_grad_enabled(False)
    model = Sam3Model.from_pretrained(WEIGHTS, attn_implementation="eager").eval()
    processor = Sam3Processor.from_pretrained(WEIGHTS)
    vision, text = Sam3Vision(model).eval(), Sam3Text(model).eval()
    detector = Sam3Detector(model, position_encoding_72(model)).eval()
    out_dir = ROOT / "reference"
    out_dir.mkdir(exist_ok=True)

    for image_name, phrases in CASES.items():
        image = Image.open(ROOT / "images" / image_name).convert("RGB")
        pixels = processor(images=image, return_tensors="pt").pixel_values
        start = time.time()
        levels = vision(pixels)
        print(f"{image_name}: vision {time.time() - start:.1f}s, levels {[tuple(l.shape) for l in levels]}")
        for phrase in phrases:
            tokens = processor(text=phrase, return_tensors="pt")
            ids, mask = tokens.input_ids, tokens.attention_mask
            features = text(ids, mask)
            split = detector(*levels, features, mask)
            whole = model(pixel_values=pixels, input_ids=ids, attention_mask=mask)
            pairs = {
                "pred_logits": (split[0], whole.pred_logits),
                "pred_boxes": (split[1], whole.pred_boxes),
                "presence_logits": (split[2], whole.presence_logits),
                "pred_masks": (split[3], whole.pred_masks),
                "semantic_seg": (split[4], whole.semantic_seg),
            }
            diffs = {k: float((a - b).abs().max()) for k, (a, b) in pairs.items()}
            scores = split[0].sigmoid() * split[2].sigmoid()
            kept = int((scores > 0.5).sum())
            print(f"  '{phrase}': ids {tuple(ids.shape)}, kept {kept}, max |split - whole| {diffs}")
            stem = f"{Path(image_name).stem}__{phrase}"
            np.savez_compressed(
                out_dir / f"{stem}.npz",
                pixel_values=pixels.numpy(),
                input_ids=ids.numpy(),
                attention_mask=mask.numpy(),
                features_288=levels[0].numpy(),
                features_144=levels[1].numpy(),
                features_72=levels[2].numpy(),
                text_features=features.numpy(),
                **{k: v[0].numpy() for k, v in pairs.items()},
            )
    sys.stdout.flush()


if __name__ == "__main__":
    main()