import cv2
import numpy as np
from PIL import Image, ImageDraw
import gradio as gr

def classify_pipe_material(image_np):
    """
    Classify overall material based on image brightness.
    Brighter images (mean intensity > 130) are assumed to be Plastic; otherwise, Metal.
    """
    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
    mean_intensity = np.mean(gray)
    return "Plastic" if mean_intensity > 130 else "Metal"

def detect_rust(roi):
    """
    Detect rust by checking for reddish-brown hues in the ROI.
    """
    hsv_roi = cv2.cvtColor(roi, cv2.COLOR_RGB2HSV)
    lower_rust = np.array([5, 50, 50])
    upper_rust = np.array([25, 255, 255])
    mask = cv2.inRange(hsv_roi, lower_rust, upper_rust)
    rust_ratio = np.count_nonzero(mask) / float(roi.shape[0] * roi.shape[1])
    return rust_ratio

def detect_mold(roi):
    """
    Detect mold by looking for greenish hues, which may indicate fungal growth.
    """
    hsv_roi = cv2.cvtColor(roi, cv2.COLOR_RGB2HSV)
    lower_mold = np.array([35, 50, 20])
    upper_mold = np.array([85, 255, 120])
    mask = cv2.inRange(hsv_roi, lower_mold, upper_mold)
    mold_ratio = np.count_nonzero(mask) / float(roi.shape[0] * roi.shape[1])
    return mold_ratio

def detect_water_damage(roi):
    """
    Detect water damage by checking for discoloration typical of stains (dark brownish-yellow).
    """
    hsv_roi = cv2.cvtColor(roi, cv2.COLOR_RGB2HSV)
    lower_water = np.array([5, 50, 50])
    upper_water = np.array([20, 200, 150])
    mask = cv2.inRange(hsv_roi, lower_water, upper_water)
    water_ratio = np.count_nonzero(mask) / float(roi.shape[0] * roi.shape[1])
    return water_ratio

def classify_defect(roi):
    """
    Classify the defect using a combination of color and texture heuristics.
    Priority is given to color cues:
      - "Rust": reddish-brown
      - "Mold": greenish
      - "Water Damage": discoloration from water stains
    Then geometric/texture analysis is used to differentiate "Crack" and "Corrosion."
    """
    area = roi.shape[0] * roi.shape[1]
    std_intensity = np.std(roi)
    
    rust_ratio = detect_rust(roi)
    mold_ratio = detect_mold(roi)
    water_ratio = detect_water_damage(roi)
    
    # Check for color-based defects first.
    if rust_ratio > 0.25:
        return "Rust"
    elif mold_ratio > 0.2:
        return "Mold"
    elif water_ratio > 0.2:
        return "Water Damage"
    
    # Then use size and intensity variation for structural defects.
    if area < 5000 and std_intensity > 50:
        return "Crack"
    elif area >= 5000 and std_intensity > 40:
        return "Corrosion"
    else:
        return "Other Defect"

def detect_infrastructure_issues(image: Image.Image):
    try:
        # Convert the PIL image to a NumPy array (RGB)
        image_np = np.array(image)
        annotated = image.copy()  # For drawing annotations
        draw = ImageDraw.Draw(annotated)
        
        # Classify the overall pipe (or structure) material
        overall_material = classify_pipe_material(image_np)
        
        # Enhance image for defect detection:
        gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        enhanced = clahe.apply(gray)
        blurred = cv2.GaussianBlur(enhanced, (5, 5), 0)
        
        # Adaptive thresholding highlights potential defect areas
        thresh = cv2.adaptiveThreshold(
            blurred, 255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY_INV,
            11, 2
        )
        
        # Morphological closing to consolidate defect regions
        kernel = np.ones((3, 3), np.uint8)
        morph = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
        
        # Edge detection to identify defect boundaries
        edges = cv2.Canny(morph, 50, 150)
        
        # Extract contours as candidate defect regions
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        detections = []
        
        # Define distinct colors for each defect type
        colors = {
            "Rust": "orange",
            "Mold": "purple",
            "Water Damage": "blue",
            "Crack": "red",
            "Corrosion": "cyan",
            "Other Defect": "gray"
        }
        
        for cnt in contours:
            if cv2.contourArea(cnt) < 100:  # Filter out noise
                continue
            x, y, w, h = cv2.boundingRect(cnt)
            roi = image_np[y:y+h, x:x+w]
            if roi.size == 0:
                continue
            defect_type = classify_defect(roi)
            detection_info = f"{defect_type} at ({x}, {y}, {w}, {h})"
            detections.append(detection_info)
            
            # Draw bounding box with defect-specific color
            box_color = colors.get(defect_type, "gray")
            draw.rectangle([x, y, x+w, y+h], outline=box_color, width=2)
            draw.text((x, y-10), defect_type, fill=box_color)
        
        # Create a textual summary
        if detections:
            summary = f"Overall Material: {overall_material}\nDetected Issues:\n" + "\n".join(detections)
        else:
            summary = f"Overall Material: {overall_material}\nNo significant defects detected."
        
        return annotated, summary
    except Exception as e:
        print("Error during detection:", e)
        return image, f"Error: {e}"

iface = gr.Interface(
    fn=detect_infrastructure_issues,
    inputs=gr.Image(type="pil", label="Upload an Infrastructure Image"),
    outputs=[gr.Image(label="Annotated Image"), gr.Textbox(label="Detection Summary")],
    title="Comprehensive Home Infrastructure Defect Detector",
    description=(
        "Upload an image of a pipe or any home infrastructure (walls, floors, etc.) to detect defects. "
        "This tool identifies issues such as Rust (orange), Mold (purple), Water Damage (blue), Cracks (red), "
        "and Corrosion (cyan), and returns both an annotated image and a detailed summary."
    )
)

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