import gradio as gr
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import numpy as np
import pickle
import os
import cv2
import logging
from sklearn.cluster import KMeans

def detect_connect_and_crop(image_source: Image.Image):
    """
    1. Detects bolt holes using Hough Transform.
    2. Separates them into Top and Bottom rows.
    3. Fits a straight horizontal line through the center of each row.
    4. Crops the region strictly between these two lines.

    """
    # Convert PIL Image to OpenCV format (RGB -> BGR usually, but we keep RGB for display)
    img_rgb = np.array(image_source)
    img_h, img_w = img_rgb.shape[:2]
    
    # Convert to Grayscale for detection
    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    
    # ── Step 1: Detect Circles (Bolt Holes) ─────────────────────────────────
    circles = cv2.HoughCircles(
        gray,
        cv2.HOUGH_GRADIENT,
        dp=1,            # Resolution ratio
        minDist=50,      # Min distance between holes (prevents double detection)
        param1=100,      # Canny edge threshold (higher = fewer edges)
        param2=30,       # Accumulator threshold (lower = more circles)
        minRadius=10,    # Min size of bolt hole
        maxRadius=40     # Max size of bolt hole
    )
    
    if circles is None:
        return image_source, image_source, "❌ No bolt holes detected. Adjust lighting or focus."
    
    # Convert coordinates to integers
    circles = np.round(circles[0]).astype(int)
    
    # ── Step 2: Separate into Top and Bottom Rows ───────────────────────────
    # We sort all circles by their Y (height) coordinate
    # The "Gap" in Y values tells us where the split between rows is.
    
    ys = [c[1] for c in circles]
    y_median = np.median(ys) # Find the middle point of the image's height logic
    
    # Split circles based on if they are above or below the median height
    top_row = [c for c in circles if c[1] < y_median]
    bot_row = [c for c in circles if c[1] >= y_median]
    
    if len(top_row) < 2 or len(bot_row) < 2:
        return image_source, image_source, "⚠️ Need at least 2 holes in both top and bottom rows to draw lines."

    # ── Step 3: Calculate Straight Lines (Average Y) ────────────────────────
    # We take the average Y of the top row to get a straight horizontal line
    y_top_line = int(np.mean([c[1] for c in top_row]))
    
    # We take the average Y of the bottom row
    y_bot_line = int(np.mean([c[1] for c in bot_row]))


    
    # ── Step 4: Visualization ───────────────────────────────────────────────
    vis_img = img_rgb.copy()
    
    # Colors
    LINE_COLOR = (0, 255, 0)     # Green
    HOLE_COLOR = (255, 0, 0)     # Red
    CENTER_COLOR = (255, 255, 0) # Yellow dots
    
    # Draw the straight horizontal lines across the full width
    cv2.line(vis_img, (0, y_top_line), (img_w, y_top_line), LINE_COLOR, 3)
    cv2.line(vis_img, (0, y_bot_line), (img_w, y_bot_line), LINE_COLOR, 3)
    
    # Draw the holes and their centers to show what we detected
    for (x, y, r) in circles:
        cv2.circle(vis_img, (x, y), r, HOLE_COLOR, 2)       # The hole outline
        cv2.circle(vis_img, (x, y), 3, CENTER_COLOR, -1)    # The center dot

    # ── Step 5: Crop the Region Between Lines ──────────────────────────────
    # Slice the array: img[y_start : y_end, x_start : x_end]
    # We add a tiny padding (e.g., +/- 0) or keep it exact.
    
    # Ensure y_top is actually above y_bot
    y_start, y_end = sorted([y_top_line, y_bot_line])
    
    # Safety check for crop size
    if y_end - y_start < 10:
        return Image.fromarray(vis_img), Image.fromarray(img_rgb), "⚠️ Calculated lines are too close."
        
    cropped_img = img_rgb[y_start:y_end, 0:img_w]
    
    # Generate Stats Text
    stats_text = (
        f"✅ **Processing Complete**\n"
        f"• Top Row Holes: {len(top_row)}\n"
        f"• Bottom Row Holes: {len(bot_row)}\n"
        f"• Top Line Y-Pos: {y_top_line} px\n"
        f"• Bottom Line Y-Pos: {y_bot_line} px\n"
        f"• Cropped Height: {y_end - y_start} px"
    )



    return Image.fromarray(vis_img), Image.fromarray(cropped_img), stats_text

# ── Gradio Interface ────────────────────────────────────────────────────────
with gr.Blocks(title="Engine Line Connector", theme=gr.themes.Default(primary_hue="blue")) as demo:
    gr.Markdown(
        """
        # 📏 Bolt-Center Line Connector
        ### Precision Cropping Tool
        1. Detects bolt holes.
        2. Draws a straight line through the average center of the top & bottom rows.
        3. Crops the saddle area exactly between these two lines.
        """
    )
    
    with gr.Row():
        input_image = gr.Image(type="pil", label="📷 Upload Engine Block Image")




        
    run_button = gr.Button("🔗 Connect Centers & Crop", variant="primary")
    
    with gr.Row():
        output_vis = gr.Image(label="✏️ Visualized Lines & Centers")
        output_crop = gr.Image(label="✂️ Final Cropped Area")
        
    stats_output = gr.Markdown()
    
    run_button.click(
        fn=detect_connect_and_crop,
        inputs=[input_image],
        outputs=[output_vis, output_crop, stats_output]
    )

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