import sys
import cv2
import json
import numpy as np
import pytesseract

# Path to the Tesseract executable
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

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 increase_contrast(image):
    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    cl = clahe.apply(l)
    limg = cv2.merge((cl, a, b))
    enhanced_image = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
    return enhanced_image

def extract_number_from_box(image, box):
    x1, y1, x2, y2 = box
    roi = image[y1:y2, x1:x2]

    try:
        # Upscale ROI for better OCR accuracy
        roi = cv2.resize(roi, None, fx=2, fy=2, interpolation=cv2.INTER_CUBIC)

        # Preprocess the ROI for better OCR results
        gray_roi = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
        gray_roi = cv2.GaussianBlur(gray_roi, (5, 5), 0)

        # Use adaptive thresholding for better binarization
        adaptive_thresh_roi = cv2.adaptiveThreshold(gray_roi, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)

        # Perform OCR on the preprocessed ROI
        custom_config = r'--oem 3 --psm 6'
        number = pytesseract.image_to_string(adaptive_thresh_roi, config=custom_config)
        number = ''.join(filter(str.isdigit, number))  # Extract only digits
        return number
    except Exception as e:
        print(f"OCR Error: {e}")
        return None

def draw_contours(image_path, colors_dict, output_path, min_contour_size=150, max_contour_size=50000):
    # 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)

    # Increase contrast of the source image
    source_image = increase_contrast(source_image)

    # 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)
            if w > 5 and h > 5:  # Only consider bounding boxes larger than a threshold
                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)

    # Extract numbers from each bounding box
    numbers_in_boxes = []
    for cnt, box in bounding_rectangles:
        center_x1 = box[0] + (box[2] - box[0]) // 4
        center_y1 = box[1] + (box[3] - box[1]) // 4
        center_x2 = center_x1 + (box[2] - box[0]) // 2
        center_y2 = center_y1 + (box[3] - box[1]) // 2
        center_box = (center_x1, center_y1, center_x2, center_y2)

        number = extract_number_from_box(source_image, center_box)
        if number:
            numbers_in_boxes.append(number)
            cv2.putText(source_image, number, (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)

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

    # Convert the array to JSON string
    json_string = json.dumps({"bounding_rectangles": bounding_rectangles, "numbers_in_boxes": numbers_in_boxes})

    # Print the JSON string
    print(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 = 'output_image_with_boxes123.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=50000)
    print(f"JSON output:\n{json_output}")