# SPDX-License-Identifier: Apache-2.0
"""Quickstart: the model card's Python snippet on the demo image.

    pip install -e code/
    python examples/quickstart.py [image] [prompt ...]

Writes examples/out/quickstart_overlay.png (masks and boxes on the image) and examples/out/quickstart.json.
The default image and the output directory are found relative to this file, so it runs from any directory; an image
path given on the command line is relative to the current directory.
"""
import json
import os
import sys
import time

import numpy as np
from PIL import Image, ImageDraw

from tt_sam3 import Sam3

HERE = os.path.dirname(os.path.abspath(__file__))
image_path = sys.argv[1] if len(sys.argv) > 1 else os.path.join(HERE, "..", "media", "source_1.png")
prompts = sys.argv[2:] or ["car"]
out_dir = os.path.join(HERE, "out")
os.makedirs(out_dir, exist_ok=True)

# --- the model card snippet -----------------------------------------------------------------------------
with Sam3.from_pretrained(device_id=0, warmup_variants={"prompts": prompts}) as model:  # warm: call 1 is fast too
    results = model(image_path, prompts)            # one Sam3Result per prompt
    t0 = time.perf_counter()
    model(image_path, prompts)                      # warm call (same image: the image part is reused)
    warm_ms = (time.perf_counter() - t0) * 1e3
# ---------------------------------------------------------------------------------------------------------

for r in results:
    print(r)
print(f"warm call (same image): {warm_ms:.1f} ms")

# overlay: one color per instance, mask blended at 50 %, box outline, score label
image = Image.open(image_path).convert("RGB")
canvas = np.asarray(image).astype(np.float32)
palette = np.array([[230, 25, 75], [60, 180, 75], [255, 225, 25], [0, 130, 200], [245, 130, 48], [145, 30, 180],
                    [70, 240, 240], [240, 50, 230], [210, 245, 60], [250, 190, 212]], dtype=np.float32)
k = 0
for r in results:
    for m in r.masks:
        canvas[m] = 0.5 * canvas[m] + 0.5 * palette[k % len(palette)]
        k += 1
overlay = Image.fromarray(canvas.astype(np.uint8))
draw = ImageDraw.Draw(overlay)
k = 0
for r in results:
    for box, score in zip(r.boxes, r.scores):
        color = tuple(int(c) for c in palette[k % len(palette)])
        draw.rectangle([float(v) for v in box], outline=color, width=3)
        draw.text((float(box[0]) + 4, float(box[1]) + 4), f"{r.prompt} {score:.2f}", fill=color)
        k += 1
overlay.save(os.path.join(out_dir, "quickstart_overlay.png"))
with open(os.path.join(out_dir, "quickstart.json"), "w") as f:
    json.dump({"image": os.path.relpath(image_path, os.path.join(HERE, "..")), "warm_call_ms": round(warm_ms, 1),
               "results": [r.to_dict() for r in results]}, f, indent=1)
print("wrote", os.path.join(out_dir, "quickstart_overlay.png"), "and quickstart.json")