import spaces  # MUST come before torch / any CUDA-touching import

import os
import sys
import math
import time

import cv2
import numpy as np
import torch
import torch.nn.functional as F

import gradio as gr
from huggingface_hub import hf_hub_download

# Make the bundled src/ package importable (SAMatcher modeling code).
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import src  # noqa: F401  (SAMatcher modeling package)
from hydra.utils import instantiate
from omegaconf import OmegaConf

MODEL_REPO = "SSSSphinx/SAMatcher"
CONFIG_PATH = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "src", "configs",
    "sam2.1", "sam2.1_hiera_l_sammatcher.yaml")
INPUT_SIZE = 1024
STRIDE_ALIGN = 16  # eval scripts use align='disk' -> stride 16 on the long side


def load_model():
    """Instantiate SAM2HQBase from the authors' config and load the released checkpoint.

    The config is loaded directly with OmegaConf (hydra's pkg:// search path is
    environment-sensitive inside the Space container), then instantiated with
    hydra's recursive instantiate helper, exactly as the authors' build script does.
    """
    cfg = OmegaConf.load(CONFIG_PATH)
    model = instantiate(cfg.model, _recursive_=True)
    model.eval()
    ckpt_path = hf_hub_download(MODEL_REPO, "weights/samatcher_best.ckpt")
    state = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    missing, unexpected = model.load_state_dict(state, strict=False)
    # The 6 fuser logit_scale params exist only on the cosine-attention fallback
    # path; they are absent from the checkpoint by design, so ignore them.
    real_missing = [k for k in missing if not k.endswith("attn.logit_scale")]
    if real_missing or unexpected:
        raise RuntimeError(
            f"Checkpoint mismatch — missing={real_missing[:5]} unexpected={list(unexpected)[:5]}"
        )
    return model


print("Loading SAMatcher (this downloads ~960 MB of weights on first boot)...")
model = load_model()
model.to("cuda")
print("SAMatcher loaded.")


# ---------------------------------------------------------------------------
# Preprocessing — mirrors dloc read_overlap_image(align='disk', resize=1024,
# pad_square=True, RGB): long side -> 1024, stride-aligned to 16, zero-padded
# to a 1024x1024 square at bottom/right, /255, NCHW.
# ---------------------------------------------------------------------------

def preprocess(image_rgb: np.ndarray):
    """Returns (model_input_tensor[1,3,S,S], content_h, content_w, scale_x, scale_y)."""
    h, w = image_rgb.shape[:2]
    scale = INPUT_SIZE / max(h, w)
    w_r, h_r = max(1, int(round(w * scale))), max(1, int(round(h * scale)))
    w_new = int(math.ceil(w_r / STRIDE_ALIGN) * STRIDE_ALIGN)
    h_new = int(math.ceil(h_r / STRIDE_ALIGN) * STRIDE_ALIGN)
    img = cv2.resize(image_rgb.astype(np.float32), (w_new, h_new), interpolation=cv2.INTER_LINEAR)
    padded = np.zeros((INPUT_SIZE, INPUT_SIZE, 3), dtype=np.float32) if max(w_new, h_new) <= INPUT_SIZE else None
    if padded is None:  # extremely wide/tall images: hard-clip to 1024
        padded = np.zeros((INPUT_SIZE, INPUT_SIZE, 3), dtype=np.float32)
        w_new, h_new = min(w_new, INPUT_SIZE), min(h_new, INPUT_SIZE)
    padded[:h_new, :w_new, :] = img[:h_new, :w_new, :]
    inp = torch.from_numpy(padded.transpose(2, 0, 1)[None] / 255.0).float()
    sx = float(w) / float(w_new)
    sy = float(h) / float(h_new)
    return inp, h_new, w_new, sx, sy


# ---------------------------------------------------------------------------
# Visualization — replicates the authors' evaluation plotting:
# magenta soft-mask overlay (alpha 0.35, keeps brightness outside the mask) +
# magenta contour + red bbox, and a cropped matching-window pair.
# ---------------------------------------------------------------------------

def render_overlay(image_rgb: np.ndarray, soft_mask: np.ndarray, box_xyxy) -> np.ndarray:
    img = image_rgb.astype(np.float32) / 255.0
    m = np.clip(soft_mask, 0.0, 1.0).astype(np.float32)
    if m.shape != image_rgb.shape[:2]:
        m = cv2.resize(m, (image_rgb.shape[1], image_rgb.shape[0]), interpolation=cv2.INTER_LINEAR)
    out = img.copy()
    a = (0.35 * m)[..., None]
    magenta = np.array([1.0, 0.0, 1.0], dtype=np.float32)
    out = (1.0 - a) * out + a * magenta
    out_u8 = np.clip(out * 255.0, 0, 255).astype(np.uint8)

    m_bin = (m > 0.5).astype(np.uint8)
    if m_bin.sum() > 0:
        contours, _ = cv2.findContours(m_bin, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(out_u8, contours, -1, (255, 0, 255), 2, lineType=cv2.LINE_AA)
    if box_xyxy is not None:
        x0, y0, x1, y1 = [int(round(float(v))) for v in box_xyxy]
        h_img, w_img = out_u8.shape[:2]
        x0 = max(0, min(w_img - 1, x0)); x1 = max(0, min(w_img, x1))
        y0 = max(0, min(h_img - 1, y0)); y1 = max(0, min(h_img, y1))
        if x1 > x0 and y1 > y0:
            cv2.rectangle(out_u8, (x0, y0), (x1, y1), (255, 0, 0), max(2, w_img // 300))
    return out_u8


def render_crop_pair(image_rgb: np.ndarray, box_xyxy, min_side: int = 224) -> np.ndarray:
    """Crop the predicted co-visible window; returns RGB crop (uint8)."""
    h, w = image_rgb.shape[:2]
    if box_xyxy is None:
        return image_rgb.copy()
    x0, y0, x1, y1 = [int(round(float(v))) for v in box_xyxy]
    x0 = max(0, min(w - 1, x0)); x1 = max(0, min(w, x1))
    y0 = max(0, min(h - 1, y0)); y1 = max(0, min(h, y1))
    if (x1 - x0) < 8 or (y1 - y0) < 8:
        return image_rgb.copy()
    crop = image_rgb[y0:y1, x0:x1]
    ch, cw = crop.shape[:2]
    if min(ch, cw) < min_side:  # upscale tiny crops for visibility
        s = min_side / min(ch, cw)
        crop = cv2.resize(crop, (max(1, int(round(cw * s))), max(1, int(round(ch * s)))),
                         interpolation=cv2.INTER_CUBIC)
    return np.clip(crop, 0, 255).astype(np.uint8)


def run_pair(image0, image1, mask_thresh=0.3, use_boxes=True):
    """Full SAMatcher inference on one image pair. Returns numpy visualizations on CPU."""
    imgs = []
    metas = []
    for image in (image0, image1):
        rgb = image if image.ndim == 3 else cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        if rgb.shape[2] == 4:
            rgb = cv2.cvtColor(rgb, cv2.COLOR_RGBA2RGB)
        inp, ch, cw, sx, sy = preprocess(rgb)
        imgs.append(inp)
        metas.append((rgb, ch, cw, sx, sy))

    batch = torch.cat([imgs[0], imgs[1]], dim=0).to("cuda")
    try:
        with torch.no_grad():
            masks, (bbox0, bbox1) = model.forward(batch[0:1], batch[1:2])
    finally:
        del batch, imgs

    return finalize(metas, masks, bbox0, bbox1, float(mask_thresh))


def finalize(metas, masks, bbox0, bbox1, mask_thresh):
    """Map raw outputs (padded-1024 coords) back to original image coords."""
    outs = []
    for idx, (rgb, ch, cw, sx, sy) in enumerate(metas):
        raw = masks[0, idx]                      # (256, 256) logits
        # Reference order: 256 -> padded 1024 -> crop content area -> original size
        raw = F.interpolate(raw[None, None], size=(INPUT_SIZE, INPUT_SIZE),
                            mode="bilinear", align_corners=False)
        raw = raw[:, :, :ch, :cw]
        raw = F.interpolate(raw, size=rgb.shape[:2], mode="bilinear",
                            align_corners=False)
        soft = raw[0, 0].sigmoid().cpu().numpy()

        # bbox: clamp to the un-padded content area, then scale to original coords
        box = bbox0[0].detach().cpu().float() if idx == 0 else bbox1[0].detach().cpu().float()
        box[0::2] = box[0::2].clamp(0, cw - 1)
        box[1::2] = box[1::2].clamp(0, ch - 1)
        box = (box * torch.tensor([sx, sy, sx, sy])).tolist()
        box[0::2] = [min(box[0], rgb.shape[1] - 1), min(box[2], rgb.shape[1] - 1)]
        box[1::2] = [min(box[1], rgb.shape[0] - 1), min(box[3], rgb.shape[0] - 1)]

        outs.append((rgb, soft, box))
    return outs


@spaces.GPU(duration=30)
def match(image0, image1, mask_thresh=0.3, show_boxes=True, progress=gr.Progress(track_tqdm=True)):
    """Estimate the co-visible (matching) region of an image pair with SAMatcher.

    Args:
        image0: Reference view (RGB image).
        image1: Second view (RGB image; often a different zoom / scale).
        mask_thresh: Sigmoid threshold applied to the co-visibility mask (0.3 matches the
            paper's evaluation default; 0.4 is the value the tuned scripts use).
        show_boxes: Draw the predicted co-visible bounding boxes on both views.

    Returns:
        (overlay_view0, overlay_view1, crop0, crop1, info_text)
    """
    t0 = time.time()
    if image0 is None or image1 is None:
        raise gr.Error("Please provide both images of the pair.")
    if isinstance(image0, str):
        image0 = cv2.cvtColor(cv2.imread(image0, cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)
    if isinstance(image1, str):
        image1 = cv2.cvtColor(cv2.imread(image1, cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)

    (rgb0, soft0, box0), (rgb1, soft1, box1) = run_pair(image0, image1, mask_thresh)

    ov0 = render_overlay(rgb0, soft0, box0 if show_boxes else None)
    ov1 = render_overlay(rgb1, soft1, box1 if show_boxes else None)
    crop0 = render_crop_pair(rgb0, box0)
    crop1 = render_crop_pair(rgb1, box1)

    # Co-visibility summary (per-view fractions; the two masks live in different
    # coordinate frames, so a cross-view IoU is not meaningful)
    a0 = (soft0 > mask_thresh)
    a1 = (soft1 > mask_thresh)
    ratio0 = float(a0.mean())
    ratio1 = float(a1.mean())
    dt = time.time() - t0
    info = (f"Co-visible fraction — view A: {ratio0 * 100:.1f}% · view B: {ratio1 * 100:.1f}% · "
            f"inference: {dt:.2f}s")
    return ov0, ov1, crop0, crop1, info


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="SAMatcher") as demo:
    gr.Markdown(
        """
        # SAMatcher: Dense Co-Visibility Modeling via Cross-View Fusion for Scale-Imbalance Image Matching

        Upload (or pick) a pair of images — ideally two views with a strong **scale imbalance** — and
        SAMatcher predicts, in a single forward pass, the **co-visible region** shared by both views:
        a dense mask and a bounding box per view, plus the cropped matching windows a downstream
        matcher would run on.

        Model: [`SSSSphinx/SAMatcher`](https://huggingface.co/SSSSphinx/SAMatcher) ·
        Paper: [SAMatcher](https://huggingface.co/papers/2606.03406) ·
        Code: [TwSphinx54/SAMatcher](https://github.com/TwSphinx54/SAMatcher)
        """
    )
    with gr.Row():
        with gr.Column():
            image0_in = gr.Image(label="Image A (reference view)", type="numpy", height=320)
        with gr.Column():
            image1_in = gr.Image(label="Image B (scale-imbalanced view)", type="numpy", height=320)
    with gr.Row():
        run_btn = gr.Button("Find co-visible region", variant="primary")
    with gr.Accordion("Advanced options", open=False):
        mask_thresh = gr.Slider(0.1, 0.9, value=0.3, step=0.05,
                                label="Co-visibility mask threshold")
        show_boxes = gr.Checkbox(value=True, label="Draw predicted bounding boxes")
    with gr.Row():
        with gr.Column():
            out0 = gr.Image(label="Image A — predicted co-visible mask + box", type="numpy", height=340)
        with gr.Column():
            out1 = gr.Image(label="Image B — predicted co-visible mask + box", type="numpy", height=340)
    with gr.Row():
        with gr.Column():
            outcrop0 = gr.Image(label="Matching window A", type="numpy", height=260)
        with gr.Column():
            outcrop1 = gr.Image(label="Matching window B", type="numpy", height=260)
    info_out = gr.Markdown()

    run_btn.click(
        fn=match,
        inputs=[image0_in, image1_in, mask_thresh, show_boxes],
        outputs=[out0, out1, outcrop0, outcrop1, info_out],
        api_name="match",
        concurrency_limit=1,
    )

    gr.Examples(
        examples=[
            ["assets/sacre_coeur_A.jpg", "assets/sacre_coeur_B.jpg", 0.3, True],
            ["assets/toronto_A.jpg", "assets/toronto_B.jpg", 0.3, True],
            ["assets/desert_dunes.jpg", "assets/desert_dunes_closeup.jpg", 0.3, True],
        ],
        inputs=[image0_in, image1_in, mask_thresh, show_boxes],
        fn=match,
        outputs=[out0, out1, outcrop0, outcrop1, info_out],
        cache_examples=True,
        cache_mode="lazy",
        label="Example pairs (the last one is a truly scale-imbalanced pair)",
    )

    gr.Markdown(
        """
        ---
        **Method**: SAMatcher fuses the two views inside a SAM2.1-Hiera-L encoder with a symmetric
        cross-view fusion module (view-aware RoPE) and decodes a shared learnable prompt into
        per-view co-visibility masks and boxes — no matching, no correspondence supervision needed.
        Example images: Sacré-Cœur & Toronto pairs from [RoMa](https://github.com/Parskatt/RoMa)
        (MIT), desert pair from [linoyts/repo-to-space-example-inputs](https://huggingface.co/datasets/linoyts/repo-to-space-example-inputs) (CC0).
        """
    )

custom_css = """
.dark .gradio-container {
    background: var(--background-fill-primary);
    color: var(--body-text-color);
}
"""

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=custom_css)