#!/opt/homebrew/bin/python3.11
"""
Multi-Stage Zoom Search for Stamps
Stage 1: Find containers (hands, envelopes, paper)
Stage 2: Search for stamps inside containers
Stage 3: Filter and rank results
"""

import os
import cv2
import json
import time
from PIL import Image
import torch
from transformers import OwlViTProcessor, OwlViTForObjectDetection

UUID = "384b0ff44aaaa1f1"
VIDEO_PATH = f"output/{UUID}/{UUID}.mp4"
OUTPUT_DIR = f"output/{UUID}/stamp_zoom_search"
os.makedirs(OUTPUT_DIR, exist_ok=True)

FRAME_INTERVAL = 5  # seconds
MIN_CONFIDENCE = 0.08
STAMP_MIN_SIZE = 15  # px
STAMP_MAX_SIZE = 200  # px

print("=" * 60)
print("🔍 Multi-Stage Zoom Search for Stamps")
print("=" * 60)

# ─── Stage 0: Extract frames ───
print("\n📹 Stage 0: Extracting frames...")
cap = cv2.VideoCapture(VIDEO_PATH)
fps = cap.get(cv2.CAP_PROP_FPS)
total_sec = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) / fps)
print(f"  Video: {fps:.1f} fps, {total_sec}s total")

frames_to_process = []
for sec in range(0, total_sec, FRAME_INTERVAL):
    cap.set(cv2.CAP_PROP_POS_MSEC, sec * 1000)
    ret, frame = cap.read()
    if ret:
        frame_path = os.path.join(OUTPUT_DIR, f"frame_{sec}s.jpg")
        cv2.imwrite(frame_path, frame, [cv2.IMWRITE_JPEG_QUALITY, 90])
        frames_to_process.append((sec, frame_path))
    else:
        print(f"  ⚠️  Failed at {sec}s")

cap.release()
print(f"  ✅ Extracted {len(frames_to_process)} frames")

# ─── Stage 1: Find containers ───
print("\n🔎 Stage 1: Loading OWL-ViT for container detection...")
container_processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
container_model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32")
container_model.eval()

CONTAINER_TERMS = [
    "hand holding object",
    "envelope",
    "piece of paper",
    "letter",
    "document",
]

# ─── Stage 2: Find stamps inside containers ───
print("\n🔬 Stage 2: Loading stamp detector...")
stamp_processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
stamp_model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32")
stamp_model.eval()

STAMP_TERMS = ["postage stamp", "stamp", "small stamp"]

# ─── Process ───
all_results = []
stamp_crops_dir = os.path.join(OUTPUT_DIR, "stamp_crops")
os.makedirs(stamp_crops_dir, exist_ok=True)

total_frames = len(frames_to_process)
start_time = time.time()

for idx, (sec, frame_path) in enumerate(frames_to_process):
    elapsed = time.time() - start_time
    eta = (elapsed / (idx + 1)) * (total_frames - idx - 1) if idx > 0 else 0
    print(f"\n  [{idx + 1}/{total_frames}] {sec}s | ETA: {eta:.0f}s")

    image = Image.open(frame_path).convert("RGB")
    img_cv = cv2.imread(frame_path)
    h, w = img_cv.shape[:2]

    # --- Stage 1: Find containers ---
    containers = []
    for term in CONTAINER_TERMS:
        try:
            inputs = container_processor(
                text=[[term]], images=image, return_tensors="pt"
            )
            with torch.no_grad():
                outputs = container_model(**inputs)
            target_sizes = torch.Tensor([h, w])
            results = container_processor.post_process_object_detection(
                outputs=outputs, target_sizes=target_sizes, threshold=0.05
            )
            for score, label, box in zip(
                results[0]["scores"], results[0]["labels"], results[0]["boxes"]
            ):
                s = float(score)
                if s > 0.08:
                    x1, y1, x2, y2 = map(int, box.tolist())
                    # Expand bbox slightly to include surrounding area
                    margin = 30
                    containers.append(
                        {
                            "term": term,
                            "score": s,
                            "bbox": [
                                max(0, x1 - margin),
                                max(0, y1 - margin),
                                min(w, x2 + margin),
                                min(h, y2 + margin),
                            ],
                        }
                    )
        except Exception:
            pass

    if not containers:
        continue

    print(f"    📦 Found {len(containers)} containers")

    # --- Stage 2: Search for stamps inside each container ---
    for container in containers:
        cx1, cy1, cx2, cy2 = container["bbox"]
        container_img = img_cv[cy1:cy2, cx1:cx2]

        if container_img.size == 0:
            continue

        # Scale up the container for better detection
        scale = max(2, 600 // max(container_img.shape[:2]))
        if scale > 1:
            scaled = cv2.resize(
                container_img, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC
            )
        else:
            scaled = container_img

        scaled_pil = Image.fromarray(cv2.cvtColor(scaled, cv2.COLOR_BGR2RGB))
        sh, sw = scaled.shape[:2]

        for stamp_term in STAMP_TERMS:
            try:
                inputs = stamp_processor(
                    text=[[stamp_term]], images=scaled_pil, return_tensors="pt"
                )
                with torch.no_grad():
                    outputs = stamp_model(**inputs)
                target_sizes = torch.Tensor([sh, sw])
                results = stamp_processor.post_process_object_detection(
                    outputs=outputs, target_sizes=target_sizes, threshold=MIN_CONFIDENCE
                )
                for score, label, box in zip(
                    results[0]["scores"], results[0]["labels"], results[0]["boxes"]
                ):
                    s = float(score)
                    if s > MIN_CONFIDENCE:
                        sx1, sy1, sx2, sy2 = box.tolist()

                        # Filter by size (in original coordinates)
                        orig_w = (sx2 - sx1) / scale
                        orig_h = (sy2 - sy1) / scale
                        if not (
                            STAMP_MIN_SIZE < orig_w < STAMP_MAX_SIZE
                            and STAMP_MIN_SIZE < orig_h < STAMP_MAX_SIZE
                        ):
                            continue

                        # Map back to original frame coordinates
                        ox1 = cx1 + int(sx1 / scale)
                        oy1 = cy1 + int(sy1 / scale)
                        ox2 = cx1 + int(sx2 / scale)
                        oy2 = cy1 + int(sy2 / scale)

                        # Crop from original frame
                        crop = img_cv[oy1:oy2, ox1:ox2]
                        if crop.size == 0:
                            continue

                        result = {
                            "timestamp": sec,
                            "container": container["term"],
                            "container_score": container["score"],
                            "stamp_term": stamp_term,
                            "score": s,
                            "bbox": [ox1, oy1, ox2, oy2],
                            "size": [int(orig_w), int(orig_h)],
                        }
                        all_results.append(result)

                        # Save crop
                        crop_name = (
                            f"stamp_{sec}s_{stamp_term.replace(' ', '_')}_{s:.2f}.jpg"
                        )
                        cv2.imwrite(os.path.join(stamp_crops_dir, crop_name), crop)

                        # Save annotated frame
                        ann = img_cv.copy()
                        cv2.rectangle(ann, (ox1, oy1), (ox2, oy2), (0, 255, 0), 3)
                        cv2.putText(
                            ann,
                            f"{stamp_term} {s:.2f}",
                            (ox1, oy1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX,
                            0.7,
                            (0, 255, 0),
                            2,
                        )
                        ann_name = (
                            f"annotated_{sec}s_{stamp_term.replace(' ', '_')}.jpg"
                        )
                        cv2.imwrite(os.path.join(OUTPUT_DIR, ann_name), ann)

                        print(
                            f"      🎯 {sec}s | {stamp_term} | {s:.2f} | {int(orig_w)}x{int(orig_h)}px"
                        )
            except Exception:
                pass

# ─── Stage 3: Filter and rank ───
print(f"\n{'=' * 60}")
print(f"📊 Results: Found {len(all_results)} stamp candidates")
print(f"{'=' * 60}")

# Sort by score
all_results.sort(key=lambda x: x["score"], reverse=True)

# Remove duplicates (same timestamp, similar bbox)
unique_results = []
seen_timestamps = {}
for r in all_results:
    ts = r["timestamp"]
    if ts not in seen_timestamps:
        seen_timestamps[ts] = []
        unique_results.append(r)
        print(
            f"  🎯 {ts}s | {r['stamp_term']} | {r['score']:.2f} | {r['size'][0]}x{r['size'][1]}px | via: {r['container']}"
        )

# Save results
results_path = os.path.join(OUTPUT_DIR, "results.json")
with open(results_path, "w") as f:
    json.dump(unique_results, f, indent=2)

print(f"\n🏁 Done. Results saved to {results_path}")
print(f"📁 Crops saved to {stamp_crops_dir}")
