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

# -----------------------------
# LPIPS — perceptual similarity (lazy-loaded so first request pays the cost)
# -----------------------------
_lpips_model = None
_lpips_torch = None

def _get_lpips():
    """Lazy-load LPIPS + torch on first use. Returns (model, torch) or (None, None)
    if the dependencies aren't installed — pipeline still works without it."""
    global _lpips_model, _lpips_torch
    if _lpips_model is not None:
        return _lpips_model, _lpips_torch
    try:
        import torch
        import lpips
        _lpips_torch = torch
        # 'alex' is the fastest backbone; 'vgg' is more accurate but slower.
        _lpips_model = lpips.LPIPS(net='alex', verbose=False).eval()
    except Exception:
        _lpips_model = False  # sentinel: tried and failed
        _lpips_torch = None
    return _lpips_model, _lpips_torch


def lpips_distance(patch1_bgr, patch2_bgr):
    """Returns LPIPS perceptual distance in [0, ~1]. Higher = more different.
    Returns None if LPIPS isn't available."""
    model, torch = _get_lpips()
    if not model:
        return None

    def to_tensor(bgr):
        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        # LPIPS expects 64x64+ float in [-1, 1], shape (1,3,H,W)
        rgb = cv2.resize(rgb, (128, 128))
        t = torch.from_numpy(rgb).float().permute(2, 0, 1) / 127.5 - 1.0
        return t.unsqueeze(0)

    with torch.no_grad():
        d = model(to_tensor(patch1_bgr), to_tensor(patch2_bgr))
    return float(d.item())


# -----------------------------
# ALIGNMENT (ORB homography + ECC sub-pixel refinement)
# -----------------------------
def align_images(img1, img2):
    gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)

    orb = cv2.ORB_create(3000)
    kp1, des1 = orb.detectAndCompute(gray1, None)
    kp2, des2 = orb.detectAndCompute(gray2, None)

    aligned = img2  # fallback

    if des1 is not None and des2 is not None:
        matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
        matches = matcher.match(des1, des2)

        if len(matches) >= 10:
            matches = sorted(matches, key=lambda x: x.distance)[:50]
            pts1 = np.float32([kp1[m.queryIdx].pt for m in matches])
            pts2 = np.float32([kp2[m.trainIdx].pt for m in matches])
            matrix, _ = cv2.findHomography(pts2, pts1, cv2.RANSAC)

            if matrix is not None:
                aligned = cv2.warpPerspective(
                    img2, matrix, (img1.shape[1], img1.shape[0])
                )

    # ---- ECC refinement, self-validating ----
    # Only keep the ECC result if it actually IMPROVED alignment
    # (measured by SSIM of the aligned images). Otherwise discard.
    try:
        g1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
        g2_before = cv2.cvtColor(aligned, cv2.COLOR_BGR2GRAY)
        before_score, _ = ssim(g1, g2_before, full=True)

        warp = np.eye(3, 3, dtype=np.float32)
        criteria = (
            cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT,
            100,
            1e-5,
        )
        _, warp = cv2.findTransformECC(
            g1, g2_before, warp, cv2.MOTION_HOMOGRAPHY, criteria, None, 5
        )
        candidate = cv2.warpPerspective(
            aligned,
            warp,
            (img1.shape[1], img1.shape[0]),
            flags=cv2.INTER_LINEAR + cv2.WARP_INVERSE_MAP,
        )
        g2_after = cv2.cvtColor(candidate, cv2.COLOR_BGR2GRAY)
        after_score, _ = ssim(g1, g2_after, full=True)

        if after_score > before_score:
            aligned = candidate
    except cv2.error:
        # ECC can fail to converge — keep ORB-only result in that case
        pass

    return aligned


# -----------------------------
# SSIM-BASED STRUCTURAL DIFF
# -----------------------------
def ssim_diff(gray1, gray2):
    """Returns a uint8 'difference' map where bright = structurally different."""
    score, diff_map = ssim(gray1, gray2, full=True)
    # ssim returns [-1, 1] where 1 == identical; invert + scale to 0..255
    diff_map = (1.0 - diff_map) * 255.0
    return score, diff_map.astype(np.uint8)


# -----------------------------
# CRACK DETECTION (IMPROVED)
# -----------------------------
def detect_cracks(gray_img):
    enhanced = cv2.equalizeHist(gray_img)

    edges = cv2.Canny(enhanced, 40, 120)

    kernel = np.ones((3,3), np.uint8)
    edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel, iterations=1)

    return edges


# -----------------------------
# MAIN FUNCTION
# -----------------------------
def compare_images(img1_path, img2_path, output_path, heatmap_path):

    img1 = cv2.imread(img1_path)
    img2 = cv2.imread(img2_path)

    img2 = cv2.resize(img2, (img1.shape[1], img1.shape[0]))

    # 🔥 ALIGN
    aligned_img2 = align_images(img1, img2)

    gray1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(aligned_img2, cv2.COLOR_BGR2GRAY)

    # Moderate blur — enough to absorb pixel noise / alignment jitter, but
    # small enough that handwritten letters and other tiny changes survive.
    blur1 = cv2.GaussianBlur(gray1, (5, 5), 0)
    blur2 = cv2.GaussianBlur(gray2, (5, 5), 0)

    # -----------------------------
    # DIFF — SSIM with a large window is the workhorse here.
    # A big win=21 window means a single new pixel doesn't trigger a hit;
    # only regions where the *local structure* genuinely changed do.
    # -----------------------------
    ssim_score, struct_diff = ssim_diff(blur1, blur2)

    diff = struct_diff
    heatmap = cv2.applyColorMap(diff, cv2.COLORMAP_JET)

    # -----------------------------
    # OBJECT DETECTION — shape-aware filtering
    # -----------------------------
    _, thresh = cv2.threshold(diff, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # OPEN first: erodes thin alignment-edge ribbons away entirely.
    # Then CLOSE: reconnects fragments of the real new object.
    # Smaller OPEN kernel preserves small changes (handwritten letters, marks).
    # Median blur first — single best speckle killer for textured surfaces
    # (wood grain, carpet, asphalt) where SSIM picks up tiny texture jitter.
    thresh = cv2.medianBlur(thresh, 7)

    open_k  = np.ones((7, 7), np.uint8)
    close_k = np.ones((15, 15), np.uint8)
    thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN,  open_k,  iterations=2)
    thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, close_k, iterations=2)

    result = aligned_img2.copy()
    img_area = gray1.shape[0] * gray1.shape[1]

    # ---- Global alignment-failure guard ----
    # If the diff covers a huge portion of the image, alignment is broken
    # (camera moved/rotated too much between shots). Bail out cleanly
    # instead of drawing one giant useless box.
    diff_ratio = cv2.countNonZero(thresh) / img_area
    alignment_failed = diff_ratio > 0.70

    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if alignment_failed:
        contours = []  # skip detection entirely
    # Allow tiny changes (single letters, marks). LPIPS verification below
    # is what keeps false positives out at this size.
    min_area = max(400, img_area * 0.0002)  # ~0.02% of image

    # Score each contour and only keep "blob-shaped" ones (real objects),
    # not "ribbon-shaped" ones (edge jitter).
    candidates = []
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area < min_area:
            continue

        x, y, w, h = cv2.boundingRect(cnt)

        # Solidity = contour area / convex hull area. Real objects ≈ 0.6+.
        # Thin ribbons along edges have very low solidity.
        hull_area = cv2.contourArea(cv2.convexHull(cnt))
        solidity = area / hull_area if hull_area > 0 else 0

        # Extent = contour area / bounding-box area. Slivers ≈ low extent.
        extent = area / (w * h) if w * h > 0 else 0

        # Aspect ratio sanity: skip extreme slivers
        aspect = max(w, h) / max(1, min(w, h))

        # LARGE-BLOB BYPASS: any region covering >2% of the image is
        # almost certainly a real change (added/removed object). Trust it
        # unconditionally — don't apply shape filters that could reject
        # irregular shapes (e.g. paper partially occluded by a plant).
        is_large = area > img_area * 0.02

        # Relaxed thresholds — small handwritten marks can be irregular.
        # LPIPS does the final perceptual gate.
        if not is_large and (solidity < 0.30 or extent < 0.18 or aspect > 12):
            continue

        candidates.append((area, x, y, w, h))

    # Keep only the top few biggest blobs — the "real" changes.
    candidates.sort(reverse=True)
    candidates = candidates[:15]  # raise cap; LPIPS will prune below

    # -----------------------------
    # LPIPS PERCEPTUAL VERIFICATION
    # For each surviving candidate region, check that the patch in img2 is
    # *perceptually* different from the same region in img1. This kills the
    # last layer of false positives (lighting shifts, residual jitter on
    # textured surfaces) that survived the SSIM + shape filter.
    # -----------------------------
    LPIPS_THRESHOLD = 0.15  # tuned: <0.10 = ~identical, >0.20 = clearly different
    verified = []
    lpips_scores = []

    for area, x, y, w, h in candidates:
        # Pad the crop slightly so LPIPS sees context, not just the object edge
        pad = 10
        x0 = max(0, x - pad)
        y0 = max(0, y - pad)
        x1 = min(img1.shape[1], x + w + pad)
        y1 = min(img1.shape[0], y + h + pad)

        patch1 = img1[y0:y1, x0:x1]
        patch2 = aligned_img2[y0:y1, x0:x1]
        if patch1.size == 0 or patch2.size == 0:
            continue

        d = lpips_distance(patch1, patch2)

        # If LPIPS is unavailable (None), trust the upstream pipeline and keep.
        # If available, only keep candidates that are perceptually different.
        if d is None or d >= LPIPS_THRESHOLD:
            verified.append((area, x, y, w, h))
            if d is not None:
                lpips_scores.append(d)

    object_changes = len(verified)

    for _, x, y, w, h in verified:
        cv2.rectangle(result, (x, y), (x + w, y + h), (0, 0, 255), 3)

    # -----------------------------
    # CRACK DETECTION (ONLY IN CHANGED AREA)
    # -----------------------------
    cracks = detect_cracks(gray2)

    # 🔥 FIX 4 APPLIED
    cracks = cv2.bitwise_and(cracks, cracks, mask=thresh)

    crack_contours, _ = cv2.findContours(cracks, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    crack_count = 0
    total_crack_area = 0
    max_crack_area = 0

    # Only run crack analysis if NO new objects were detected. When a real
    # object appears, every edge of it would otherwise be flagged as a "crack".
    if object_changes == 0:
        for cnt in crack_contours:
            area = cv2.contourArea(cnt)

            if 80 < area < 1500:
                x, y, w, h = cv2.boundingRect(cnt)

                if w < 60 and h < 60:
                    crack_count += 1
                    total_crack_area += area
                    max_crack_area = max(max_crack_area, area)

                    cv2.rectangle(result, (x, y), (x+w, y+h), (255,0,0), 2)

    # -----------------------------
    # SEVERITY SCORING
    # >=1 real object change → HIGH. Otherwise → NONE.
    # -----------------------------
    if object_changes >= 1 or crack_count >= 1:
        severity = "HIGH"
    else:
        severity = "NONE"

    # -----------------------------
    # DRAW LABEL
    # -----------------------------
    color = (0,255,0)

    if severity == "MEDIUM":
        color = (0,255,255)
    elif severity == "HIGH":
        color = (0,0,255)

    label = f"Severity: {severity}"
    if alignment_failed:
        label = "ALIGNMENT FAILED — retake from same angle"
        color = (0, 165, 255)
    cv2.putText(result, label, (20, 40),
                cv2.FONT_HERSHEY_SIMPLEX, 1, color, 3)

    # Save outputs
    cv2.imwrite(output_path, result)
    cv2.imwrite(heatmap_path, heatmap)

    return {
        "crack_count": crack_count,
        "total_crack_area": int(total_crack_area),
        "max_crack_area": int(max_crack_area),
        "object_changes": object_changes,
        "ssim_score": float(ssim_score),
        "lpips_max": float(max(lpips_scores)) if lpips_scores else None,
        "lpips_available": lpips_scores != [] or _get_lpips()[0] is not False,
        "alignment_failed": bool(alignment_failed),
        "severity": severity
    }