"""
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. REMOVED and ADDED objects are matched against each other to detect
     MOVED objects, using a globally-greedy bipartite match over DINOv2
     similarity (see _match_moved_objects)
  6. Draw color-coded bounding boxes from the SAM masks

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 usable as fallback output

All models are lazy-loaded. Falls back to classical pipeline (diff.py) if any
required model is unavailable — see app.py for the wiring.
"""
import cv2
import numpy as np
from skimage.metrics import structural_similarity as ssim

# =====================================================================
# CONFIG / 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
MAX_LONG_EDGE = 960

DINO_SIM_THRESHOLD = 0.92   # same-location patches: >= this => unchanged, drop
MOVED_SIM_THRESHOLD = 0.75  # cross-location patches: >= this => same object, moved
EDGE_DENSITY_THRESHOLD = 1.0

REGION_PAD = 12       # padding when cropping a change region for verification
MOVE_MATCH_PAD = 12   # padding when cropping removed/added objects for move-matching


# =====================================================================
# LAZY MODEL LOADERS
# =====================================================================

_loftr = None

def _get_loftr():
    """Lazy-load kornia LoFTR (transformer dense matcher)."""
    global _loftr
    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
    try:
        from ultralytics import SAM
        _sam = SAM("sam2_t.pt")  # ~150MB, auto-downloads
        print("[diff_ai] SAM2 loaded")
    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
    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))
    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), scale


def align_loftr(img1_bgr, img2_bgr):
    """Warp img2 onto img1 using LoFTR matches + RANSAC homography.

    Returns (aligned_bgr, valid_mask) on success, or (None, None) on
    failure. valid_mask is a uint8 mask (255 = valid, 0 = invalid) the
    same size as img1, marking which pixels of `aligned_bgr` actually
    came from img2 vs. the black border warpPerspective fills in when
    the two images aren't framed identically. Callers should restrict
    any pixel-level comparison to valid_mask == 255, otherwise that
    black border gets treated as a huge fake "removed" region."""
    bundle = _get_loftr()
    if not bundle:
        return None, None
    matcher = bundle["matcher"]
    torch   = bundle["torch"]
    device  = bundle["device"]

    try:
        t1, s1 = _to_loftr_tensor(img1_bgr, torch, device)
        t2, s2 = _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 > 0.5
        kp1 = kp1[mask] / s1
        kp2 = kp2[mask] / s2

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

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

        out_size = (img1_bgr.shape[1], img1_bgr.shape[0])
        aligned = cv2.warpPerspective(img2_bgr, H, out_size)

        # Warp an all-white mask through the same H to find which pixels
        # of `aligned` are real img2 content vs. the black fill border.
        valid = np.full(img2_bgr.shape[:2], 255, dtype=np.uint8)
        valid_mask = cv2.warpPerspective(valid, H, out_size)
        # Erode a bit so interpolation-blurred edge pixels near the
        # border aren't counted as valid either.
        valid_mask = cv2.erode(valid_mask, np.ones((9, 9), np.uint8))

        return aligned, valid_mask
    except Exception as e:
        print(f"[diff_ai] LoFTR alignment error: {e}")
        return None, None


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

def extract_change_regions(img1_bgr, aligned_bgr, valid_mask=None):
    """Returns (heatmap_bgr, list of {bbox, area, centroid}, binmask,
    alignment_failed). Each region is a connected blob of pixels that
    significantly differ.

    valid_mask (optional): uint8 mask from align_loftr marking which
    pixels of aligned_bgr are real warped content vs. warpPerspective's
    black fill border. When provided, the border is excluded from the
    change mask entirely so it can never be picked up as a fake
    "removed" region, and low valid coverage also counts toward
    alignment_failed."""
    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)

    img_area = img1_bgr.shape[0] * img1_bgr.shape[1]

    if valid_mask is not None:
        # Never let the warpPerspective border be treated as a change
        binmask = cv2.bitwise_and(binmask, valid_mask)
        # If a large chunk of the frame has no real img2 coverage at all,
        # the homography barely overlaps the frame — treat as a failure
        # the same way a too-large diff_ratio would be.
        coverage_ratio = cv2.countNonZero(valid_mask) / img_area
        if coverage_ratio < 0.5:
            print(f"[diff_ai] low valid coverage after warp ({coverage_ratio:.2f}) — "
                  f"treating as alignment failure")
            return heatmap, [], binmask, True

    # Global alignment-failure guard
    diff_ratio = cv2.countNonZero(binmask) / img_area
    alignment_failed = diff_ratio > 0.55

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

    min_area = max(400, int(img_area * 0.0005))   # >= 0.05% of image
    max_area = int(img_area * 0.6)                # < 60% (skip background)

    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
# =====================================================================

def dinov2_similarity(patch1_bgr, patch2_bgr):
    """Returns cosine similarity in [0, 1] between two patches.
    Higher = more perceptually similar. Used to drop change regions caused
    by lighting/shadow rather than real object changes, and to match
    removed/added objects for MOVED detection."""
    bundle = _get_dinov2()
    if not bundle:
        return None
    model     = bundle["model"]
    processor = bundle["processor"]
    torch     = bundle["torch"]
    device    = bundle["device"]

    if patch1_bgr.size == 0 or patch2_bgr.size == 0:
        return None
    try:
        rgb1 = cv2.cvtColor(patch1_bgr, cv2.COLOR_BGR2RGB)
        rgb2 = cv2.cvtColor(patch2_bgr, cv2.COLOR_BGR2RGB)
        inputs = processor(images=[rgb1, rgb2], return_tensors="pt").to(device)
        with torch.no_grad():
            out = model(**inputs)
        feats = out.last_hidden_state[:, 0]            # CLS token, (2, D)
        feats = feats / feats.norm(dim=1, keepdim=True)
        sim = float((feats[0] * feats[1]).sum().item())
        return sim
    except Exception as e:
        print(f"[diff_ai] DINOv2 sim error: {e}")
        return None


# =====================================================================
# 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."""
    sam = _get_sam()
    if not sam:
        return None
    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:
        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)."""
    mask = _sam_prompt_single(img_bgr, point_xy)
    if mask is not None and int(mask.sum()) > 100:
        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


def bbox_from_mask(mask):
    """Tight bbox (x, y, w, h) around a binary mask, or None if empty."""
    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_DENSITY_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"


# =====================================================================
# STAGE 6 — MOVED-OBJECT MATCHING (removed <-> added)
# =====================================================================

def _crop_patch(img_bgr, bbox, pad):
    """Crop bbox from img_bgr with padding, clipped to image bounds."""
    x, y, w, h = bbox
    x0 = max(0, x - pad)
    y0 = max(0, y - pad)
    x1 = min(img_bgr.shape[1], x + w + pad)
    y1 = min(img_bgr.shape[0], y + h + pad)
    return img_bgr[y0:y1, x0:x1]


def _bboxes_overlap(bbox_a, bbox_b):
    """AABB intersection test. Overlapping removed/added regions are
    treated as an in-place change rather than a move, so they're never
    considered as a MOVED pair."""
    ax, ay, aw, ah = bbox_a
    bx, by, bw, bh = bbox_b
    return ax < bx + bw and ax + aw > bx and ay < by + bh and ay + ah > by


def _match_moved_objects(removed_objects, added_objects, img1_bgr, aligned_bgr,
                          pad=MOVE_MATCH_PAD, sim_threshold=MOVED_SIM_THRESHOLD):
    """Match REMOVED objects against ADDED objects to find MOVED objects.

    Fix vs. the previous implementation: the old version iterated
    `removed_objects` in (arbitrary, scan-order-derived) list order and let
    each removed object greedily grab its own single best-matching added
    object before moving to the next. That's order-dependent — an earlier
    removed object could claim a mediocre match and starve a later removed
    object of the added object it should have matched to, and swapping the
    order of `removed_objects` could change the result.

    This version computes the *entire* removed x added similarity matrix
    up front (skipping overlapping pairs), then commits pairs in
    descending similarity order — the single best candidate pair anywhere
    in the matrix is matched first, then the next best remaining pair, and
    so on. This is order-independent and much closer to an optimal
    bipartite matching than the old per-item greedy approach.

    Returns (moved_objects, surviving_removed, surviving_added).
    """
    if not removed_objects or not added_objects:
        return [], list(removed_objects), list(added_objects)

    candidates = []  # (similarity, removed_idx, added_idx)
    for ri, rem in enumerate(removed_objects):
        rem_patch = _crop_patch(img1_bgr, rem["bbox"], pad)
        for ai, add in enumerate(added_objects):
            if _bboxes_overlap(rem["bbox"], add["bbox"]):
                continue
            add_patch = _crop_patch(aligned_bgr, add["bbox"], pad)
            sim = dinov2_similarity(rem_patch, add_patch)
            if sim is not None and sim >= sim_threshold:
                candidates.append((sim, ri, ai))

    # Best pairs first, order-independent
    candidates.sort(key=lambda c: c[0], reverse=True)

    matched_removed = set()
    matched_added = set()
    moved_objects = []
    for sim, ri, ai in candidates:
        if ri in matched_removed or ai in matched_added:
            continue
        matched_removed.add(ri)
        matched_added.add(ai)
        moved_objects.append({
            "from": removed_objects[ri],
            "to": added_objects[ai],
            "similarity": sim,
        })
        print(f"  MOVED: from {removed_objects[ri]['centroid']} "
              f"-> to {added_objects[ai]['centroid']}  sim={sim:.3f}")

    surviving_removed = [r for i, r in enumerate(removed_objects) if i not in matched_removed]
    surviving_added   = [a for i, a in enumerate(added_objects)   if i not in matched_added]

    for i, rem in enumerate(removed_objects):
        if i not in matched_removed:
            print(f"  REMOVED (unmatched) @{rem['centroid']}")

    return moved_objects, surviving_removed, surviving_added


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

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


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

def compare_images_ai(img1_path, img2_path, output_path, heatmap_path):
    """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")

    # Normalize sizes + downscale to keep models fast
    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, valid_mask = align_loftr(img1, img2)
    if aligned is None:
        raise RuntimeError("LoFTR alignment failed")

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

    # 3) For each region: DINOv2 verify + SAM2 prompted segmentation
    added_objects = []
    removed_objects = []

    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)
        patch1 = img1[y0:y1, x0:x1]
        patch2 = aligned[y0:y1, x0:x1]

        # 3a) DINOv2 perceptual gate
        sim = dinov2_similarity(patch1, patch2)
        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) if mask is not None else 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}")

    # 4) MOVED object detection: globally-greedy match REMOVED <-> ADDED
    moved_objects, removed_objects, added_objects = _match_moved_objects(
        removed_objects, added_objects, img1, aligned
    )

    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) Draw boxes
    result_img = aligned.copy()

    for a in added_objects:
        _draw_box(result_img, a["bbox"], COLOR_ADDED, "ADDED")
    for r in removed_objects:
        _draw_box(result_img, r["bbox"], COLOR_REMOVED, "REMOVED")
    for m in moved_objects:
        _draw_box(result_img, m["from"]["bbox"], COLOR_MOVED, "MOVED")
        _draw_box(result_img, m["to"]["bbox"],   COLOR_MOVED, "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)

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