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/OneDrive - PTW TU Darmstadt/Dissertation/NFAG/polygon_map_test/GCSM"

# Define a directory where you have write access for the output CSV
output_directory = "C:/Users/O.Oezen/OneDrive - PTW TU Darmstadt/Dissertation/NFAG/polygon_map_test/GCSM_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 factory
    (0, 192, 0),        # green ecological
    (255, 255, 255),    # white roads
    (210, 104, 165),    # purple storage
    (5, 255, 243),      # yellow social
    (210, 0, 0),        # blue energy supply systems
    (54, 100, 184),      # brown office
    (128, 128, 128),     # gray unused
]

# Corresponding names for the colors
color_names = ["red factory", "green ecological", "white roads", "purple storage", "yellow social", "blue energy", "brown office", "gray unused"]

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')):
        filepath = os.path.join(input_directory, file)
        img = cv2.imread(filepath)

        if img is None:
            print(f"Error loading image: {filepath}")
            continue

        for target_color, color_name in zip(target_colors, color_names):
            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)

            mask_filename = os.path.join(output_directory, f'mask_{color_name}_{file}')
            cv2.imwrite(mask_filename, mask)

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

            filenames.append(file)
            colorcode.append(target_color)
            colorname.append(color_name)
            percentages.append(np.round(percentage, 5))

# Create DataFrame
df = pd.DataFrame({
    "filename": filenames,
    "colorcode": colorcode,
    "colorname": colorname,
    "percentage": percentages
})

# Calculate scaled_percentage per image to account for irregular shapes, as these sections are put in as black areas
df["scaled_percentage"] = df.groupby("filename")["percentage"].transform(lambda x: (x / x.sum()) * 100)
df["scaled_percentage"] = df["scaled_percentage"].round(2)

# Save the DataFrame to CSV
df.to_csv(filename, index=False, sep=';')

print(f"{filename} has been created with headers 'filename', 'colorcode', 'colorname', 'percentage', and 'scaled_percentage'.")