import numpy as np
import cv2
import gradio as gr
from skimage.morphology import convex_hull_image

# =========================================================
# Utility Functions
# =========================================================

def orient_to_portrait(img):
    h, w = img.shape[:2]
    if w > h:
        img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
    return img

def analyze_cv_only(image, threshold_pct, min_defect_size):
    """
    Simulates the ML pipeline using traditional image processing.
    """
    if image is None:
        return None, "0.0%", "0.0%", None, None

    # 1. Preprocessing
    img_color = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    img_color = orient_to_portrait(img_color)
    gray = cv2.cvtColor(img_color, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)

    # 2. Simulate Print Segmentation (Thresholding)
    # Instead of a model, we use simple binary thresholding
    val = int((threshold_pct / 100.0) * 255)
    _, print_mask = cv2.threshold(blurred, val, 255, cv2.THRESH_BINARY_INV)

    # Find largest contour to represent the print
    contours, _ = cv2.findContours(print_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return image, "0.0%", "0.0%", None, None
        
    largest_cnt = max(contours, key=cv2.contourArea)
    print_mask_clean = np.zeros_like(print_mask)
    cv2.drawContours(print_mask_clean, [largest_cnt], -1, 255, thickness=-1)

    # 3. Simulate Defect Detection (Edge Detection within the print)
    # Detects high-contrast areas (cracks/voids) inside the print mask
    edges = cv2.Canny(blurred, 50, 150)
    defect_mask = cv2.bitwise_and(edges, print_mask_clean)
    
    # Filter defects by size
    final_defect_mask = np.zeros_like(defect_mask)
    def_cnts, _ = cv2.findContours(defect_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for cnt in def_cnts:
        if cv2.contourArea(cnt) >= min_defect_size:
            cv2.drawContours(final_defect_mask, [cnt], -1, 255, thickness=-1)

    # 4. Metrics
    # Integrity: Ratio of print to its convex hull
    try:
        hull = convex_hull_image(print_mask_clean > 0)
        integrity_score = np.sum(print_mask_clean > 0) / np.sum(hull)
    except:
        integrity_score = 0.0

    # Quality: Inverse of defect area
    print_area = np.sum(print_mask_clean > 0)
    defect_area = np.sum(final_defect_mask > 0)
    quality_score = 100.0 - ((defect_area / print_area) * 100.0) if print_area > 0 else 0.0

    # 5. Visualization
    vis = img_color.copy()
    cv2.drawContours(vis, [largest_cnt], -1, (255, 0, 0), 2) # Blue outline
    cv2.drawContours(vis, [cv2.convexHull(largest_cnt)], -1, (0, 0, 255), 2) # Red hull
    
    # Draw defects in green
    vis[final_defect_mask > 0] = [0, 255, 0]

    # Save outputs
    mask_path = "combined_mask.png"
    cv2.imwrite(mask_path, final_defect_mask)
    vis_path = "analysis_output.png"
    cv2.imwrite(vis_path, vis)

    return cv2.cvtColor(vis, cv2.COLOR_BGR2RGB), f"{integrity_score*100:.1f}%", f"{quality_score:.1f}%", vis_path, mask_path

# =========================================================
# UI Layout
# =========================================================
custom_css = """
.score-box { background-color: #f8f9fa; border-radius: 10px; padding: 15px; text-align: center; }
.score-value { font-size: 2rem !important; font-weight: 700 !important; color: #2b6cb0 !important; }
"""

with gr.Blocks(css=custom_css) as app:
    gr.Markdown("# CV-Only Print Analysis (No ML)")
    
    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="numpy", label="Input Image")
            thresh_slider = gr.Slider(50, 255, value=127, label="Brightness Threshold")
            size_slider = gr.Slider(0, 500, value=20, label="Min Defect Size")
            run_btn = gr.Button("Analyze", variant="primary")
            
        with gr.Column():
            img_output = gr.Image(label="Analysis Result")
            with gr.Row():
                area_out = gr.Markdown("0.0%", elem_classes=["score-value"])
                surface_out = gr.Markdown("0.0%", elem_classes=["score-value"])

    run_btn.click(
        fn=analyze_cv_only,
        inputs=[img_input, thresh_slider, size_slider],
        outputs=[img_output, area_out, surface_out]
    )

app.queue().launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)