import gradio as gr
import cv2
import numpy as np
from PIL import Image
from sklearn.cluster import DBSCAN
from scipy.spatial import ConvexHull
import warnings
warnings.filterwarnings("ignore")


# ─────────────────────────────────────────────────────────────────
# CORE UTILITIES
# ─────────────────────────────────────────────────────────────────

def preprocess_image(img_rgb: np.ndarray) -> list[np.ndarray]:
    """
    Generate multiple preprocessed versions to maximize detection coverage
    across different lighting, contrast, and color conditions.
    """
    h, w = img_rgb.shape[:2]
    long_side = max(h, w)
    scale = 1.0
    if long_side > 1200:
        scale = 1200 / long_side
        img_rgb = cv2.resize(img_rgb, (int(w * scale), int(h * scale)))

    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    variants = []

    # 1. CLAHE – best for industrial / low-contrast images
    clahe = cv2.createCLAHE(clipLimit=9.0, tileGridSize=(8, 8))
    clahe_gray = clahe.apply(gray)
    variants.append(clahe_gray)

    # 2. Gaussian blur (reduce noise in sharp shop photos)
    blurred = cv2.GaussianBlur(gray, (5, 5), 1.5)
    variants.append(blurred)

    # 3. Bilateral filter (edge-preserving smoothing)
    bilateral = cv2.bilateralFilter(gray, 9, 75, 75)
    variants.append(bilateral)

    # 4. Equalized histogram (handles B&W / washed-out images)
    equalized = cv2.equalizeHist(gray)
    variants.append(equalized)

    # 5. Inverted CLAHE (holes appear bright in some lighting)
    variants.append(cv2.bitwise_not(clahe_gray))

    return variants, img_rgb, scale


def hough_multi_params(gray: np.ndarray, img_shape: tuple) -> np.ndarray:
    """
    Run HoughCircles across a sweep of parameters.
    Radius bounds are computed as a fraction of image dimensions.
    """
    h, w = img_shape[:2]
    short_side = min(h, w)

    # Radius range: bolt holes are typically 1–8% of image width
    r_min = max(8,  int(short_side * 0.012))
    r_max = max(40, int(short_side * 0.09))

    param_grid = [
        dict(dp=1.0, minDist=int(short_side*0.04), param1=80,  param2=28, minRadius=r_min, maxRadius=r_max),
        dict(dp=1.2, minDist=int(short_side*0.04), param1=60,  param2=22, minRadius=r_min, maxRadius=r_max),
        dict(dp=1.5, minDist=int(short_side*0.05), param1=100, param2=35, minRadius=r_min, maxRadius=r_max),
        dict(dp=1.0, minDist=int(short_side*0.03), param1=50,  param2=18, minRadius=r_min, maxRadius=r_max),
        dict(dp=2.0, minDist=int(short_side*0.05), param1=70,  param2=25, minRadius=r_min, maxRadius=r_max),
    ]

    all_circles = []
    for p in param_grid:
        c = cv2.HoughCircles(gray, cv2.HOUGH_GRADIENT, **p)
        if c is not None:
            all_circles.extend(np.round(c[0]).astype(int).tolist())

    return np.array(all_circles) if all_circles else np.empty((0, 3), dtype=int)


def contour_circle_fallback(gray: np.ndarray) -> np.ndarray:
    """
    Contour-based circle detector as a fallback.
    Finds circular blobs via adaptive thresholding + circularity filter.
    """
    h, w = gray.shape
    short = min(h, w)
    r_min = int(short * 0.012)
    r_max = int(short * 0.09)

    thresh = cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 21, 4
    )
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)

    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    circles = []
    for cnt in contours:
        area = cv2.contourArea(cnt)
        perim = cv2.arcLength(cnt, True)
        if perim == 0:
            continue
        circularity = 4 * np.pi * area / (perim ** 2)
        if circularity < 0.55:
            continue
        (cx, cy), radius = cv2.minEnclosingCircle(cnt)
        if r_min <= int(radius) <= r_max:
            circles.append([int(cx), int(cy), int(radius)])

    return np.array(circles) if circles else np.empty((0, 3), dtype=int)


def deduplicate_circles(circles: np.ndarray, tol_frac: float = 0.04, img_shape=None) -> np.ndarray:
    """
    Remove duplicate/overlapping detections using NMS-style spatial dedup.
    Tolerance is a fraction of the shorter image dimension.
    """
    if len(circles) == 0:
        return circles
    h, w = img_shape[:2]
    tol = int(min(h, w) * tol_frac)
    circles = circles[np.argsort(circles[:, 2])[::-1]]  # sort by radius desc
    kept = []
    for c in circles:
        if all(np.linalg.norm(c[:2] - k[:2]) > tol for k in kept):
            kept.append(c)
    return np.array(kept)


# ─────────────────────────────────────────────────────────────────
# ROTATION-AWARE CLUSTERING
# ─────────────────────────────────────────────────────────────────

def detect_orientation(points: np.ndarray) -> float:
    """
    PCA on hole centres → returns dominant axis angle in degrees.
    Works regardless of camera rotation.
    """
    if len(points) < 2:
        return 0.0
    centered = points - points.mean(axis=0)
    _, _, vt = np.linalg.svd(centered)
    angle = np.degrees(np.arctan2(vt[0, 1], vt[0, 0]))
    return angle


def cluster_holes(circles: np.ndarray, img_shape: tuple):
    """
    DBSCAN clustering along the perpendicular axis to the detected orientation.
    Returns dict of cluster_id → list of (x, y).
    Falls back to 2-cluster top/bottom split if DBSCAN finds only 1 cluster.
    """
    pts = circles[:, :2].astype(float)
    angle = detect_orientation(pts)

    # Project onto the axis perpendicular to the main orientation
    angle_rad = np.radians(angle)
    perp = np.array([-np.sin(angle_rad), np.cos(angle_rad)])
    proj = pts @ perp  # 1-D projection

    h, w = img_shape[:2]
    eps = min(h, w) * 0.08
    labels = DBSCAN(eps=eps, min_samples=1).fit_predict(proj.reshape(-1, 1))

    clusters = {}
    for lbl, pt in zip(labels, circles):
        clusters.setdefault(int(lbl), []).append((int(pt[0]), int(pt[1])))

    # Sort clusters by median projection (top → bottom / left → right)
    medians = {k: np.median([pts[i] @ perp for i, l in enumerate(labels) if l == k])
               for k in clusters}
    sorted_keys = sorted(clusters.keys(), key=lambda k: medians[k])
    ordered = {i: clusters[k] for i, k in enumerate(sorted_keys)}

    return ordered, angle


# ─────────────────────────────────────────────────────────────────
# BOUNDING BOX
# ─────────────────────────────────────────────────────────────────

def build_bbox(all_pts: list, extend_px: int, img_shape: tuple):
    all_x = [p[0] for p in all_pts]
    all_y = [p[1] for p in all_pts]
    h, w = img_shape[:2]
    x1 = max(0, min(all_x) - extend_px)
    y1 = max(0, min(all_y) - extend_px)
    x2 = min(w, max(all_x) + extend_px)
    y2 = min(h, max(all_y) + extend_px)
    return x1, y1, x2, y2


# ─────────────────────────────────────────────────────────────────
# VISUALIZATION
# ─────────────────────────────────────────────────────────────────

PALETTE = [
    (0, 230, 80),    # green
    (255, 180, 0),   # amber
    (0, 200, 255),   # cyan
    (255, 80, 200),  # magenta
    (180, 255, 80),  # lime
    (255, 140, 60),  # orange
]

def draw_results(vis: np.ndarray, circles: np.ndarray, clusters: dict,
                 angle: float, extend_px: int):
    h, w = vis.shape[:2]
    all_pts = []

    for row_idx, pts in clusters.items():
        color = PALETTE[row_idx % len(PALETTE)]
        pts_sorted = sorted(pts, key=lambda p: p[0])
        all_pts.extend(pts)

        # Draw connecting line between holes in same cluster
        if len(pts_sorted) >= 2:
            cv2.line(vis, pts_sorted[0], pts_sorted[-1], color, 2, cv2.LINE_AA)

        for x, y in pts:
            r_draw = 22
            cv2.circle(vis, (x, y), r_draw, color, 3, cv2.LINE_AA)
            cv2.circle(vis, (x, y), 5,      (255, 60, 60), -1, cv2.LINE_AA)

    if not all_pts:
        return vis, None

    x1, y1, x2, y2 = build_bbox(all_pts, extend_px, vis.shape)

    # Draw bounding box
    cv2.rectangle(vis, (x1, y1), (x2, y2), (255, 80, 80), 3, cv2.LINE_AA)

    # Draw orientation arrow in corner of bbox
    cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
    arrow_len = 50
    ang = np.radians(angle)
    ex, ey = int(cx + arrow_len * np.cos(ang)), int(cy + arrow_len * np.sin(ang))
    cv2.arrowedLine(vis, (cx, cy), (ex, ey), (255, 255, 255), 2,
                    cv2.LINE_AA, tipLength=0.3)

    # Overlay stats
    font, fs, thick = cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1
    for row_idx, pts in clusters.items():
        c = PALETTE[row_idx % len(PALETTE)]
        label = f"Row {row_idx+1}: {len(pts)} holes"
        tx = x1 + 6
        ty = y1 + 22 + row_idx * 22
        ty = min(ty, h - 8)
        cv2.putText(vis, label, (tx, ty), font, fs, (0,0,0), thick+2, cv2.LINE_AA)
        cv2.putText(vis, label, (tx, ty), font, fs, c,     thick,   cv2.LINE_AA)

    return vis, (x1, y1, x2, y2)


# ─────────────────────────────────────────────────────────────────
# MAIN PIPELINE
# ─────────────────────────────────────────────────────────────────

def detect_bolt_holes(image, extend_px: int = 40, min_holes: int = 2):
    """
    Robust, rotation-invariant bolt-hole localizer.

    Pipeline:
    1. Multi-variant preprocessing (CLAHE, bilateral, equalized, inverted)
    2. Hough + contour-based detection across parameter sweeps
    3. Aggregate and deduplicate across all attempts
    4. PCA-based orientation detection
    5. DBSCAN clustering along perpendicular axis
    6. Draw bounding box + per-row annotations
    7. Crop and return ROI
    """
    if image is None:
        return None, None, "⚠️ No image provided."

    img_rgb = np.array(image)
    if img_rgb.ndim == 2:
        img_rgb = cv2.cvtColor(img_rgb, cv2.COLOR_GRAY2RGB)
    elif img_rgb.shape[2] == 4:
        img_rgb = cv2.cvtColor(img_rgb, cv2.COLOR_RGBA2RGB)

    variants, img_rgb, scale = preprocess_image(img_rgb)

    # ── Aggregate circles from ALL preprocessed variants ────────────
    raw_circles = []
    for gray_variant in variants:
        c = hough_multi_params(gray_variant, img_rgb.shape)
        if len(c):
            raw_circles.extend(c.tolist())
        c2 = contour_circle_fallback(gray_variant)
        if len(c2):
            raw_circles.extend(c2.tolist())

    if not raw_circles:
        return Image.fromarray(img_rgb), Image.fromarray(img_rgb), \
               "❌ No circles detected. Try a closer or clearer image."

    circles = deduplicate_circles(np.array(raw_circles), img_shape=img_rgb.shape)

    if len(circles) < min_holes:
        return Image.fromarray(img_rgb), Image.fromarray(img_rgb), \
               f"⚠️ Only {len(circles)} unique circles found (need ≥{min_holes})."

    # ── Cluster + orient ─────────────────────────────────────────────
    clusters, angle = cluster_holes(circles, img_rgb.shape)

    # ── Draw ─────────────────────────────────────────────────────────
    vis = img_rgb.copy()
    vis, bbox = draw_results(vis, circles, clusters, angle, extend_px)

    # ── Crop ROI ──────────────────────────────────────────────────────
    if bbox:
        x1, y1, x2, y2 = bbox
        cropped = img_rgb[y1:y2, x1:x2]
    else:
        cropped = img_rgb

    # ── Stats ─────────────────────────────────────────────────────────
    total = sum(len(v) for v in clusters.values())
    rows_info = "\n".join(
        f"  • Row {i+1} ({len(pts)} holes): {pts}"
        for i, pts in clusters.items()
    )
    stats = (
        f"✅ **{total} bolt holes detected** across **{len(clusters)} row(s)**\n"
        f"📐 Dominant orientation: **{angle:.1f}°**\n"
        f"{rows_info}\n"
        + (f"📦 Bounding Box: ({x1},{y1}) → ({x2},{y2})  |  "
           f"ROI: {x2-x1}×{y2-y1} px" if bbox else "")
    )

    return Image.fromarray(vis), Image.fromarray(cropped), stats


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

CSS = """
body { font-family: 'Courier New', monospace; background: #0e0e0e; color: #e8e8e8; }
.gradio-container { max-width: 1200px; margin: 0 auto; }
.title-bar { text-align: center; padding: 18px 0 4px; }
.title-bar h1 { font-size: 1.7rem; letter-spacing: 3px; color: #00e5a0; margin: 0; }
.title-bar p { color: #888; font-size: 0.85rem; margin: 4px 0 0; }
footer { display: none !important; }
"""

with gr.Blocks(
    title="Bolt Hole Localizer v2 — Rotation-Invariant",
    theme=gr.themes.Base(
        primary_hue="green",
        neutral_hue="gray",
        font=[gr.themes.GoogleFont("Space Mono"), "monospace"],
    ),
    css=CSS,
) as demo:

    gr.HTML("""
    <div class="title-bar">
      <h1>🔩 BOLT HOLE LOCALIZER v2</h1>
      <p>Rotation-Invariant · Multi-Scale · DBSCAN Clustering · Any Engine Configuration</p>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            inp_img = gr.Image(type="pil", label="📷 Input Engine Image",
                               sources=["upload", "clipboard"])

            with gr.Accordion("⚙️ Advanced Settings", open=False):
                extend_slider = gr.Slider(
                    10, 120, value=40, step=5,
                    label="Bounding Box Margin (px)",
                    info="Extra padding around detected bolt cluster"
                )
                min_holes_slider = gr.Slider(
                    2, 8, value=2, step=1,
                    label="Minimum Holes Required",
                    info="Reject results below this count"
                )

            run_btn = gr.Button("🔍 Detect Bolt Holes", variant="primary", size="lg")

        with gr.Column(scale=2):
            with gr.Row():
                out_annotated = gr.Image(label="📌 Annotated — Holes + Orientation + BBox")
                out_cropped   = gr.Image(label="✂️ Cropped ROI")
            out_stats = gr.Markdown(label="Detection Stats")

    run_btn.click(
        fn=detect_bolt_holes,
        inputs=[inp_img, extend_slider, min_holes_slider],
        outputs=[out_annotated, out_cropped, out_stats],
    )

    gr.Markdown("""
    ---
    **Detection Strategy:**  
    Multi-variant preprocessing (CLAHE · bilateral · equalized · inverted) →
    Hough + contour sweep →
    Spatial deduplication →
    PCA orientation →
    DBSCAN row clustering →
    Rotation-invariant bounding box
    """)

if __name__ == "__main__":
    demo.launch()