import cv2
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
import webcolors
from scipy.spatial import KDTree

import webcolors
from scipy.spatial import KDTree

def get_closest_color_name(rgb_tuple):
    """
    Finds the nearest human-readable color name for any given RGB tuple.
    Compatible with both newer and older versions of the 'webcolors' library.
    """
    # 1. Fetch color names and RGB mapping based on webcolors version
    try:
        # Modern webcolors syntax
        color_names = webcolors.names("css3")
        rgb_values = [webcolors.name_to_rgb(name, spec="css3") for name in color_names]
    except AttributeError:
        # Legacy webcolors fallback (v1.11 or older)
        css3_db = getattr(webcolors, "CSS3_HEX_TO_NAMES", webcolors.css3_hex_to_names)
        color_names = list(css3_db.values())
        rgb_values = [webcolors.hex_to_rgb(hex_code) for hex_code in css3_db.keys()]

    # 2. Query nearest neighbor using KDTree
    kdt_db = KDTree(rgb_values)
    _, index = kdt_db.query(rgb_tuple)
    return color_names[index]

def visualize_named_color_segmentation(image_path, num_colors=5, output_file="named_segmented_result.png"):
    """
    Segments an image by dominant colors, identifies human-readable color names,
    creates individual region masks, and plots a visual report.
    """
    # 1. Load image and convert to RGB
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"Could not load image at path: {image_path}")
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # 2. Reshape for K-Means
    pixels = image_rgb.reshape((-1, 3))
    
    # 3. Apply K-Means Clustering
    kmeans = KMeans(n_clusters=num_colors, n_init=10, random_state=42)
    labels = kmeans.fit_predict(pixels)
    colors = kmeans.cluster_centers_.astype(int)
    
    # Calculate pixel counts and percentages
    counts = np.bincount(labels)
    total_pixels = len(pixels)
    sorted_indices = np.argsort(counts)[::-1] # Sort by area size (largest first)
    
    # 4. Create the full segmented image (Quantized Image)
    segmented_pixels = colors[labels]
    segmented_image = segmented_pixels.reshape(image_rgb.shape)
    
    # 5. Build Subplot Grid for Visualization
    cols = 3
    rows = int(np.ceil((num_colors + 2) / cols))
    plt.figure(figsize=(16, 4.5 * rows))
    
    # Display Original Image
    plt.subplot(rows, cols, 1)
    plt.imshow(image_rgb)
    plt.title("Original Image", fontsize=12, fontweight='bold')
    plt.axis("off")
    
    # Display Full Color-Segmented Image
    plt.subplot(rows, cols, 2)
    plt.imshow(segmented_image)
    plt.title(f"Segmented Image ({num_colors} Colors)", fontsize=12, fontweight='bold')
    plt.axis("off")
    
    # Display Individual Color Region Masks with Color Names
    labels_2d = labels.reshape(image_rgb.shape[:2])
    
    for i, idx in enumerate(sorted_indices):
        cluster_color = tuple(colors[idx])
        percentage = (counts[idx] / total_pixels) * 100
        hex_code = f"#{cluster_color[0]:02x}{cluster_color[1]:02x}{cluster_color[2]:02x}"
        
        # Get human-readable color name
        color_name = get_closest_color_name(cluster_color).capitalize()
        
        # Create an isolated view for this specific color region
        region_mask = (labels_2d == idx)
        isolated_region = np.zeros_like(image_rgb)
        isolated_region[region_mask] = image_rgb[region_mask]
        
        # Plot individual segmented mask
        plt.subplot(rows, cols, i + 3)
        plt.imshow(isolated_region)
        plt.title(
            f"Region {i+1}: {color_name} ({percentage:.2f}%)\nRGB: {cluster_color} | HEX: {hex_code}", 
            fontsize=11, fontweight='bold'
        )
        plt.axis("off")
        
    plt.tight_layout()
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"Segmented output successfully saved to '{output_file}'")
    plt.show()

# --- Example Usage ---
if __name__ == "__main__":
    IMAGE_FILE = "/media/suman/Backup_of_extra_/Sasi/Flowers_images/imgR_nobg.png" # Replace with your image file path
    
    # Run segmentation with color name detection
    visualize_named_color_segmentation(IMAGE_FILE, num_colors=4, output_file="named_segmented_regions.png")