"""
What a port of SAM 3.1's still-picture paths is held to: the makers' own paths, in fp32, on their own sample pictures.

  reference.json   for each picture, each click case's points, labels, IoUs and object score, and each phrase's
                   scores and boxes (x0 y0 x1 y1, picture pixels); the tokens of each of PHRASES
  <picture>-click-<case>.png, <picture>-find-<case>-<n>.png   each mask at the picture's size, 1-bit

usage: sam_reference.py <checkpoint> <pictures folder> <reference folder>
"""
import json, sys
from pathlib import Path

import numpy as np
from PIL import Image

from sam_graphs import CASES, PHRASES, Sam


def save_mask(mask, path: Path):
    Image.fromarray(np.asarray(mask, bool)).convert("1").save(path, optimize=True)


def main():
    checkpoint, pictures, out = sys.argv[1], Path(sys.argv[2]), Path(sys.argv[3])
    out.mkdir(parents=True, exist_ok=True)
    sam = Sam(checkpoint)
    reference = {"pictures": {}, "tokens": {}}
    for name, (click_cases, find_cases) in CASES.items():
        picture = Image.open(pictures / name).convert("RGB")
        w, h = picture.size
        stem = Path(name).stem
        entry = {"width": w, "height": h, "clicks": [], "finds": []}
        features = sam.click_features(picture)
        for case, points, labels in click_cases:
            mask, ious, obj = sam.click(features, points, labels, w, h)
            file = f"{stem}-click-{case}.png"
            save_mask((mask > 0).cpu().numpy(), out / file)
            entry["clicks"].append({"case": case, "points": points, "labels": labels, "ious": ious.cpu().tolist(),
                                    "object": obj, "mask": file})
        for case, phrase, boxes, labels in find_cases:
            scores, xyxy, masks = sam.find(picture, phrase, boxes, labels)
            files = []
            for i, m in enumerate(masks.cpu().numpy()):
                files.append(f"{stem}-find-{case}-{i}.png")
                save_mask(m, out / files[-1])
            entry["finds"].append({"case": case, "phrase": phrase, "boxes": boxes, "labels": labels,
                                   "tokens": sam.tokens(phrase)[0].tolist(), "scores": scores.cpu().tolist(),
                                   "found": xyxy.cpu().tolist(), "masks": files})
        reference["pictures"][name] = entry
        print(f"{name}: {len(click_cases)} click case(s), {len(find_cases)} phrase(s)", flush=True)
    for phrase in PHRASES:
        reference["tokens"][phrase] = sam.tokens(phrase)[0].tolist()
    (out / "reference.json").write_text(json.dumps(reference, indent=1, ensure_ascii=False), encoding="utf-8")


if __name__ == "__main__":
    main()