import sys
import cv2
import json
import numpy as np
import os

def hex_to_hsv(hex_code):
    # Convert hex code to RGB format
    rgb_color = np.array([int(hex_code[i:i+2], 16) for i in (1, 3, 5)])
    # Convert RGB to BGR format
    bgr_color = np.flip(rgb_color)
    # Convert BGR to HSV format
    hsv_color = cv2.cvtColor(np.uint8([[bgr_color]]), cv2.COLOR_BGR2HSV)[0][0]
    return hsv_color

def draw_contours(image_path, colors_dict, output_path, min_contour_size=150, max_contour_size=10000):
    # Load the source image
    source_image = cv2.imread(image_path)

    if source_image is None:
        print(f"Error: Unable to load image at {image_path}")
        sys.exit(1)

    # Convert the image to HSV color space
    hsv_image = cv2.cvtColor(source_image, cv2.COLOR_BGR2HSV)

    # Initialize a blank mask to store combined ROIs for all colors
    combined_mask = np.zeros_like(hsv_image[:, :, 0])

    # Loop through each color
    for color_hex_code, bounds in colors_dict.items():
        # Create a mask for the specific color
        mask = cv2.inRange(hsv_image, bounds['lower'], bounds['upper'])

        # Combine the mask with the previous masks
        combined_mask = cv2.bitwise_or(combined_mask, mask)

    # Find contours in the combined mask
    contours, _ = cv2.findContours(combined_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    bounding_rectangles = []

    # Iterate through each contour to get bounding rectangle coordinates
    for cnt, contour in enumerate(contours):
        contour_area = cv2.contourArea(contour)
        if min_contour_size <= contour_area <= max_contour_size:
            x, y, w, h = cv2.boundingRect(contour)
            x1, y1 = x, y
            x2, y2 = x + w, y + h
            bounding_rectangles.append((cnt, (x1, y1, x2, y2)))  # Append each rectangle individually

            # Draw rectangle around each contour
            cv2.rectangle(source_image, (x1, y1), (x2, y2), (0, 255, 0), 2)

    # Save the resulting image
    cv2.imwrite(output_path, source_image)

    # Convert the array to JSON string
    json_string = json.dumps(bounding_rectangles)

    # Print the JSON string
    return json_string

if __name__ == "__main__":
    # Define the colors and their corresponding lower and upper bounds in HSV format
    colors = {
        '#FFFBDB': {'lower': np.array([20, 10, 200]), 'upper': np.array([30, 100, 255])},  # Light Yellow
        '#AFCB04': {'lower': np.array([24, 200, 150]), 'upper': np.array([44, 255, 255])}  # Lime Green


    }

    # Check if the script is run with the correct number of arguments
    if len(sys.argv) != 2:
        print("Usage: python script.py <image_file>")
        sys.exit(1)

    # Get the image file path from command-line arguments
    image_file = sys.argv[1]

    # Output path for the resulting image
    output_path = 'floor_plan_with_bbox789.png'

    # Draw contours on areas filled with specified colors
    json_output = draw_contours(image_file, colors, output_path, min_contour_size=150, max_contour_size=10000)
    print(json_output)