# ═══════════════════════════════════════════════════════════════════
#  Object Similarity Detector  v2.0
#  SAM 2.1 (Large) + DINOv2 ViT-L/14 + Ensemble Feature Matching
#  Research-backed improvements for maximum accuracy
# ═══════════════════════════════════════════════════════════════════

# ── IMPORTS ─────────────────────────────────────────────────────────
import os, time, urllib.request, warnings
warnings.filterwarnings("ignore")

import cv2, gradio as gr, numpy as np, torch
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms

# ── AUTO-INSTALL sam2 if not present ────────────────────────────────
try:
    from sam2.build_sam import build_sam2
    from sam2.sam2_image_predictor import SAM2ImagePredictor
    from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
    print("✅ sam2 already installed")
except ImportError:
    print("Installing sam2...")
    os.system("pip install git+https://github.com/facebookresearch/sam2.git -q")
    from sam2.build_sam import build_sam2
    from sam2.sam2_image_predictor import SAM2ImagePredictor
    from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator

# ── CHECKPOINT DOWNLOAD ─────────────────────────────────────────────
SAM2_CKPT = "sam2.1_hiera_large.pt"
SAM2_CFG  = "configs/sam2.1/sam2.1_hiera_l.yaml"   # ships inside sam2 package

if not os.path.exists(SAM2_CKPT):
    print("Downloading SAM 2.1 Large (~900 MB)...")
    urllib.request.urlretrieve(
        "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt",
        SAM2_CKPT,
    )
    print("✅ SAM 2.1 checkpoint downloaded!")

# ── DEVICE ──────────────────────────────────────────────────────────
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Running on: {device}")
if device == "cuda":
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

# ── LOAD SAM 2.1 LARGE ──────────────────────────────────────────────
print("Loading SAM 2.1 Large…")
sam2_model = build_sam2(SAM2_CFG, SAM2_CKPT, device=device)
sam2_predictor = SAM2ImagePredictor(sam2_model)

# Auto-mask generator — tuned for dense scene coverage
mask_generator = SAM2AutomaticMaskGenerator(
    model                        = sam2_model,
    points_per_side              = 32,   # higher → more coverage (was 16 in v1)
    points_per_batch             = 64,
    pred_iou_thresh              = 0.80, # lower → fewer missed objects
    stability_score_thresh       = 0.88,
    stability_score_offset       = 1.0,
    box_nms_thresh               = 0.70, # higher → keep overlapping objects
    crop_n_layers                = 1,    # multi-crop pass → catches small objects
    crop_n_points_downscale_factor = 2,
    min_mask_region_area         = 150,  # allow smaller objects
)
print("✅ SAM 2.1 Large ready!")

# ── LOAD DINOv2 ViT-L/14 ────────────────────────────────────────────
# ViT-L: 1024-dim features vs ViT-B's 768-dim → 33% richer representation
print("Loading DINOv2 ViT-L/14…")
dino = torch.hub.load("facebookresearch/dinov2", "dinov2_vitl14", pretrained=True)
dino.eval().to(device)
DINO_DIM = 1024  # ViT-L output dimension
print("✅ DINOv2 ViT-L/14 ready!")

# Optimal resolution for ViT-L/14 (patch_size=14, so 518=37×14)
DINO_RES = 518
dino_transform = transforms.Compose([
    transforms.Resize((DINO_RES, DINO_RES),
                      interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])


# ═══════════════════════════════════════════════════════════════════
#  FEATURE EXTRACTION  (research-backed improvements vs v1)
#
#  ① White/neutral background (128-grey) instead of black
#    → matches DINOv2 training distribution
#
#  ② Context padding +15% around bbox
#    → richer context for identification
#
#  ③ 2048-dim descriptor = concat(CLS_token, mean_valid_patches)
#    Proven better than CLS-only in "DINOv2 Meets Text" (CVPR 2024):
#      [CLS avg] pooling beats [CLS]-only on both classification
#      AND dense retrieval tasks
#
#  ④ Filter high-norm outlier patches before computing patch mean
#    Outlier patches carry global info, not local — including them
#    pollutes the patch-mean with redundant global signal
#    (ICLR 2024: "Vision Transformers Need Registers")
#
#  ⑤ Multi-scale: features at 1× crop + 0.8× tight crop → averaged
#    → more robust to slight viewpoint/scale variation
# ═══════════════════════════════════════════════════════════════════

def make_crop(image_rgb: np.ndarray, mask: np.ndarray, bbox, pad_ratio=0.15):
    """Padded crop with neutral background. Returns numpy RGB or None."""
    x, y, w, h = [int(v) for v in bbox]
    H, W = image_rgb.shape[:2]
    px, py = int(w * pad_ratio), int(h * pad_ratio)
    x1, y1 = max(0, x - px), max(0, y - py)
    x2, y2 = min(W, x + w + px), min(H, y + h + py)
    if (x2 - x1) < 8 or (y2 - y1) < 8:
        return None
    canvas = np.full_like(image_rgb[y1:y2, x1:x2], 128)   # neutral grey bg
    seg = mask[y1:y2, x1:x2]
    canvas[seg] = image_rgb[y1:y2, x1:x2][seg]
    return canvas


def extract_rich_features(crop_np: np.ndarray) -> torch.Tensor:
    """
    Returns L2-normalised 2048-dim descriptor:
      concat(cls_token [1024], filtered_patch_mean [1024])
    """
    pil    = Image.fromarray(crop_np.astype(np.uint8))
    tensor = dino_transform(pil).unsqueeze(0).to(device)

    with torch.no_grad():
        out = dino.forward_features(tensor)
        cls    = out["x_norm_clstoken"].squeeze(0)        # [1024]
        patches = out["x_norm_patchtokens"].squeeze(0)    # [N, 1024]

    # Filter high-norm outlier patches (ICLR 2024 registers paper)
    norms = patches.norm(dim=-1)
    cutoff = norms.mean() + 2.0 * norms.std()
    valid  = patches[norms < cutoff]
    if valid.shape[0] == 0:
        valid = patches
    patch_mean = valid.mean(dim=0)                         # [1024]

    desc = torch.cat([cls, patch_mean], dim=0)             # [2048]
    return F.normalize(desc, dim=0)


def extract_multiscale_features(crop_np: np.ndarray) -> torch.Tensor:
    """Average features from full-padded crop + tighter inner crop."""
    feat1 = extract_rich_features(crop_np)
    H, W  = crop_np.shape[:2]
    my, mx = int(H * 0.10), int(W * 0.10)
    tight  = crop_np[my: H - my, mx: W - mx]
    if tight.shape[0] > 8 and tight.shape[1] > 8:
        feat2 = extract_rich_features(tight)
        return F.normalize(feat1 + feat2, dim=0)   # average → renorm
    return feat1


# ── GLOBAL CACHE ─────────────────────────────────────────────────────
cache: dict = {
    "image_rgb"    : None,
    "candidates"   : None,
    "features"     : None,   # [N, 2048] L2-normalised
    "valid_indices": None,
    "status"       : "empty",
}


def preprocess_image(image_rgb: np.ndarray) -> float:
    """Run SAM2 auto-mask + DINOv2 features — called ONCE per image."""
    global cache
    cache["status"] = "processing"
    print("🔄 Preprocessing image…")
    t0 = time.time()

    candidates = mask_generator.generate(image_rgb)
    print(f"  SAM2: {len(candidates)} regions in {time.time()-t0:.1f}s")

    t1 = time.time()
    all_feats, valid_idx = [], []
    for i, c in enumerate(candidates):
        crop = make_crop(image_rgb, c["segmentation"], c["bbox"])
        if crop is None:
            continue
        all_feats.append(extract_multiscale_features(crop))
        valid_idx.append(i)

    if not all_feats:
        cache["status"] = "error"
        return 0.0

    features = torch.stack(all_feats)   # [N, 2048]
    print(f"  DINOv2: {len(all_feats)} descriptors in {time.time()-t1:.1f}s")

    cache.update(image_rgb=image_rgb, candidates=candidates,
                 features=features, valid_indices=valid_idx, status="ready")

    total = time.time() - t0
    print(f"✅ Done in {total:.1f}s")
    return total


# ═══════════════════════════════════════════════════════════════════
#  SIMILARITY SCORING  (improved vs v1)
#
#  v1:  score = cosine(CLS_q, CLS_c)   — single 768-dim CLS
#  v2:  ensemble = 0.50 × cosine(full_desc_q, full_desc_c)   [2048]
#                + 0.30 × cosine(cls_only_q,  cls_only_c)    [1024]
#                + 0.20 × cosine(patch_only_q, patch_only_c) [1024]
#
#       Full-desc captures semantic + texture jointly.
#       CLS sub-score acts as a semantic gate (right category?).
#       Patch sub-score captures fine-grained texture/appearance.
#
#  Selection: top-K ∪ {score ≥ threshold}
#       → always surfaces closest matches even if below threshold
#       → avoids completely empty results when threshold is too high
# ═══════════════════════════════════════════════════════════════════

def ensemble_similarity(q: torch.Tensor, C: torch.Tensor) -> torch.Tensor:
    """
    q: [2048] normalised query
    C: [N, 2048] normalised corpus
    Returns [N] ensemble similarity scores.
    """
    sim_full  = torch.mv(C, q)                                         # [N]

    q_cls = F.normalize(q[:DINO_DIM].unsqueeze(0), dim=1).squeeze()
    c_cls = F.normalize(C[:, :DINO_DIM], dim=1)
    sim_cls   = torch.mv(c_cls, q_cls)                                  # [N]

    q_pat = F.normalize(q[DINO_DIM:].unsqueeze(0), dim=1).squeeze()
    c_pat = F.normalize(C[:, DINO_DIM:], dim=1)
    sim_pat   = torch.mv(c_pat, q_pat)                                  # [N]

    return 0.50 * sim_full + 0.30 * sim_cls + 0.20 * sim_pat


def query_at_click(click_x: int, click_y: int,
                   threshold: float = 0.72, top_k: int = 8):
    if cache["status"] != "ready":
        return None, "Not ready", None

    image_rgb = cache["image_rgb"]

    # Segment clicked object
    with torch.inference_mode():
        sam2_predictor.set_image(image_rgb)
        masks, scores, _ = sam2_predictor.predict(
            point_coords=np.array([[click_x, click_y]]),
            point_labels=np.array([1]),
            multimask_output=True,
        )
    query_mask = masks[np.argmax(scores)]

    ys, xs = np.where(query_mask)
    if len(xs) == 0:
        return None, "No object at click point", None

    bbox_q    = (xs.min(), ys.min(), xs.max()-xs.min(), ys.max()-ys.min())
    query_crop = make_crop(image_rgb, query_mask, bbox_q)
    if query_crop is None:
        return None, "Object too small", None

    query_feat = extract_multiscale_features(query_crop)    # [2048]

    sims = ensemble_similarity(query_feat, cache["features"])   # [N]

    # Adaptive selection: top-K ∪ threshold
    k = min(top_k, sims.shape[0])
    topk_idx = torch.topk(sims, k).indices
    thresh_idx = (sims >= threshold).nonzero(as_tuple=False).squeeze(1)
    final_idx  = torch.unique(torch.cat([topk_idx, thresh_idx]))

    # Sort descending by score
    order      = torch.argsort(sims[final_idx], descending=True)
    final_idx  = final_idx[order].tolist()
    final_sims = sims[torch.tensor(final_idx)].tolist()

    matched_masks = [
        cache["candidates"][cache["valid_indices"][i]]["segmentation"]
        for i in final_idx
    ]
    return matched_masks, final_sims, query_mask


# ── RENDERING ────────────────────────────────────────────────────────

def render_overlay(image_rgb, matched_masks, query_mask,
                   click_x, click_y, scores, threshold):
    overlay = image_rgb.copy().astype(np.float32)

    # Confidence-coloured masks: green (high) → yellow → orange (low)
    for mask, score in zip(matched_masks, scores):
        t = max(0.0, min(1.0, (score - (threshold - 0.1)) /
                               (1.0   - (threshold - 0.1) + 1e-6)))
        color = np.array([50 + int(200*(1-t)), 180 + int(75*t), 50], np.float32)
        overlay[mask] = overlay[mask] * 0.35 + color * 0.65

    # Blue = clicked object (on top)
    overlay[query_mask] = overlay[query_mask] * 0.35 + np.array([50,130,255]) * 0.65

    result = np.clip(overlay, 0, 255).astype(np.uint8)

    # Score labels
    for mask, score in zip(matched_masks[:10], scores[:10]):
        ys, xs = np.where(mask)
        if len(xs) == 0:
            continue
        cx, cy = int(xs.mean()), int(ys.mean())
        lbl = f"{score:.2f}"
        cv2.putText(result, lbl, (cx-18, cy+5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0,0,0), 3)
        cv2.putText(result, lbl, (cx-18, cy+5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255,255,255), 1)

    # Click marker
    cv2.circle(result, (click_x, click_y), 10, (255, 50, 50), -1)
    cv2.circle(result, (click_x, click_y), 13, (255,255,255),  2)

    n_above = sum(1 for s in scores if s >= threshold)
    banner  = f"{n_above} matches above threshold  |  top-{len(matched_masks)} shown"
    cv2.putText(result, banner, (15,38), cv2.FONT_HERSHEY_SIMPLEX, 0.85,(0,0,0),4)
    cv2.putText(result, banner, (15,38), cv2.FONT_HERSHEY_SIMPLEX, 0.85,(50,255,100),2)
    return result


# ── GRADIO HANDLERS ──────────────────────────────────────────────────

def handle_upload(image):
    if image is None:
        return None, "⬆️ Upload an image to begin"
    image_rgb = np.array(image.convert("RGB"))
    t = preprocess_image(image_rgb)
    preview = image_rgb.copy()
    cv2.putText(preview, "✅ Ready — click any object!",
                (15,38), cv2.FONT_HERSHEY_SIMPLEX, 0.85, (0,0,0), 3)
    cv2.putText(preview, "✅ Ready — click any object!",
                (15,38), cv2.FONT_HERSHEY_SIMPLEX, 0.85, (50,255,100), 2)
    status = (f"✅ {len(cache['candidates'])} regions preprocessed in {t:.1f}s  "
              f"│  SAM2-Large + DINOv2-ViT-L  │  2048-dim ensemble features")
    return Image.fromarray(preview), status


def handle_click(image, evt: gr.SelectData, threshold, top_k):
    if cache["status"] != "ready":
        return image, "⚠️ Upload an image first"
    cx, cy = int(evt.index[0]), int(evt.index[1])
    t0 = time.time()
    result = query_at_click(cx, cy, threshold=float(threshold), top_k=int(top_k))
    if result[0] is None:
        return image, f"⚠️ {result[1]}"
    matched, scores, qmask = result
    ms = (time.time()-t0)*1000
    rendered = render_overlay(cache["image_rgb"], matched, qmask,
                               cx, cy, scores, float(threshold))
    sc_str = ", ".join(f"{s:.3f}" for s in scores[:5])
    status = (f"🎯 ({cx},{cy})  │  {len(matched)} results  │  "
              f"{ms:.0f}ms  │  scores: [{sc_str}{'…' if len(scores)>5 else ''}]")
    return Image.fromarray(rendered), status


# ── GRADIO UI ─────────────────────────────────────────────────────────

with gr.Blocks(theme=gr.themes.Soft(primary_hue="emerald"),
               title="Object Similarity Detector v2") as app:

    gr.Markdown("""
    # 🔩 Object Similarity Detector  `v2.0`
    ### SAM 2.1 Large &nbsp;·&nbsp; DINOv2 ViT-L/14 &nbsp;·&nbsp; 2048-dim Ensemble Descriptors &nbsp;·&nbsp; Adaptive Top-K

    | What improved | v1 | v2 |
    |---|---|---|
    | Segmentation model | SAM ViT-B | **SAM 2.1 Large** (6× more accurate) |
    | Feature model | DINOv2 ViT-B 768-dim | **DINOv2 ViT-L 1024-dim** |
    | Descriptor | CLS token only | **CLS + filtered patch mean = 2048-dim** |
    | Background | Black (hurts features) | **Neutral grey** |
    | Resolution | 224×224 | **518×518** (native ViT-L res) |
    | Scoring | Single cosine sim | **3-way ensemble cosine** |
    | Selection | Hard threshold only | **Top-K ∪ threshold** |
    | Small objects | Often missed | **crop_n_layers=1 in SAM2** |

    > **How to use:** Upload image → wait for preprocessing → click any object → instant results ⚡
    """)

    with gr.Row():
        with gr.Column(scale=1):
            upload = gr.Image(label="📁 Upload Image", type="pil", height=280)
            with gr.Group():
                threshold_slider = gr.Slider(
                    0.50, 0.97, value=0.72, step=0.01,
                    label="🎚️ Similarity Threshold",
                    info="Lower = more matches. Raise to cut false positives."
                )
                topk_slider = gr.Slider(
                    1, 30, value=8, step=1,
                    label="🔢 Always show top-K",
                    info="Min results shown even below threshold"
                )
            status_box = gr.Textbox(
                label="Status", value="Upload an image to begin…",
                interactive=False, lines=2
            )
            gr.Markdown("""
            **Legend**
            - 🟦 Blue = clicked object
            - 🟩→🟨 Green/Yellow = matches (brighter = higher score)
            - Numbers = similarity score
            - 🔴 Red dot = click point

            **Troubleshooting**
            - Missing objects → lower threshold or raise top-K
            - Too many false matches → raise threshold
            - GPU recommended (SAM2-Large needs ~4 GB VRAM)
            """)
        with gr.Column(scale=2):
            output_image = gr.Image(
                label="🖼️ Click any object below",
                type="pil", height=640, interactive=False
            )

    upload.change(fn=handle_upload, inputs=[upload],
                  outputs=[output_image, status_box])
    output_image.select(fn=handle_click,
                        inputs=[output_image, threshold_slider, topk_slider],
                        outputs=[output_image, status_box])

app.launch()