import os
import numpy as np
import pandas as pd
import cv2  # Import OpenCV

# Define the directory where you want to scan for image files
input_directory = "C:/Users/O.Oezen/Documents/polygon_map_test"  # Change to your actual directory

# Define a directory where you have write access for the output CSV
output_directory = "C:/Users/O.Oezen/Documents/polygon_map_test/Output"  
os.makedirs(output_directory, exist_ok=True)

# Define the full path for the output file
filename = os.path.join(output_directory, 'Map_Color_Percentage.csv')

# Define the target colors (in BGR format)
target_colors = [
    (0, 0, 192),      # red
    (46, 167, 78),    # green
    (0, 0, 0),        # white
    (160, 48, 112),   # purple
    (9, 255, 243)     # yellow
]
color_names = ["red", "green", "white", "purple", "yellow"]  # Corresponding names for the colors
color_range = 50  # Range of color matching

# Initialize lists to hold the data
filenames = []
colorcode = []
colorname = []
percentages = []

# Loop through files in the specified input directory
for file in os.listdir(input_directory):
    if file.lower().endswith(('.jpg', '.jpeg', '.png')):  # Check for image file extensions
        # Construct the full path to the image
        filepath = os.path.join(input_directory, file)  # Use input_directory

        # Load the image
        img = cv2.imread(filepath)

        # Check if the image was loaded successfully
        if img is None:
            print(f"Error loading image: {filepath}")
            continue  # Skip to the next file if there is an error

        # Loop through each target color
        for target_color, color_name in zip(target_colors, color_names):
            # Calculate the percentage of the image that is within the color range of the target 
            lower_bound = np.clip(np.array(target_color) - color_range, 0, 255)
            upper_bound = np.clip(np.array(target_color) + color_range, 0, 255)
            mask = cv2.inRange(img, lower_bound, upper_bound)

            # Save the mask for the current image with a unique filename
            mask_filename = os.path.join(output_directory, f'mask_{color_name}_{file}')  # Unique mask filename
            cv2.imwrite(mask_filename, mask)  # Save the mask

            num_pixels = img.shape[0] * img.shape[1]
            target_pixels = cv2.countNonZero(mask)
            percentage = (target_pixels / num_pixels) * 100

            # Append data to lists
            filenames.append(file)  # Add the file name to the list
            colorcode.append(target_color)  # Add color code to the list
            colorname.append(color_name)  # Get the name for the target color
            percentages.append(np.round(percentage, 2))  # Append calculated percentage

# Create a DataFrame from the collected data
data = {
    "filename": filenames,
    "colorcode": colorcode,
    "colorname": colorname,
    "percentage": percentages
}

df = pd.DataFrame(data)  # Create DataFrame using the data dictionary

# Save the DataFrame to a CSV file
df.to_csv(filename, index=False)

print(f"{filename} has been created with headers 'filename', 'colorcode', 'colorname', and 'percentage' and populated with data from the directory.")