"""
This file generates a color legend image for building component visualization.
It creates a PNG image showing color swatches and labels for two categories:
1. Gestalt Colors - for various building components like roof, walls, windows, etc.
2. Edge Colors - for architectural edges like ridges, eaves, hips, valleys, etc.
The legend helps visualize the color mappings used in building analysis and annotation.
"""

import cv2
import numpy as np

# Color mappings (as provided in the context file)
gestalt_color_mapping = {
    "unclassified": (215, 62, 138),
    "apex": (235, 88, 48),
    "eave_end_point": (248, 130, 228),
    "flashing_end_point": (71, 11, 161),
    "ridge": (214, 251, 248),
    "rake": (13, 94, 47),
    "eave": (54, 243, 63),
    "post": (187, 123, 236),
    "ground_line": (136, 206, 14),
    "flashing": (162, 162, 32),
    "step_flashing": (169, 255, 219),
    "hip": (8, 89, 52),
    "valley": (85, 27, 65),
    "roof": (215, 232, 179),
    "door": (110, 52, 23),
    "garage": (50, 233, 171),
    "window": (230, 249, 40),
    "shutter": (122, 4, 233),
    "fascia": (95, 230, 240),
    "soffit": (2, 102, 197),
    "horizontal_siding": (131, 88, 59),
    "vertical_siding": (110, 187, 198),
    "brick": (171, 252, 7),
    "concrete": (32, 47, 246),
    "other_wall": (112, 61, 240),
    "trim": (151, 206, 58),
    "unknown": (127, 127, 127),
    "transition_line": (0,0,0),
}

edge_color_mapping = {
    'cornice_return': (215, 62, 138),
    'cornice_strip': (235, 88, 48),
    'eave':  (54, 243, 63),
    "flashing": (162, 162, 32),
    'hip': (8, 89, 52),
    'rake': (13, 94, 47),
    'ridge': (214, 251, 248),
    "step_flashing": (169, 255, 219),
    'transition_line': (200,0,50),
    'valley': (85, 27, 65),
}

# Parameters for the legend image
swatch_width = 150
swatch_height = 25
text_color_bgr = (0, 0, 0)  # Black
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = 0.4
font_thickness = 1
horizontal_padding = 20
vertical_padding = 15 # Padding at top/bottom of sections and between items
text_offset_x = 10 # Space between swatch and text
item_row_height = swatch_height + 5 # Total height for one swatch + text row, including small spacing
section_title_font_scale = 0.6
section_title_font_thickness = 1
section_title_area_height = 40 # Space allocated for section title

all_mappings_with_titles = [
    ("Gestalt Colors", gestalt_color_mapping),
    ("Edge Colors", edge_color_mapping),
]

# Calculate maximum width needed for labels to determine overall image width
max_label_width = 0
for _, color_map in all_mappings_with_titles:
    for label in color_map.keys():
        (text_w, _), _ = cv2.getTextSize(label, font, font_scale, font_thickness)
        if text_w > max_label_width:
            max_label_width = text_w

image_width = horizontal_padding + swatch_width + text_offset_x + max_label_width + horizontal_padding

# Calculate total image height
image_height = vertical_padding # Initial top padding
for i, (_, color_map) in enumerate(all_mappings_with_titles):
    image_height += section_title_area_height
    image_height += len(color_map) * item_row_height
    if i < len(all_mappings_with_titles) - 1:
        image_height += vertical_padding # Padding between sections
image_height += vertical_padding # Final bottom padding

# Create a white canvas
legend_image = np.full((image_height, image_width, 3), 255, dtype=np.uint8) # White background

current_y = vertical_padding

for title_text, color_map in all_mappings_with_titles:
    # Draw section title
    (title_w, title_h), title_baseline = cv2.getTextSize(title_text, font, section_title_font_scale, section_title_font_thickness)
    title_x = horizontal_padding # Align title to the left
    # Center title vertically within its allocated space
    title_y_baseline = current_y + (section_title_area_height - title_h) // 2 + title_h - title_baseline // 2
    cv2.putText(legend_image, title_text, (title_x, title_y_baseline), font, section_title_font_scale, text_color_bgr, section_title_font_thickness)
    current_y += section_title_area_height

    # Draw color swatches and labels for the current map
    for label, rgb_color in color_map.items():
        # OpenCV uses BGR, so convert RGB to BGR
        bgr_color = (rgb_color[2], rgb_color[1], rgb_color[0])

        swatch_x1 = horizontal_padding
        swatch_y1 = current_y
        swatch_x2 = swatch_x1 + swatch_width
        swatch_y2 = swatch_y1 + swatch_height

        # Draw the color swatch (filled rectangle)
        cv2.rectangle(legend_image, (swatch_x1, swatch_y1), (swatch_x2, swatch_y2), bgr_color, -1)
        # Draw a thin black border around the swatch for clarity
        cv2.rectangle(legend_image, (swatch_x1, swatch_y1), (swatch_x2, swatch_y2), (0,0,0), 1)


        # Get text size for precise vertical centering
        (text_w, text_h), baseline = cv2.getTextSize(label, font, font_scale, font_thickness)
        
        # Calculate position for the text (vertically centered with the swatch)
        text_x = swatch_x2 + text_offset_x
        # The 'org' for putText is the bottom-left corner of the text string.
        # To center text_h within swatch_height:
        text_y_baseline = swatch_y1 + (swatch_height - text_h) // 2 + text_h
        
        cv2.putText(legend_image, label, (text_x, text_y_baseline), font, font_scale, text_color_bgr, font_thickness)
        
        current_y += item_row_height # Move to the next item row
    
    current_y += vertical_padding # Add padding after the items in this section

# Save the image
output_filename = "gestalt_color_legend.png"
try:
    cv2.imwrite(output_filename, legend_image)
    print(f"Legend image saved as {output_filename}")
except Exception as e:
    print(f"Error saving image: {e}")

# To display the image (optional, uncomment if you have a display environment)
# cv2.imshow("Gestalt Color Legend", legend_image)
# cv2.waitKey(0)
# cv2.destroyAllWindows()