import cv2
import numpy as np
import gradio as gr

def image_inference(image_path, mode, epsilon_thresh):
    # Read the image
    img = cv2.cvtColor(image_path, cv2.COLOR_RGB2BGR)
    if img is None:
        raise FileNotFoundError(f"Image at path '{image_path}' not found.")
    
    # Create a copy of the original image
    img_copy = img.copy()
    
    # Convert the image to grayscale
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # Apply thresholding
    _, img_thresh = cv2.threshold(img_gray, 200, 255, cv2.THRESH_BINARY_INV)

    # Find the contours
    contours, _ = cv2.findContours(img_thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    contour_info = f"Number of contours found: {len(contours)}"  # Text output for Gradio
    # print(f"Number of contours found = {len(contours)}")
    
    def convert_color(hsv):
        # Utility to convert a single HSV color tuple into BGR
        pixel_img = np.uint8([[hsv]])
        return tuple(int(i) for i in cv2.cvtColor(pixel_img, cv2.COLOR_HSV2BGR).flatten())

    if mode == "contour":
        if len(contours) > 0:
            for i, single_contour in enumerate(contours):
                hsv = (int(i/len(contours) * 180), 255, 255)
                color = convert_color(hsv)
                # Draw the contour
                img_final = cv2.drawContours(img_copy, contours, i, color, 4)
                # Calculate and display the area
                # area = cv2.contourArea(single_contour)
                # x, y, w, h = cv2.boundingRect(single_contour)
                # img_final = cv2.putText(img_final, f"Area: {area}", (x, y - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1)
        else:
            img_final = img_copy  # No contours found, return the original image
    elif mode == "rotated rectangle":
        for cnt in contours:
            # Rotated bounding box
            box = cv2.minAreaRect(cnt)
            box_pts = np.intp(cv2.boxPoints(box))
            img_final = cv2.drawContours(img_copy, [box_pts], -1, (0, 255, 0), 4)
            
            # Calculate and display the area
            area = cv2.contourArea(cnt)
            x, y, w, h = cv2.boundingRect(cnt)
            img_final = cv2.putText(img_final, f"Area: {area}", (x, y - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1)
            
    elif mode == "rectangle":
        for cnt in contours:
            # Bounding rectangle
            x, y, w, h = cv2.boundingRect(cnt)
            img_final = cv2.rectangle(img_copy, (x, y), (x + w, y + h), (0, 255, 0), 4)
            
            # Calculate and display the area
            area = cv2.contourArea(cnt)
            img_final = cv2.putText(img_final, f"Area: {area}", (x, y - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1)
            
    elif mode == "circle":
        for cnt in contours:
            # Enclosing circle
            ((x, y), radius) = cv2.minEnclosingCircle(cnt)
            img_final = cv2.circle(img_copy, (int(x), int(y)), int(round(radius)), (0, 255, 0), 4)
            
            # Calculate and display the perimeter
            perimeter = cv2.arcLength(cnt, True)
            x, y, w, h = cv2.boundingRect(cnt)
            img_final = cv2.putText(img_final, f"Perimeter: {perimeter:.2f}", (x, y - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1)
            
    elif mode == "ellipse":
        for cnt in contours:
            if len(cnt) >= 5:  # fitEllipse requires at least 5 points
                ellipse = cv2.fitEllipse(cnt)
                img_final = cv2.ellipse(img_copy, ellipse, (0, 255, 0), 4)
                
                # Calculate and display the perimeter
                perimeter = cv2.arcLength(cnt, True)
                x, y, w, h = cv2.boundingRect(cnt)
                img_final = cv2.putText(img_final, f"Perimeter: {perimeter:.2f}", (x, y - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 0), 1)
            else:
                img_final = img_copy
    elif mode == "centroid":
        for cnt in contours:
            M = cv2.moments(cnt)
            if M["m00"] != 0:
                x = int(round(M["m10"] / M["m00"]))
                y = int(round(M["m01"] / M["m00"]))
                img_final = cv2.circle(img_copy, (x, y), 10, (255, 0, 0), -1)
            else:
                img_final = img_copy
    elif mode == "contour approx":
        if len(contours) > 0:
            c = max(contours, key=cv2.contourArea)
            peri = cv2.arcLength(c, True)
            approx = cv2.approxPolyDP(c, epsilon_thresh * peri, True)
            img_final = cv2.drawContours(img_copy, [approx], -1, (0, 255, 0), 3)
            x, y, w, h = cv2.boundingRect(c)
            text = f"eps={epsilon_thresh:.4f}, num_pts={len(approx)}"
            img_final = cv2.putText(img_copy, text, (x, y - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
        else:
            img_final = img_copy
    else:
        img_final = img_copy
        print(f"Mode '{mode}' not recognized. Please choose a valid mode.")
    
    return contour_info, cv2.cvtColor(img_final, cv2.COLOR_BGR2RGB)

# Gradio interface
input_image = gr.Image(type="numpy", label="Input Image")
epsilon_thresh = gr.Slider(
    0.01,
    0.1,
    step=0.01,
    value=0.05,
    label="Epsilon Threshold",
    info="Adjust the Contour Threshold according to the object size that you want to detect. Only Applicable for Contour Approx",
)
mode = gr.Radio(
    ["contour", "rectangle", "rotated rectangle", "circle", "ellipse", "centroid", "contour approx"],
    label="Contour Type",
    info="Choose the MODE",
)
output_text = gr.Textbox(label="Contour Information")
output_image = gr.Image(label="Output Image")

app = gr.Interface(
    fn=image_inference,
    inputs=[input_image, mode, epsilon_thresh],
    outputs=[output_text, output_image],
    title="Contour Detection using OpenCV",
    description="A Gradio app for dynamic image analysis using Contour detection techniques.",
    allow_flagging="never",
    examples=[["./sample/tictactoe.png", "contour", float(0.01)], ["./sample/tetris_blocks.png", "rectangle", float(0.00)], ["./sample/shape.png", "contour approx", float(0.05)]],
    cache_examples=False,
)
app.queue().launch()