#!/opt/homebrew/bin/python3.11
"""
Smart Stamp Score v2 - Pure OpenCV but with better stamp signatures
Key insight: stamps have BORDER + CENTER pattern with different colors
"""

import os
import cv2
import json
import time
import numpy as np

UUID = "384b0ff44aaaa1f1"
VIDEO_PATH = f"output/{UUID}/{UUID}.mp4"
OUTPUT_DIR = f"output/{UUID}/smart_stamp_v2"
os.makedirs(OUTPUT_DIR, exist_ok=True)
CROPS_DIR = os.path.join(OUTPUT_DIR, "crops")
os.makedirs(CROPS_DIR, exist_ok=True)

FRAME_INTERVAL = 5
print("=" * 60)
print("🔍 Smart Stamp Search v2 - Better Scoring")
print("=" * 60)

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: {total_sec}s")


def compute_stamp_score(roi, frame):
    """
    Compute how likely a region is a stamp.
    Stamps have:
    1. Border pattern (edge density high around perimeter)
    2. Color diversity (multiple hues)
    3. Moderate texture (not solid, not pure noise)
    4. Rectangular shape
    """
    h, w = roi.shape[:2]
    if h < 10 or w < 10 or h > 200 or w > 200:
        return 0.0

    aspect = w / h
    if not (0.3 < aspect < 3.0):
        return 0.0

    score = 0.0

    # 1. Color diversity (hues)
    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
    hue = hsv[:, :, 0]
    sat = hsv[:, :, 1]
    val = hsv[:, :, 2]

    # Count significant hues (sat > 30 to ignore grays)
    mask_color = sat > 30
    if np.sum(mask_color) < h * w * 0.1:
        return 0.0  # Too gray/white

    unique_hues = len(np.unique(hue[mask_color]))
    hue_score = min(1.0, unique_hues / 40)
    score += hue_score * 0.3

    # 2. Edge density (stamps have patterns/lines)
    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 30, 100)
    edge_ratio = np.sum(edges > 0) / (h * w)
    edge_score = min(1.0, edge_ratio / 0.15)
    score += edge_score * 0.2

    # 3. Border vs center contrast (stamps have borders)
    border_thickness = max(2, min(h, w) // 6)
    border = np.ones((h, w), dtype=np.uint8) * 255
    border[border_thickness:-border_thickness, border_thickness:-border_thickness] = 0
    center = 255 - border

    border_mean = np.mean(gray[border > 0])
    center_mean = np.mean(gray[center > 0])
    border_center_diff = abs(border_mean - center_mean)
    contrast_score = min(1.0, border_center_diff / 40)
    score += contrast_score * 0.2

    # 4. Hue variance between border and center
    border_hue = hue[border > 0][mask_color[border > 0]]
    center_hue = hue[center > 0][mask_color[center > 0]]

    if len(border_hue) > 5 and len(center_hue) > 5:
        border_hue_mean = np.mean(border_hue)
        center_hue_mean = np.mean(center_hue)
        hue_diff = min(
            abs(border_hue_mean - center_hue_mean),
            180 - abs(border_hue_mean - center_hue_mean),
        )
        hue_diff_score = min(1.0, hue_diff / 60)
        score += hue_diff_score * 0.3
    else:
        score += 0.1

    return min(1.0, score)


all_results = []
start_time = time.time()

for sec in range(0, total_sec, FRAME_INTERVAL):
    cap.set(cv2.CAP_PROP_POS_MSEC, sec * 1000)
    ret, frame = cap.read()
    if not ret:
        continue

    elapsed = time.time() - start_time
    progress = sec / total_sec * 100

    h, w = frame.shape[:2]
    frame_results = []

    # Collect candidate regions (hands + paper)
    candidates = []

    # Skin/hand
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    skin = cv2.inRange(hsv, np.array([0, 20, 60]), np.array([25, 180, 255]))
    skin += cv2.inRange(hsv, np.array([160, 20, 60]), np.array([179, 180, 255]))
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9))
    skin = cv2.morphologyEx(skin, cv2.MORPH_CLOSE, kernel)
    skin = cv2.morphologyEx(skin, cv2.MORPH_OPEN, kernel)

    contours, _ = cv2.findContours(skin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if 1500 < area < h * w * 0.35:
            x, y, cw, ch = cv2.boundingRect(cnt)
            margin = 50
            candidates.append(
                {
                    "type": "hand",
                    "bbox": [
                        max(0, x - margin),
                        max(0, y - margin),
                        min(w, x + cw + margin),
                        min(h, y + ch + margin),
                    ],
                }
            )

    # Paper/envelope
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    _, bright = cv2.threshold(gray, 175, 255, cv2.THRESH_BINARY)
    bright = cv2.morphologyEx(
        bright, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    )
    contours, _ = cv2.findContours(bright, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if 3000 < area < h * w * 0.5:
            x, y, cw, ch = cv2.boundingRect(cnt)
            aspect = cw / ch if ch > 0 else 0
            if 0.2 < aspect < 4.0:
                margin = 40
                candidates.append(
                    {
                        "type": "paper",
                        "bbox": [
                            max(0, x - margin),
                            max(0, y - margin),
                            min(w, x + cw + margin),
                            min(h, y + ch + margin),
                        ],
                    }
                )

    # In each candidate, find small stamp-like regions
    for container in candidates:
        cx1, cy1, cx2, cy2 = container["bbox"]
        region = frame[cy1:cy2, cx1:cx2]

        if region.size == 0:
            continue

        rh, rw = region.shape[:2]
        region_gray = cv2.cvtColor(region, cv2.COLOR_BGR2GRAY)

        # Find small rectangular shapes via edges
        edges = cv2.Canny(region_gray, 30, 100)
        contours_s, _ = cv2.findContours(
            edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )

        for cnt in contours_s:
            area = cv2.contourArea(cnt)
            if 150 < area < 20000:
                x, y, sw, sh = cv2.boundingRect(cnt)
                if not (15 < sw < 150 and 15 < sh < 150):
                    continue

                aspect = sw / sh if sh > 0 else 0
                if not (0.3 < aspect < 3.0):
                    continue

                roi = region[y : y + sh, x : x + sw]
                if roi.size == 0:
                    continue

                stamp_score = compute_stamp_score(roi, frame)

                if stamp_score > 0.4:
                    ox1 = cx1 + x
                    oy1 = cy1 + y
                    ox2 = cx1 + x + sw
                    oy2 = cy1 + y + sh

                    crop = frame[oy1:oy2, ox1:ox2]
                    if crop.size == 0:
                        continue

                    frame_results.append(
                        {
                            "timestamp": sec,
                            "container": container["type"],
                            "score": stamp_score,
                            "bbox": [ox1, oy1, ox2, oy2],
                            "size": [sw, sh],
                            "crop": crop,
                        }
                    )

    if frame_results:
        frame_results.sort(key=lambda x: x["score"], reverse=True)
        # Keep top 3 per frame
        top = frame_results[:3]
        all_results.extend(top)

        print(
            f"  [{sec}s | {progress:.0f}%] Found {len(top)} candidates (top score: {top[0]['score']:.2f})"
        )

        # Save top crops
        for r in top:
            crop_name = f"stamp_{sec}s_{r['container']}_{r['score']:.2f}.jpg"
            cv2.imwrite(os.path.join(CROPS_DIR, crop_name), r["crop"])

            # Annotate
            cv2.rectangle(
                frame,
                (r["bbox"][0], r["bbox"][1]),
                (r["bbox"][2], r["bbox"][3]),
                (0, 255, 0),
                2,
            )
            cv2.putText(
                frame,
                f"{r['score']:.2f}",
                (r["bbox"][0], r["bbox"][1] - 5),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (0, 255, 0),
                1,
            )

        ann_path = os.path.join(OUTPUT_DIR, f"annotated_{sec}s.jpg")
        cv2.imwrite(ann_path, frame)
    else:
        if sec % 120 == 0:
            print(f"  [{sec // 60}min | {progress:.0f}%] Scanning...")

cap.release()

# Sort and deduplicate
all_results.sort(key=lambda x: x["score"], reverse=True)
seen = set()
unique = []
for r in all_results:
    ts = r["timestamp"]
    if ts not in seen:
        seen.add(ts)
        # Remove crop from serializable result
        result_out = {k: v for k, v in r.items() if k != "crop"}
        unique.append(result_out)

print(f"\n{'=' * 60}")
print(f"📊 Found {len(unique)} stamp candidates (score > 0.4)")
for r in unique:
    print(
        f"  🎯 {r['timestamp']}s | {r['container']} | score:{r['score']:.2f} | {r['size'][0]}x{r['size'][1]}px"
    )

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

print(f"\n🏁 Done. Crops: {CROPS_DIR}")
