import gradio as gr
from inference_sdk import InferenceHTTPClient
from PIL import Image, ImageDraw, ImageFont
import os
from collections import defaultdict

API_KEY = os.getenv("ROBOFLOW_API_KEY") 

CLIENT = InferenceHTTPClient(
    api_url="https://detect.roboflow.com",
    api_key=API_KEY  
)


# Model settings
MODEL_ID = "hvacsym/5"
CONFIDENCE_THRESHOLD = 0.3  # Confidence threshold for filtering predictions
GRID_SIZE = (3, 3)  # 4x4 segmentation

def process_image(image_path):
    """Processes an uploaded image and returns the final image with bounding boxes & symbol counts."""
    original_image = Image.open(image_path)
    width, height = original_image.size
    seg_w, seg_h = width // GRID_SIZE[1], height // GRID_SIZE[0]

    # Create a copy of the image for bounding boxes
    final_image = original_image.copy()
    draw_final = ImageDraw.Draw(final_image)

    # Load font
    try:
        font = ImageFont.truetype("arial.ttf", 10)
    except:
        font = ImageFont.load_default()

    # Dictionary for total counts
    total_counts = defaultdict(int)

    # Process each segment
    for row in range(GRID_SIZE[0]):
        for col in range(GRID_SIZE[1]):
            x1, y1 = col * seg_w, row * seg_h
            x2, y2 = (col + 1) * seg_w, (row + 1) * seg_h

            segment = original_image.crop((x1, y1, x2, y2))
            segment_path = f"segment_{row}_{col}.png"
            segment.save(segment_path)

            # Run inference
            result = CLIENT.infer(segment_path, model_id=MODEL_ID)
            filtered_predictions = [
                pred for pred in result["predictions"] if pred["confidence"] * 100 >= CONFIDENCE_THRESHOLD
            ]

            # Draw bounding boxes and update counts
            for obj in filtered_predictions:
                sx, sy, sw, sh = obj["x"], obj["y"], obj["width"], obj["height"]
                class_name = obj["class"]
                confidence = obj["confidence"]
                total_counts[class_name] += 1

                # Adjust coordinates for final image
                x_min, y_min = x1 + (sx - sw // 2), y1 + (sy - sh // 2)
                x_max, y_max = x1 + (sx + sw // 2), y1 + (sy + sh // 2)

                # Draw bounding box
                draw_final.rectangle([x_min, y_min, x_max, y_max], outline="green", width=2)

                # Draw label
                text = f"{class_name} {confidence:.2f}"
                draw_final.text((x_min + 2, y_min - 10), text, fill="white", font=font)

    # Save final image with bounding boxes
    final_image_path = "final_detected_image.png"
    final_image.save(final_image_path)

    return final_image_path, total_counts

def process_uploaded_image(image):
    """Gradio wrapper function that calls `process_image` and formats the output."""
    final_image_path, total_counts = process_image(image)

    # Convert count dictionary to readable text
    count_text = "\n".join([f"{label}: {count}" for label, count in total_counts.items()])

    return final_image_path, count_text

# Deploy with Gradio
iface = gr.Interface(
    fn=process_uploaded_image,
    inputs=gr.Image(type="filepath"),  # Gradio input expects a file path
    outputs=[gr.Image(type="filepath"), gr.Text()],
    title="HVAC Symbol Detector",
    description="Upload an HVAC blueprint image. The model will detect symbols and return the final image with bounding boxes along with symbol counts."
)

# Launch the Gradio app
iface.launch()