"""
AI-first image comparison pipeline (change-region-driven).

Strategy (different from naive segment-everything-then-match):
  1. LoFTR (kornia) aligns img2 onto img1
  2. SSIM produces a binary "change mask" — where pixels actually differ
  3. Connected components on the change mask give us discrete change regions
  4. For each region:
       a. Crop a context patch from both images
       b. DINOv2 verifies perceptual difference (kills shadow/lighting noise)
       c. SAM2 in PROMPT mode (point prompt at the region centroid) returns
          a tight object mask in img2
       d. Edge-density heuristic classifies as ADDED vs REMOVED
  5. Draw color-coded bounding boxes / masks (falling back to the raw change
     bbox if SAM couldn't produce a mask, so a single failed SAM call never
     causes a real change to vanish from the output)

Why this is better than the previous attempt:
  - No SAM auto-mask-generation (10x faster — only prompted on change regions)
  - No cross-image DINO matching (no phantom added/removed pairs)
  - Change mask is the source of truth; SAM and DINO act as refinement layers
  - Pipeline degrades gracefully: if SAM/DINO fail, the change-region bboxes
    are still drawn as fallback output (not silently dropped)
  - DINOv2 embeddings are batched: the perceptual gate embeds all regions in
    one forward pass, and moved-object matching embeds each patch once and
    compares via a similarity matrix instead of one forward pass per pair

All models are lazy-loaded (thread-safely). Falls back to classical pipeline
(diff.py) if any required model is unavailable — see app.py for the wiring.
"""
import threading
try:
    import spaces
    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False

import cv2
import numpy as np
from skimage.metrics import structural_similarity as ssim

# =====================================================================
# TUNABLE CONSTANTS
# =====================================================================

# Drawing colors (BGR)
COLOR_ADDED   = (0, 200,   0)
COLOR_REMOVED = (0,   0, 220)
COLOR_MOVED   = (255, 255,  0)   # cyan in BGR  (drawn on both from & to boxes)

LOFTR_LONG_EDGE   = 640          # working resolution for LoFTR matching
LOFTR_MIN_MATCHES = 10           # below this, alignment is considered failed
LOFTR_CONF_THRESH = 0.5

MAX_LONG_EDGE = 960              # cap on working image size for the pipeline

CHANGE_MIN_AREA_FRAC = 0.0005    # region must be >= this fraction of image area
CHANGE_MAX_AREA_FRAC = 0.6       # region must be < this fraction (skip background)
ALIGNMENT_FAIL_DIFF_RATIO = 0.55 # if more of the image "changed" than this, alignment likely failed

DINO_SIM_THRESHOLD  = 0.85       # >= this means perceptually identical → drop region
MOVED_SIM_THRESHOLD = 0.75       # >= this cosine sim links a removed/added pair as "moved"
REGION_PAD = 12                  # px padding around a region when cropping context patches

EDGE_VARIANCE_THRESHOLD = 1.0    # Laplacian-variance floor for added/removed classification

# SAM mask is considered too small/noisy below this fraction of the prompt
# image's area (relative, so it scales with input resolution).
SAM_MIN_MASK_AREA_FRAC = 0.0005

_MODEL_LOCK = threading.Lock()


# =====================================================================
# LAZY MODEL LOADERS (thread-safe)
# =====================================================================

_loftr = None

def _get_loftr():
    """Lazy-load kornia LoFTR (transformer dense matcher)."""
    global _loftr
    if _loftr is not None:
        return _loftr
    with _MODEL_LOCK:
        if _loftr is not None:
            return _loftr
        try:
            import torch
            import kornia.feature as KF
            device = "cuda" if torch.cuda.is_available() else "cpu"
            matcher = KF.LoFTR(pretrained="outdoor").eval().to(device)
            _loftr = {"matcher": matcher, "torch": torch, "device": device}
            print("[diff_ai] LoFTR loaded")
        except Exception as e:
            print(f"[diff_ai] LoFTR unavailable: {e}")
            _loftr = False
    return _loftr


_sam = None

def _get_sam():
    """Lazy-load ultralytics SAM2 (used in prompted mode, not auto-mask)."""
    global _sam
    if _sam is not None:
        return _sam
    with _MODEL_LOCK:
        if _sam is not None:
            return _sam
        try:
            import torch
            from ultralytics import SAM
            model = SAM("sam2_t.pt")  # ~150MB, auto-downloads
            device = "cuda" if torch.cuda.is_available() else "cpu"
            try:
                model.to(device)
            except Exception:
                pass  # some ultralytics versions take device per-call instead
            _sam = {"model": model, "device": device}
            print(f"[diff_ai] SAM2 loaded on {device}")
        except Exception as e:
            print(f"[diff_ai] SAM2 unavailable: {e}")
            _sam = False
    return _sam


_dinov2 = None

def _get_dinov2():
    """Lazy-load DINOv2 via HuggingFace transformers (Python 3.9 compatible)."""
    global _dinov2
    if _dinov2 is not None:
        return _dinov2
    with _MODEL_LOCK:
        if _dinov2 is not None:
            return _dinov2
        try:
            import torch
            from transformers import AutoModel, AutoImageProcessor
            device = "cuda" if torch.cuda.is_available() else "cpu"
            processor = AutoImageProcessor.from_pretrained(
                "facebook/dinov2-small", use_fast=True
            )
            model = AutoModel.from_pretrained("facebook/dinov2-small").eval().to(device)
            _dinov2 = {
                "model": model, "processor": processor,
                "torch": torch, "device": device,
            }
            print("[diff_ai] DINOv2 loaded")
        except Exception as e:
            print(f"[diff_ai] DINOv2 unavailable: {e}")
            _dinov2 = False
    return _dinov2


# =====================================================================
# STAGE 1 — LoFTR ALIGNMENT
# =====================================================================

def _to_loftr_tensor(bgr, torch_mod, device):
    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
    h, w = gray.shape
    scale = LOFTR_LONG_EDGE / max(h, w)
    nh, nw = int(round(h * scale)), int(round(w * scale))
    # Round down to a multiple of 8 (LoFTR requirement) without ever hitting 0.
    nh = max(8, nh - nh % 8)
    nw = max(8, nw - nw % 8)
    resized = cv2.resize(gray, (nw, nh))
    t = torch_mod.from_numpy(resized).float()[None, None] / 255.0
    return t.to(device), (nw / w, nh / h)  # per-axis scale, resize isn't isotropic-safe


def align_loftr(img1_bgr, img2_bgr):
    """Warp img2 onto img1 using LoFTR matches + RANSAC homography."""
    bundle = _get_loftr()
    if not bundle:
        return None
    matcher = bundle["matcher"]
    torch   = bundle["torch"]
    device  = bundle["device"]

    try:
        t1, (sx1, sy1) = _to_loftr_tensor(img1_bgr, torch, device)
        t2, (sx2, sy2) = _to_loftr_tensor(img2_bgr, torch, device)
        with torch.no_grad():
            corr = matcher({"image0": t1, "image1": t2})

        kp1 = corr["keypoints0"].cpu().numpy()
        kp2 = corr["keypoints1"].cpu().numpy()
        confidence = corr["confidence"].cpu().numpy()

        mask = confidence > LOFTR_CONF_THRESH
        kp1 = kp1[mask] / np.array([sx1, sy1])
        kp2 = kp2[mask] / np.array([sx2, sy2])

        if len(kp1) < LOFTR_MIN_MATCHES:
            print(f"[diff_ai] LoFTR: only {len(kp1)} confident matches")
            return None

        H, _ = cv2.findHomography(kp2, kp1, cv2.RANSAC, 5.0)
        if H is None:
            return None

        return cv2.warpPerspective(
            img2_bgr, H, (img1_bgr.shape[1], img1_bgr.shape[0])
        )
    except Exception as e:
        print(f"[diff_ai] LoFTR alignment error: {e}")
        return None


# =====================================================================
# STAGE 2 — CHANGE REGION EXTRACTION (SSIM + connected components)
# =====================================================================

def extract_change_regions(img1_bgr, aligned_bgr):
    """Returns (heatmap_bgr, list of {bbox, area, centroid}, binmask, alignment_failed).
    Each region is a connected blob of pixels that significantly differ."""
    g1 = cv2.cvtColor(img1_bgr,   cv2.COLOR_BGR2GRAY)
    g2 = cv2.cvtColor(aligned_bgr, cv2.COLOR_BGR2GRAY)
    g1 = cv2.GaussianBlur(g1, (5, 5), 0)
    g2 = cv2.GaussianBlur(g2, (5, 5), 0)

    _, diff = ssim(g1, g2, full=True)
    diff_u8 = np.clip((1.0 - diff) * 255.0, 0, 255).astype(np.uint8)
    heatmap = cv2.applyColorMap(diff_u8, cv2.COLORMAP_JET)

    # Otsu binarize + clean noise
    _, binmask = cv2.threshold(diff_u8, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    binmask = cv2.medianBlur(binmask, 7)
    open_k  = np.ones((5, 5), np.uint8)
    close_k = np.ones((15, 15), np.uint8)
    binmask = cv2.morphologyEx(binmask, cv2.MORPH_OPEN,  open_k,  iterations=1)
    binmask = cv2.morphologyEx(binmask, cv2.MORPH_CLOSE, close_k, iterations=2)

    # Global alignment-failure guard
    img_area = img1_bgr.shape[0] * img1_bgr.shape[1]
    diff_ratio = cv2.countNonZero(binmask) / img_area
    alignment_failed = diff_ratio > ALIGNMENT_FAIL_DIFF_RATIO

    # Connected components → individual change regions
    n_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
        binmask, connectivity=8
    )

    min_area = max(400, int(img_area * CHANGE_MIN_AREA_FRAC))
    max_area = int(img_area * CHANGE_MAX_AREA_FRAC)

    regions = []
    if not alignment_failed:
        for i in range(1, n_labels):  # skip label 0 (background)
            x = int(stats[i, cv2.CC_STAT_LEFT])
            y = int(stats[i, cv2.CC_STAT_TOP])
            w = int(stats[i, cv2.CC_STAT_WIDTH])
            h = int(stats[i, cv2.CC_STAT_HEIGHT])
            area = int(stats[i, cv2.CC_STAT_AREA])
            if area < min_area or area > max_area:
                continue
            cx = float(centroids[i, 0])
            cy = float(centroids[i, 1])
            regions.append({
                "bbox": (x, y, w, h),
                "area": area,
                "centroid": (cx, cy),
            })

    return heatmap, regions, binmask, alignment_failed


# =====================================================================
# STAGE 3 — DINOv2 PERCEPTUAL VERIFICATION (batched)
# =====================================================================

def dinov2_embed_batch(patches_bgr):
    """Embed a list of BGR patches in a single forward pass.
    Returns an (N, D) L2-normalized numpy array, or None if the model is
    unavailable or every patch is empty. Empty patches get an all-zero
    embedding (guaranteed cosine sim of 0 with anything)."""
    bundle = _get_dinov2()
    if not bundle:
        return None
    model, processor = bundle["model"], bundle["processor"]
    torch, device = bundle["torch"], bundle["device"]

    valid_idx = [i for i, p in enumerate(patches_bgr) if p is not None and p.size > 0]
    if not valid_idx:
        return None

    try:
        rgb_imgs = [cv2.cvtColor(patches_bgr[i], cv2.COLOR_BGR2RGB) for i in valid_idx]
        inputs = processor(images=rgb_imgs, return_tensors="pt").to(device)
        with torch.no_grad():
            out = model(**inputs)
        feats = out.last_hidden_state[:, 0]  # CLS token, (n_valid, D)
        feats = feats / feats.norm(dim=1, keepdim=True)
        feats = feats.cpu().numpy()
    except Exception as e:
        print(f"[diff_ai] DINOv2 batch embed error: {e}")
        return None

    dim = feats.shape[1]
    result = np.zeros((len(patches_bgr), dim), dtype=feats.dtype)
    for slot, i in enumerate(valid_idx):
        result[i] = feats[slot]
    return result


def dinov2_similarity(patch1_bgr, patch2_bgr):
    """Cosine similarity in [-1, 1] between two patches (single-pair
    convenience wrapper around the batched embedder). Higher = more
    perceptually similar."""
    embs = dinov2_embed_batch([patch1_bgr, patch2_bgr])
    if embs is None:
        return None
    if not np.any(embs[0]) or not np.any(embs[1]):
        return None
    return float(np.dot(embs[0], embs[1]))


# =====================================================================
# STAGE 4 — SAM2 PROMPTED SEGMENTATION
# =====================================================================

def _sam_prompt_single(img_bgr, point_xy):
    """Run SAM2 with one point prompt. Returns binary uint8 mask or None."""
    bundle = _get_sam()
    if not bundle:
        return None
    sam = bundle["model"]
    try:
        results = sam(
            img_bgr,
            points=[[float(point_xy[0]), float(point_xy[1])]],
            labels=[1],
            verbose=False,
        )
        if not results:
            return None
        r = results[0]
        if r.masks is None or len(r.masks.data) == 0:
            return None
        masks_np = r.masks.data.cpu().numpy()
        areas = [int((m > 0.5).sum()) for m in masks_np]
        if not areas:
            return None
        best_idx = int(np.argmax(areas))
        return (masks_np[best_idx] > 0.5).astype(np.uint8)
    except Exception as e:
        print(f"[diff_ai] SAM2 prompt error @{point_xy}: {e}")
        return None


def sam_mask_at_point(img_bgr, point_xy, fallback_bbox=None):
    """Run SAM2 with a single positive point prompt. Returns the largest
    returned mask as a binary uint8 array, or None on failure.

    If the centroid prompt produces an empty/noisy mask and fallback_bbox
    is provided, tries additional points across the region and picks the
    largest mask (multi-point fallback for edge cases where the centroid
    lands on background — e.g. ring-shaped change regions)."""
    img_area = img_bgr.shape[0] * img_bgr.shape[1]
    min_mask_area = max(50, int(img_area * SAM_MIN_MASK_AREA_FRAC))

    mask = _sam_prompt_single(img_bgr, point_xy)
    if mask is not None and int(mask.sum()) > min_mask_area:
        return mask

    if fallback_bbox is not None:
        x, y, w, h = fallback_bbox
        margin = 0.25
        candidates = [
            (x + w * margin, y + h * margin),
            (x + w * (1.0 - margin), y + h * margin),
            (x + w * margin, y + h * (1.0 - margin)),
            (x + w * (1.0 - margin), y + h * (1.0 - margin)),
            (x + w * 0.5, y + h * 0.5),
        ]
        best_mask = None
        best_area = 0
        for px, py in candidates:
            m = _sam_prompt_single(img_bgr, (px, py))
            if m is not None:
                a = int(m.sum())
                if a > best_area:
                    best_area = a
                    best_mask = m
        if best_mask is not None:
            return best_mask

    return mask  # possibly None, possibly small — caller falls back to bbox


def bbox_from_mask(mask):
    """Tight bbox (x, y, w, h) around a binary mask, or None if empty."""
    if mask is None:
        return None
    ys, xs = np.where(mask > 0)
    if len(ys) == 0:
        return None
    x0, x1 = int(xs.min()), int(xs.max())
    y0, y1 = int(ys.min()), int(ys.max())
    return (x0, y0, x1 - x0 + 1, y1 - y0 + 1)


# =====================================================================
# STAGE 5 — ADDED vs REMOVED CLASSIFICATION
# =====================================================================

def classify_added_or_removed(patch1_bgr, patch2_bgr,
                               edge_threshold=EDGE_VARIANCE_THRESHOLD):
    """Edge-density heuristic. The image with more edges in this region
    is the one that has the 'object'.
    Returns 'added' (img2 has more), 'removed' (img1 has more),
    or None if both patches have near-zero edge variance (noise)."""
    g1 = cv2.cvtColor(patch1_bgr, cv2.COLOR_BGR2GRAY)
    g2 = cv2.cvtColor(patch2_bgr, cv2.COLOR_BGR2GRAY)
    e1 = float(cv2.Laplacian(g1, cv2.CV_64F).var())
    e2 = float(cv2.Laplacian(g2, cv2.CV_64F).var())
    if e1 < edge_threshold and e2 < edge_threshold:
        return None
    return "added" if e2 >= e1 else "removed"


# =====================================================================
# DRAWING
# =====================================================================

def _draw_box(img, bbox, color, label=None):
    x, y, w, h = bbox
    cv2.rectangle(img, (x, y), (x + w, y + h), color, 3)
    if label:
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)
        ty = max(th + 8, y)
        cv2.rectangle(img, (x, ty - th - 8), (x + tw + 8, ty), color, -1)
        cv2.putText(img, label, (x + 4, ty - 4),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)


def _composite_object(base, object_img, mask, bbox, color):
    """Paste object pixels from object_img onto a darkened base using mask,
    with a colored outline. If mask is missing/empty (SAM failed or was
    unavailable), fall back to drawing the raw change-region bbox so the
    change is never silently dropped from the output. Modifies base in place."""
    if mask is not None and int(mask.sum()) > 0:
        mask_bin = mask.astype(bool)
        base[mask_bin] = object_img[mask_bin]
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(base, contours, -1, color, 2)
    elif bbox is not None:
        x, y, w, h = bbox
        base[y:y + h, x:x + w] = object_img[y:y + h, x:x + w]
        _draw_box(base, bbox, color)
    return base


# =====================================================================
# MAIN ENTRY
# =====================================================================

def _compare_images_ai_impl(img1_path, img2_path, output_path, heatmap_path, output_box_path=None):
    """End-to-end change-region-driven AI pipeline. Raises RuntimeError
    if any required model is unavailable."""
    if not _get_loftr() or not _get_sam() or not _get_dinov2():
        raise RuntimeError("AI pipeline unavailable — required model missing")

    img1 = cv2.imread(img1_path)
    img2 = cv2.imread(img2_path)
    if img1 is None or img2 is None:
        raise RuntimeError("Failed to read input images")

    if img1.shape[:2] != img2.shape[:2]:
        print(f"[diff_ai] warning: input sizes differ {img1.shape[:2]} vs "
              f"{img2.shape[:2]}; resizing img2 to match img1 (may distort)")
        img2 = cv2.resize(img2, (img1.shape[1], img1.shape[0]))

    h0, w0 = img1.shape[:2]
    if max(h0, w0) > MAX_LONG_EDGE:
        s = MAX_LONG_EDGE / max(h0, w0)
        new_size = (int(w0 * s), int(h0 * s))
        img1 = cv2.resize(img1, new_size, interpolation=cv2.INTER_AREA)
        img2 = cv2.resize(img2, new_size, interpolation=cv2.INTER_AREA)

    # 1) LoFTR alignment
    aligned = align_loftr(img1, img2)
    if aligned is None:
        raise RuntimeError("LoFTR alignment failed")

    # 2) Change region extraction (SSIM + CC)
    heatmap, regions, _, alignment_failed = extract_change_regions(img1, aligned)
    if alignment_failed:
        regions = []
    print(f"[diff_ai] {len(regions)} change region(s) before verification")

    # 3) Batched DINOv2 perceptual gate across all regions at once
    patches1, patches2 = [], []
    for r in regions:
        x, y, w, h = r["bbox"]
        x0 = max(0, x - REGION_PAD); y0 = max(0, y - REGION_PAD)
        x1 = min(img1.shape[1], x + w + REGION_PAD)
        y1 = min(img1.shape[0], y + h + REGION_PAD)
        r["patch_bounds"] = (x0, y0, x1, y1)
        patches1.append(img1[y0:y1, x0:x1])
        patches2.append(aligned[y0:y1, x0:x1])

    emb1 = dinov2_embed_batch(patches1)
    emb2 = dinov2_embed_batch(patches2)

    added_objects = []
    removed_objects = []

    for i, r in enumerate(regions):
        x0, y0, x1, y1 = r["patch_bounds"]
        patch1, patch2 = patches1[i], patches2[i]

        # 3a) DINOv2 perceptual gate
        sim = None
        if emb1 is not None and emb2 is not None and np.any(emb1[i]) and np.any(emb2[i]):
            sim = float(np.dot(emb1[i], emb2[i]))
        if sim is not None and sim >= DINO_SIM_THRESHOLD:
            print(f"  region @{r['centroid']}: DINOv2 sim={sim:.3f} — perceptually same, dropping")
            continue

        # 3b) Classify as added vs removed via edge density
        label = classify_added_or_removed(patch1, patch2)
        if label is None:
            print(f"  region @{r['centroid']}: edge variance too low on both sides, skipping")
            continue

        # 3c) SAM2 prompted at the centroid for a clean object mask
        target_img = aligned if label == "added" else img1
        mask = sam_mask_at_point(target_img, r["centroid"], r["bbox"])
        tight_bbox = bbox_from_mask(mask) or r["bbox"]

        entry = {
            "bbox": tight_bbox,
            "centroid": r["centroid"],
            "dino_sim": sim,
            "mask": mask,
            "source": target_img,
        }
        if label == "added":
            added_objects.append(entry)
        else:
            removed_objects.append(entry)

        print(f"  region @{r['centroid']}: sim={sim} -> {label}"
              f"{' (no SAM mask, using bbox fallback)' if mask is None else ''}")

    # 4) MOVED object detection: match REMOVED <-> ADDED pairs.
    # Embed every candidate patch once, then compare via a similarity matrix
    # instead of one DINOv2 forward pass per (removed, added) pair.
    rem_patches, add_patches = [], []
    for rem in removed_objects:
        rx, ry, rw, rh = rem["bbox"]
        r_x0 = max(0, rx - REGION_PAD); r_y0 = max(0, ry - REGION_PAD)
        r_x1 = min(img1.shape[1], rx + rw + REGION_PAD)
        r_y1 = min(img1.shape[0], ry + rh + REGION_PAD)
        rem_patches.append(img1[r_y0:r_y1, r_x0:r_x1])
    for add in added_objects:
        ax, ay, aw, ah = add["bbox"]
        a_x0 = max(0, ax - REGION_PAD); a_y0 = max(0, ay - REGION_PAD)
        a_x1 = min(aligned.shape[1], ax + aw + REGION_PAD)
        a_y1 = min(aligned.shape[0], ay + ah + REGION_PAD)
        add_patches.append(aligned[a_y0:a_y1, a_x0:a_x1])

    rem_emb = dinov2_embed_batch(rem_patches) if rem_patches else None
    add_emb = dinov2_embed_batch(add_patches) if add_patches else None

    moved_objects = []
    surviving_removed = []
    remaining_added = list(range(len(added_objects)))

    if rem_emb is not None and add_emb is not None and len(remaining_added) > 0:
        sim_matrix = rem_emb @ add_emb.T  # (R, A), both rows L2-normalized
    else:
        sim_matrix = None

    for ri, rem in enumerate(removed_objects):
        rx, ry, rw, rh = rem["bbox"]
        best_sim, best_j = -1.0, None
        if sim_matrix is not None:
            for j in remaining_added:
                add = added_objects[j]
                ax, ay, aw, ah = add["bbox"]
                # Skip if bboxes overlap (different objects at same position)
                if rx < ax + aw and rx + rw > ax and ry < ay + ah and ry + rh > ay:
                    continue
                sim = float(sim_matrix[ri, j])
                if sim > best_sim:
                    best_sim, best_j = sim, j

        if best_j is not None and best_sim >= MOVED_SIM_THRESHOLD:
            add = added_objects[best_j]
            moved_objects.append({"from": rem, "to": add, "similarity": best_sim})
            remaining_added.remove(best_j)
            print(f"  MOVED: from {rem['centroid']} -> to {add['centroid']}  sim={best_sim:.3f}")
        else:
            surviving_removed.append(rem)
            if best_j is not None:
                print(f"  REMOVED (unmatched) @{rem['centroid']}: best sim={best_sim:.3f} < {MOVED_SIM_THRESHOLD}")
            else:
                print(f"  REMOVED (unmatched) @{rem['centroid']}: no non-overlapping ADDED found")

    removed_objects = surviving_removed
    added_objects = [added_objects[j] for j in remaining_added]

    n_moved   = len(moved_objects)
    n_added   = len(added_objects)
    n_removed = len(removed_objects)
    n_total   = n_added + n_removed + n_moved
    severity  = "HIGH" if n_total > 0 else "NONE"

    # 5) Render output image - Segment Overlay (overlay)
    DARK_FACTOR = 0.3
    result_img = cv2.multiply(aligned, np.array([DARK_FACTOR] * 3, dtype=np.float64))
    result_img = np.clip(result_img, 0, 255).astype(np.uint8)

    for a in added_objects:
        _composite_object(result_img, aligned, a["mask"], a["bbox"], COLOR_ADDED)
    for r in removed_objects:
        _composite_object(result_img, r["source"], r["mask"], r["bbox"], COLOR_REMOVED)
    for m in moved_objects:
        _composite_object(result_img, m["from"]["source"], m["from"]["mask"], m["from"]["bbox"], COLOR_MOVED)
        _composite_object(result_img, aligned,              m["to"]["mask"],   m["to"]["bbox"],   COLOR_MOVED)

    # Draw severity label in top-left
    label = f"{severity}  +{n_added}  -{n_removed}  ~{n_moved}"
    cv2.putText(result_img, label, (20, 42),
                cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), 5)
    cv2.putText(result_img, label, (20, 42),
                cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0), 2)

    cv2.imwrite(output_path, result_img)
    cv2.imwrite(heatmap_path, heatmap)

    # 6) Render output image - Bounding Box (box)
    if output_box_path:
        box_img = aligned.copy()
        for a in added_objects:
            _draw_box(box_img, a["bbox"], COLOR_ADDED, "ADDED")
        for r in removed_objects:
            _draw_box(box_img, r["bbox"], COLOR_REMOVED, "REMOVED")
        for m in moved_objects:
            _draw_box(box_img, m["from"]["bbox"], COLOR_MOVED, "MOVED")
            _draw_box(box_img, m["to"]["bbox"],   COLOR_MOVED, "MOVED")

        cv2.putText(box_img, label, (20, 42),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), 5)
        cv2.putText(box_img, label, (20, 42),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 0), 2)
        cv2.imwrite(output_box_path, box_img)

    return {
        "pipeline": "ai",
        "added": n_added,
        "removed": n_removed,
        "moved": n_moved,
        "object_changes": n_total,
        "severity": severity,
        "alignment_failed": alignment_failed,
    }


if HAS_SPACES:
    @spaces.GPU
    def compare_images_ai(img1_path, img2_path, output_path, heatmap_path, output_box_path=None):
        return _compare_images_ai_impl(img1_path, img2_path, output_path, heatmap_path, output_box_path=output_box_path)
else:
    def compare_images_ai(img1_path, img2_path, output_path, heatmap_path, output_box_path=None):
        return _compare_images_ai_impl(img1_path, img2_path, output_path, heatmap_path, output_box_path=output_box_path)