#%%
import numpy as np
import cv2
import matplotlib.pyplot as plt
from PIL import Image

def merge_intervals(intervals):
    # Merge overlapping intervals.
    if not intervals:
        return []
    intervals.sort(key=lambda x: x[0])
    merged = [intervals[0]]
    for current in intervals[1:]:
        last = merged[-1]
        if current[0] <= last[1]:
            merged[-1] = (last[0], max(last[1], current[1]))
        else:
            merged.append(current)
    return merged

def get_allowed_intervals(neighbor_values, delta):
    # Compute forbidden intervals for this channel from neighbor values.
    forbidden = []
    for c in neighbor_values:
        low = max(0, c - delta)
        high = min(1, c + delta)
        forbidden.append((low, high))
    # Merge overlapping forbidden intervals.
    merged = merge_intervals(forbidden)
    allowed = []
    # Allowed region before the first forbidden interval.
    if merged:
        if merged[0][0] > 0:
            allowed.append((0, merged[0][0]))
        # Allowed regions between forbidden intervals.
        for i in range(1, len(merged)):
            allowed.append((merged[i-1][1], merged[i][0]))
        # Allowed region after the last forbidden interval.
        if merged[-1][1] < 1:
            allowed.append((merged[-1][1], 1))
    else:
        allowed.append((0, 1))
    return allowed

def sample_from_allowed(allowed):
    # Choose a random value from the allowed intervals with probability proportional to interval length.
    lengths = [end - start for start, end in allowed]
    total_length = sum(lengths)
    if total_length <= 0:
        return None
    point = np.random.rand() * total_length
    cumulative = 0
    for (start, end), length in zip(allowed, lengths):
        cumulative += length
        if point <= cumulative:
            return np.random.uniform(start, end)
    return allowed[-1][0]

def sample_channel(neighbor_values, delta):
    # Compute allowed intervals and sample a value.
    allowed = get_allowed_intervals(neighbor_values, delta)
    if not allowed:
        return None
    return sample_from_allowed(allowed)

def generate_valid_color(neighbors, r):
    """
    Generate a random color that satisfies, for each neighbor,
    |R - R_neighbor| >= sqrt(r), and similarly for G and B.
    neighbors: list of neighbor colors (each is an array of shape (3,))
    r: threshold for squared difference.
    """
    delta = np.sqrt(r)
    color = np.zeros(3)
    for channel in range(3):
        neighbor_channel = [n[channel] for n in neighbors]
        sampled = sample_channel(neighbor_channel, delta)
        # If the allowed intervals are empty, use a fallback.
        if sampled is None:
            sampled = np.random.rand()
        color[channel] = sampled
    return color

def generate_color_pattern(h, w, d, r):
    """
    Generate a dense color pattern.
    h, w: height and width of the output image in pixels.
    d: patch size in pixels (each patch is d x d).
    r: randomness regulation factor. For each neighbor the condition is:
       (|channel_new - channel_neighbor|)^2 >= r.
    """
    pattern = np.zeros((h, w, 3), dtype=np.float32)
    num_patches_x = h // d
    num_patches_y = w // d

    for i in range(num_patches_x):
        for j in range(num_patches_y):
            # Collect neighbor patch colors (using mean because patches have uniform color).
            neighbors = []
            if i > 0:
                patch_above = pattern[(i-1)*d:i*d, j*d:(j+1)*d]
                neighbors.append(np.mean(patch_above, axis=(0, 1)))
            if j > 0:
                patch_left = pattern[i*d:(i+1)*d, (j-1)*d:j*d]
                neighbors.append(np.mean(patch_left, axis=(0, 1)))
                
            if len(neighbors) == 0:
                # For the first patch use a random color.
                new_color = np.random.rand(3)
            else:
                new_color = generate_valid_color(neighbors, r)
                
            pattern[i*d:(i+1)*d, j*d:(j+1)*d] = new_color
            # print('Patch ({}, {})'.format(i, j))
    return pattern

def convert_to_grayscale(color_pattern):
    """
    Convert color pattern to grayscale using standard luminance weights.
    
    Parameters:
    color_pattern: Input color image (H x W x 3) with values in range [0, 1]
    
    Returns:
    Grayscale image (H x W) with values in range [0, 1]
    """
    # Use standard RGB to grayscale conversion: 0.299*R + 0.587*G + 0.114*B
    grayscale = 0.299 * color_pattern[:, :, 2] + 0.587 * color_pattern[:, :, 1] + 0.114 * color_pattern[:, :, 0]
    return grayscale

def optimize_grayscale_patches(grayscale, d, min_diff=0.1, max_iterations=100):
    """
    Optimize grayscale image to improve value difference between neighboring patches.
    
    Parameters:
    grayscale: Input grayscale image (H x W) with values in range [0, 1]
    d: patch size in pixels (each patch is d x d)
    min_diff: minimum absolute difference required between neighboring patches
    max_iterations: maximum number of optimization iterations
    
    Returns:
    Optimized grayscale image with improved patch differences
    """
    h, w = grayscale.shape
    num_patches_x = h // d
    num_patches_y = w // d
    
    # Create optimized version
    optimized = grayscale.copy()
    
    for iteration in range(max_iterations):
        improved = False
        
        for i in range(num_patches_x):
            for j in range(num_patches_y):
                # Get current patch value (mean of patch)
                current_patch = optimized[i*d:(i+1)*d, j*d:(j+1)*d]
                current_value = np.mean(current_patch)
                
                # Collect neighbor patch values
                neighbor_values = []
                if i > 0:
                    patch_above = optimized[(i-1)*d:i*d, j*d:(j+1)*d]
                    neighbor_values.append(np.mean(patch_above))
                if j > 0:
                    patch_left = optimized[i*d:(i+1)*d, (j-1)*d:j*d]
                    neighbor_values.append(np.mean(patch_left))
                if i < num_patches_x - 1:
                    patch_below = optimized[(i+1)*d:(i+2)*d, j*d:(j+1)*d]
                    neighbor_values.append(np.mean(patch_below))
                if j < num_patches_y - 1:
                    patch_right = optimized[i*d:(i+1)*d, (j+1)*d:(j+2)*d]
                    neighbor_values.append(np.mean(patch_right))
                
                if len(neighbor_values) == 0:
                    continue
                
                # Check if current value is too close to any neighbor
                min_neighbor_diff = min([abs(current_value - nv) for nv in neighbor_values])
                
                if min_neighbor_diff < min_diff:
                    # Need to adjust this patch
                    # Find valid range that maintains min_diff from all neighbors
                    forbidden = []
                    for nv in neighbor_values:
                        forbidden.append((max(0, nv - min_diff), min(1, nv + min_diff)))
                    
                    # Merge forbidden intervals
                    merged_forbidden = merge_intervals(forbidden)
                    
                    # Find allowed intervals
                    allowed = []
                    if merged_forbidden:
                        if merged_forbidden[0][0] > 0:
                            allowed.append((0, merged_forbidden[0][0]))
                        for k in range(1, len(merged_forbidden)):
                            gap = merged_forbidden[k][0] - merged_forbidden[k-1][1]
                            if gap > 0:
                                allowed.append((merged_forbidden[k-1][1], merged_forbidden[k][0]))
                        if merged_forbidden[-1][1] < 1:
                            allowed.append((merged_forbidden[-1][1], 1))
                    else:
                        allowed.append((0, 1))
                    
                    # Sample new value from allowed intervals
                    if allowed:
                        new_value = sample_from_allowed(allowed)
                        if new_value is not None:
                            optimized[i*d:(i+1)*d, j*d:(j+1)*d] = new_value
                            improved = True
        
        if not improved:
            print(f"Converged after {iteration + 1} iterations")
            break
    
    return optimized

def analyze_patch_differences(image, d):
    """
    Analyze and report the minimum differences between neighboring patches.
    
    Parameters:
    image: Grayscale image (H x W) with values in range [0, 1]
    d: patch size in pixels
    
    Returns:
    Dictionary with statistics about patch differences
    """
    h, w = image.shape
    num_patches_x = h // d
    num_patches_y = w // d
    
    horizontal_diffs = []
    vertical_diffs = []
    
    for i in range(num_patches_x):
        for j in range(num_patches_y):
            current_patch = image[i*d:(i+1)*d, j*d:(j+1)*d]
            current_value = np.mean(current_patch)
            
            # Check horizontal neighbor (right)
            if j < num_patches_y - 1:
                patch_right = image[i*d:(i+1)*d, (j+1)*d:(j+2)*d]
                right_value = np.mean(patch_right)
                horizontal_diffs.append(abs(current_value - right_value))
            
            # Check vertical neighbor (below)
            if i < num_patches_x - 1:
                patch_below = image[(i+1)*d:(i+2)*d, j*d:(j+1)*d]
                below_value = np.mean(patch_below)
                vertical_diffs.append(abs(current_value - below_value))
    
    stats = {
        'min_diff': min(horizontal_diffs + vertical_diffs) if (horizontal_diffs + vertical_diffs) else 0,
        'max_diff': max(horizontal_diffs + vertical_diffs) if (horizontal_diffs + vertical_diffs) else 0,
        'mean_diff': np.mean(horizontal_diffs + vertical_diffs) if (horizontal_diffs + vertical_diffs) else 0,
        'median_diff': np.median(horizontal_diffs + vertical_diffs) if (horizontal_diffs + vertical_diffs) else 0,
    }
    
    return stats

def rgb_to_cmyk(rgb_image):
    """
    Convert RGB image to CMYK color space for printing.
    
    Parameters:
    rgb_image: RGB image (H x W x 3) with values in range [0, 1], BGR order (OpenCV format)
    
    Returns:
    CMYK image (H x W x 4) with values in range [0, 1]
    """
    # OpenCV uses BGR, so we need to handle it as B, G, R in channels 0, 1, 2
    b = rgb_image[:, :, 0]
    g = rgb_image[:, :, 1]
    r = rgb_image[:, :, 2]
    
    # Calculate K (black) channel
    k = 1 - np.maximum(np.maximum(r, g), b)
    
    # Avoid division by zero
    k_inv = 1 - k
    k_inv = np.where(k_inv == 0, 1e-10, k_inv)
    
    # Calculate CMY channels
    c = (1 - r - k) / k_inv
    m = (1 - g - k) / k_inv
    y = (1 - b - k) / k_inv
    
    # Stack into CMYK image
    cmyk = np.stack([c, m, y, k], axis=2)
    cmyk = np.clip(cmyk, 0, 1)
    
    return cmyk

def cmyk_to_rgb(cmyk_image):
    """
    Convert CMYK image back to RGB color space.
    
    Parameters:
    cmyk_image: CMYK image (H x W x 4) with values in range [0, 1]
    
    Returns:
    RGB image (H x W x 3) with values in range [0, 1], BGR order (OpenCV format)
    """
    c = cmyk_image[:, :, 0]
    m = cmyk_image[:, :, 1]
    y = cmyk_image[:, :, 2]
    k = cmyk_image[:, :, 3]
    
    # Convert back to RGB
    r = (1 - c) * (1 - k)
    g = (1 - m) * (1 - k)
    b = (1 - y) * (1 - k)
    
    # Stack in BGR order for OpenCV
    rgb = np.stack([b, g, r], axis=2)
    rgb = np.clip(rgb, 0, 1)
    
    return rgb

def optimize_cmyk_patches(cmyk_image, d, min_diff=0.1, max_iterations=100):
    """
    Optimize CMYK image to improve color difference between neighboring patches.
    Uses Euclidean distance in CMYK space.
    
    Parameters:
    cmyk_image: Input CMYK image (H x W x 4) with values in range [0, 1]
    d: patch size in pixels (each patch is d x d)
    min_diff: minimum Euclidean distance required between neighboring patches in CMYK space
    max_iterations: maximum number of optimization iterations
    
    Returns:
    Optimized CMYK image with improved patch differences
    """
    h, w = cmyk_image.shape[:2]
    num_patches_x = h // d
    num_patches_y = w // d
    
    optimized = cmyk_image.copy()
    
    for iteration in range(max_iterations):
        improved = False
        
        for i in range(num_patches_x):
            for j in range(num_patches_y):
                # Get current patch CMYK values (mean of patch)
                current_patch = optimized[i*d:(i+1)*d, j*d:(j+1)*d]
                current_cmyk = np.mean(current_patch, axis=(0, 1))
                
                # Collect neighbor patch CMYK values
                neighbor_cmyk = []
                if i > 0:
                    patch_above = optimized[(i-1)*d:i*d, j*d:(j+1)*d]
                    neighbor_cmyk.append(np.mean(patch_above, axis=(0, 1)))
                if j > 0:
                    patch_left = optimized[i*d:(i+1)*d, (j-1)*d:j*d]
                    neighbor_cmyk.append(np.mean(patch_left, axis=(0, 1)))
                if i < num_patches_x - 1:
                    patch_below = optimized[(i+1)*d:(i+2)*d, j*d:(j+1)*d]
                    neighbor_cmyk.append(np.mean(patch_below, axis=(0, 1)))
                if j < num_patches_y - 1:
                    patch_right = optimized[i*d:(i+1)*d, (j+1)*d:(j+2)*d]
                    neighbor_cmyk.append(np.mean(patch_right, axis=(0, 1)))
                
                if len(neighbor_cmyk) == 0:
                    continue
                
                # Calculate minimum Euclidean distance to neighbors in CMYK space
                distances = [np.linalg.norm(current_cmyk - n) for n in neighbor_cmyk]
                min_distance = min(distances)
                
                if min_distance < min_diff:
                    # Need to adjust this patch - try to move it away from neighbors
                    # Find the direction away from the closest neighbor
                    closest_idx = np.argmin(distances)
                    closest_neighbor = neighbor_cmyk[closest_idx]
                    
                    # Move away from closest neighbor
                    direction = current_cmyk - closest_neighbor
                    direction_norm = np.linalg.norm(direction)
                    
                    if direction_norm > 0:
                        direction = direction / direction_norm
                        # Move by the deficit amount
                        move_amount = (min_diff - min_distance) * 1.2  # 1.2 factor for faster convergence
                        new_cmyk = current_cmyk + direction * move_amount
                        new_cmyk = np.clip(new_cmyk, 0, 1)
                        
                        # Check if new value is better
                        new_distances = [np.linalg.norm(new_cmyk - n) for n in neighbor_cmyk]
                        new_min_distance = min(new_distances)
                        
                        if new_min_distance > min_distance:
                            optimized[i*d:(i+1)*d, j*d:(j+1)*d] = new_cmyk
                            improved = True
        
        if not improved:
            print(f"CMYK optimization converged after {iteration + 1} iterations")
            break
    
    return optimized

def analyze_cmyk_patch_differences(cmyk_image, d):
    """
    Analyze CMYK patch differences using Euclidean distance in CMYK space.
    
    Parameters:
    cmyk_image: CMYK image (H x W x 4) with values in range [0, 1]
    d: patch size in pixels
    
    Returns:
    Dictionary with statistics about CMYK patch differences
    """
    h, w = cmyk_image.shape[:2]
    num_patches_x = h // d
    num_patches_y = w // d
    
    horizontal_diffs = []
    vertical_diffs = []
    
    for i in range(num_patches_x):
        for j in range(num_patches_y):
            current_patch = cmyk_image[i*d:(i+1)*d, j*d:(j+1)*d]
            current_cmyk = np.mean(current_patch, axis=(0, 1))
            
            # Check horizontal neighbor (right)
            if j < num_patches_y - 1:
                patch_right = cmyk_image[i*d:(i+1)*d, (j+1)*d:(j+2)*d]
                right_cmyk = np.mean(patch_right, axis=(0, 1))
                horizontal_diffs.append(np.linalg.norm(current_cmyk - right_cmyk))
            
            # Check vertical neighbor (below)
            if i < num_patches_x - 1:
                patch_below = cmyk_image[(i+1)*d:(i+2)*d, j*d:(j+1)*d]
                below_cmyk = np.mean(patch_below, axis=(0, 1))
                vertical_diffs.append(np.linalg.norm(current_cmyk - below_cmyk))
    
    all_diffs = horizontal_diffs + vertical_diffs
    stats = {
        'min_diff': min(all_diffs) if all_diffs else 0,
        'max_diff': max(all_diffs) if all_diffs else 0,
        'mean_diff': np.mean(all_diffs) if all_diffs else 0,
        'median_diff': np.median(all_diffs) if all_diffs else 0,
    }
    
    return stats

def grayscale_to_cmyk_optimized(grayscale_target, d, min_cmyk_diff=0.2, max_iterations=100, color_diversity=0.3):
    """
    Generate a CMYK image that when converted to grayscale matches the target grayscale,
    while maintaining good color separation between patches.
    
    Parameters:
    grayscale_target: Target grayscale image (H x W) with values in range [0, 1]
    d: patch size in pixels
    min_cmyk_diff: minimum Euclidean distance required between neighboring patches in CMYK space
    max_iterations: maximum number of optimization iterations
    color_diversity: amount of color diversity to inject (0.0 to 1.0)
    
    Returns:
    CMYK image (H x W x 4) optimized for printing
    """
    h, w = grayscale_target.shape
    num_patches_x = h // d
    num_patches_y = w // d
    
    # Initialize CMYK image
    cmyk_image = np.zeros((h, w, 4), dtype=np.float32)
    
    # Generate initial CMYK values for each patch
    for i in range(num_patches_x):
        for j in range(num_patches_y):
            # Get target grayscale value for this patch
            target_gray = np.mean(grayscale_target[i*d:(i+1)*d, j*d:(j+1)*d])
            
            # Generate random CMY values with diversity
            c = np.random.rand() * color_diversity
            m = np.random.rand() * color_diversity
            y = np.random.rand() * color_diversity
            
            # Calculate K based on target grayscale
            # Grayscale from CMYK: gray ≈ 1 - K when CMY are small
            # For more accurate: gray = (1-C)*(1-K)*0.114 + (1-M)*(1-K)*0.587 + (1-Y)*(1-K)*0.299
            k = 1 - target_gray  # Initial approximation
            
            # Adjust to match target grayscale
            # Create temporary CMYK and test
            temp_cmyk = np.array([c, m, y, k])
            temp_rgb = np.array([(1-c)*(1-k), (1-m)*(1-k), (1-y)*(1-k)])  # R, G, B
            temp_gray = 0.299 * temp_rgb[0] + 0.587 * temp_rgb[1] + 0.114 * temp_rgb[2]
            
            # Adjust K to match target
            k_adjustment = target_gray - temp_gray
            k = np.clip(k - k_adjustment, 0, 1)
            
            cmyk_image[i*d:(i+1)*d, j*d:(j+1)*d] = [c, m, y, k]
    
    # Optimize CMYK to improve patch differences while maintaining grayscale target
    print("Optimizing CMYK for color separation while matching grayscale...")
    for iteration in range(max_iterations):
        improved = False
        
        for i in range(num_patches_x):
            for j in range(num_patches_y):
                # Get current patch
                current_patch = cmyk_image[i*d:(i+1)*d, j*d:(j+1)*d]
                current_cmyk = np.mean(current_patch, axis=(0, 1))
                
                # Get target grayscale for this patch
                target_gray = np.mean(grayscale_target[i*d:(i+1)*d, j*d:(j+1)*d])
                
                # Collect neighbor CMYK values
                neighbor_cmyk = []
                if i > 0:
                    neighbor_cmyk.append(np.mean(cmyk_image[(i-1)*d:i*d, j*d:(j+1)*d], axis=(0, 1)))
                if j > 0:
                    neighbor_cmyk.append(np.mean(cmyk_image[i*d:(i+1)*d, (j-1)*d:j*d], axis=(0, 1)))
                if i < num_patches_x - 1:
                    neighbor_cmyk.append(np.mean(cmyk_image[(i+1)*d:(i+2)*d, j*d:(j+1)*d], axis=(0, 1)))
                if j < num_patches_y - 1:
                    neighbor_cmyk.append(np.mean(cmyk_image[i*d:(i+1)*d, (j+1)*d:(j+2)*d], axis=(0, 1)))
                
                if len(neighbor_cmyk) == 0:
                    continue
                
                # Check CMYK distance to neighbors
                distances = [np.linalg.norm(current_cmyk - n) for n in neighbor_cmyk]
                min_distance = min(distances)
                
                if min_distance < min_cmyk_diff:
                    # Try to improve color separation
                    # Vary CMY channels while keeping grayscale constant
                    for attempt in range(5):
                        # Generate new CMY values
                        new_c = np.clip(current_cmyk[0] + (np.random.rand() - 0.5) * 0.2, 0, 1)
                        new_m = np.clip(current_cmyk[1] + (np.random.rand() - 0.5) * 0.2, 0, 1)
                        new_y = np.clip(current_cmyk[2] + (np.random.rand() - 0.5) * 0.2, 0, 1)
                        
                        # Calculate K to maintain grayscale
                        # gray = (1-C)*(1-K)*0.114 + (1-M)*(1-K)*0.587 + (1-Y)*(1-K)*0.299
                        # gray = (1-K) * [0.114*(1-C) + 0.587*(1-M) + 0.299*(1-Y)]
                        rgb_factor = 0.114*(1-new_c) + 0.587*(1-new_m) + 0.299*(1-new_y)
                        if rgb_factor > 0:
                            new_k = 1 - (target_gray / rgb_factor)
                            new_k = np.clip(new_k, 0, 1)
                        else:
                            new_k = current_cmyk[3]
                        
                        new_cmyk = np.array([new_c, new_m, new_y, new_k])
                        
                        # Verify grayscale match
                        test_rgb = np.array([(1-new_c)*(1-new_k), (1-new_m)*(1-new_k), (1-new_y)*(1-new_k)])
                        test_gray = 0.299 * test_rgb[0] + 0.587 * test_rgb[1] + 0.114 * test_rgb[2]
                        
                        # Check if grayscale is close enough and CMYK distance improved
                        if abs(test_gray - target_gray) < 0.05:  # Grayscale tolerance
                            new_distances = [np.linalg.norm(new_cmyk - n) for n in neighbor_cmyk]
                            new_min_distance = min(new_distances)
                            
                            if new_min_distance > min_distance:
                                cmyk_image[i*d:(i+1)*d, j*d:(j+1)*d] = new_cmyk
                                improved = True
                                break
        
        if not improved:
            print(f"Grayscale-to-CMYK optimization converged after {iteration + 1} iterations")
            break
    
    return cmyk_image

def cmyk_to_grayscale(cmyk_image):
    """
    Convert CMYK image to grayscale.
    
    Parameters:
    cmyk_image: CMYK image (H x W x 4) with values in range [0, 1]
    
    Returns:
    Grayscale image (H x W) with values in range [0, 1]
    """
    # Convert to RGB first, then to grayscale
    rgb = cmyk_to_rgb(cmyk_image)
    grayscale = 0.299 * rgb[:, :, 2] + 0.587 * rgb[:, :, 1] + 0.114 * rgb[:, :, 0]
    return grayscale

def main():
    # Parameters for the pattern.
    # Image dimensions in millimeters
    width_mm = 50  # Width in millimeters
    height_mm = 50  # Height in millimeters
    dpi = 600  # Dots per inch for printing
    patch_size = 1 # 3 for 300dpi, 6 for 600dpi
    
    # Convert mm to pixels: pixels = (mm / 25.4) * dpi
    width = int((width_mm / 25.4) * dpi)
    height = int((height_mm / 25.4) * dpi)
    
    print(f"Image dimensions: {width_mm}mm x {height_mm}mm at {dpi} DPI")
    print(f"Pixel dimensions: {width} x {height} pixels")
    
    
    # Set r to a value that gives sufficient difference.
    # Note: r is the threshold for squared difference per channel, and the condition is
    # |c - c_neighbor| >= sqrt(r). Very high r may make it impossible to satisfy.
    r = 1

    # Step 1: Generate initial color pattern
    print("Generating color pattern...")
    pattern = generate_color_pattern(height, width, patch_size, r)
    pattern_display = (pattern * 255).astype(np.uint8)
    
    # Step 2: Convert to grayscale
    print("\nConverting to grayscale...")
    grayscale = convert_to_grayscale(pattern)
    grayscale_display = (grayscale * 255).astype(np.uint8)
    
    # Step 3: Analyze original grayscale differences
    print("\nAnalyzing original grayscale patch differences...")
    original_stats = analyze_patch_differences(grayscale, patch_size)
    print(f"Original - Min: {original_stats['min_diff']:.4f}, Max: {original_stats['max_diff']:.4f}, "
          f"Mean: {original_stats['mean_diff']:.4f}, Median: {original_stats['median_diff']:.4f}")
    
    # Step 4: Optimize grayscale to improve patch differences
    print("\nOptimizing grayscale patch differences...")
    min_diff_threshold = 0.15  # Adjust this value to control minimum difference
    optimized_grayscale = optimize_grayscale_patches(grayscale, patch_size, min_diff=min_diff_threshold, max_iterations=50)
    optimized_display = (optimized_grayscale * 255).astype(np.uint8)
    
    # Step 5: Analyze optimized grayscale differences
    print("\nAnalyzing optimized grayscale patch differences...")
    optimized_stats = analyze_patch_differences(optimized_grayscale, patch_size)
    print(f"Optimized - Min: {optimized_stats['min_diff']:.4f}, Max: {optimized_stats['max_diff']:.4f}, "
          f"Mean: {optimized_stats['mean_diff']:.4f}, Median: {optimized_stats['median_diff']:.4f}")
    
    # Step 6: Generate CMYK image from optimized grayscale
    print("\n" + "="*60)
    print("GENERATING CMYK FROM OPTIMIZED GRAYSCALE")
    print("="*60)
    print("\nGenerating CMYK image to match optimized grayscale...")
    cmyk_from_grayscale = grayscale_to_cmyk_optimized(
        optimized_grayscale, 
        patch_size, 
        min_cmyk_diff=0.2, 
        max_iterations=50,
        color_diversity=0.5
    )
    
    # Analyze CMYK differences
    print("\nAnalyzing CMYK patch differences...")
    cmyk_stats = analyze_cmyk_patch_differences(cmyk_from_grayscale, patch_size)
    print(f"CMYK - Min: {cmyk_stats['min_diff']:.4f}, Max: {cmyk_stats['max_diff']:.4f}, "
          f"Mean: {cmyk_stats['mean_diff']:.4f}, Median: {cmyk_stats['median_diff']:.4f}")
    
    # Convert CMYK back to RGB for display
    rgb_from_cmyk = cmyk_to_rgb(cmyk_from_grayscale)
    rgb_from_cmyk_display = (rgb_from_cmyk * 255).astype(np.uint8)
    
    # Convert CMYK back to grayscale to verify
    grayscale_from_cmyk = cmyk_to_grayscale(cmyk_from_grayscale)
    grayscale_from_cmyk_display = (grayscale_from_cmyk * 255).astype(np.uint8)
    
    # Analyze how well grayscale was preserved
    grayscale_from_cmyk_stats = analyze_patch_differences(grayscale_from_cmyk, patch_size)
    print(f"Grayscale from CMYK - Min: {grayscale_from_cmyk_stats['min_diff']:.4f}")
    
    # Calculate grayscale reconstruction error
    gray_error = np.mean(np.abs(optimized_grayscale - grayscale_from_cmyk))
    print(f"Grayscale reconstruction error: {gray_error:.4f}")
    
    # Display all images with matplotlib
    fig, axes = plt.subplots(2, 3, figsize=(18, 12), dpi=300)
    
    # Row 1: Original grayscale optimization
    # Color pattern (convert BGR to RGB for matplotlib)
    pattern_rgb = cv2.cvtColor(pattern_display, cv2.COLOR_BGR2RGB)
    axes[0, 0].imshow(pattern_rgb)
    axes[0, 0].set_title(f'1. Original Color Pattern\nSize: {height}x{width} pixels', fontsize=10)
    axes[0, 0].axis('on')
    axes[0, 0].grid(True, alpha=0.3)
    
    # Grayscale from color
    axes[0, 1].imshow(grayscale_display, cmap='gray', vmin=0, vmax=255)
    axes[0, 1].set_title(f'2. Grayscale (from color)\nSize: {height}x{width} pixels\nMin diff: {original_stats["min_diff"]:.4f}', fontsize=10)
    axes[0, 1].axis('on')
    axes[0, 1].grid(True, alpha=0.3)
    
    # Optimized grayscale
    axes[0, 2].imshow(optimized_display, cmap='gray', vmin=0, vmax=255)
    axes[0, 2].set_title(f'3. Optimized Grayscale (Target)\nSize: {height}x{width} pixels\nMin diff: {optimized_stats["min_diff"]:.4f}', fontsize=10, fontweight='bold')
    axes[0, 2].axis('on')
    axes[0, 2].grid(True, alpha=0.3)
    
    # Row 2: CMYK generation from optimized grayscale
    # RGB from CMYK
    rgb_from_cmyk_viz = cv2.cvtColor(rgb_from_cmyk_display, cv2.COLOR_BGR2RGB)
    axes[1, 0].imshow(rgb_from_cmyk_viz)
    axes[1, 0].set_title(f'4. RGB from Generated CMYK\nCMYK Min diff: {cmyk_stats["min_diff"]:.4f}\n(Best for printing)', fontsize=10, fontweight='bold')
    axes[1, 0].axis('on')
    axes[1, 0].grid(True, alpha=0.3)
    
    # Grayscale from CMYK (verification)
    axes[1, 1].imshow(grayscale_from_cmyk_display, cmap='gray', vmin=0, vmax=255)
    axes[1, 1].set_title(f'5. Grayscale from CMYK\nMin diff: {grayscale_from_cmyk_stats["min_diff"]:.4f}\nError: {gray_error:.4f}', fontsize=10)
    axes[1, 1].axis('on')
    axes[1, 1].grid(True, alpha=0.3)
    
    # Show CMYK channels visualization
    cmyk_viz = np.zeros((height, width, 3), dtype=np.uint8)
    cmyk_viz[:, :, 0] = (cmyk_from_grayscale[:, :, 0] * 255).astype(np.uint8)  # Cyan in B channel
    cmyk_viz[:, :, 1] = (cmyk_from_grayscale[:, :, 1] * 255).astype(np.uint8)  # Magenta in G channel
    cmyk_viz[:, :, 2] = (cmyk_from_grayscale[:, :, 2] * 255).astype(np.uint8)  # Yellow in R channel
    axes[1, 2].imshow(cv2.cvtColor(cmyk_viz, cv2.COLOR_BGR2RGB))
    axes[1, 2].set_title(f'6. CMYK Channels (CMY only)\nK channel separate', fontsize=10)
    axes[1, 2].axis('on')
    axes[1, 2].grid(True, alpha=0.3)
    
    plt.tight_layout()
    print(f"\nDisplay info: Figure DPI={fig.dpi}, Size={fig.get_size_inches()} inches")
    plt.show()

    # Save all images
    cv2.imwrite('color_pattern.png', pattern_display)
    cv2.imwrite('grayscale_from_color.png', grayscale_display)
    cv2.imwrite('optimized_grayscale.png', optimized_display)
    cv2.imwrite('cmyk_generated_rgb.png', rgb_from_cmyk_display)
    cv2.imwrite('grayscale_from_cmyk.png', grayscale_from_cmyk_display)
    
    # Save CMYK data
    np.save('cmyk_from_grayscale.npy', cmyk_from_grayscale)
    
    # Save CMYK as TIFF (proper format for printing)
    print("\nSaving CMYK as TIFF...")
    try:
        # Convert CMYK to uint8 (0-255 range)
        cmyk_uint8 = (cmyk_from_grayscale * 255).astype(np.uint8)
        
        # For TIFF CMYK, we need to ensure proper channel order
        # PIL CMYK expects channels in C, M, Y, K order
        # Split channels
        c_channel = cmyk_uint8[:, :, 0]
        m_channel = cmyk_uint8[:, :, 1]
        y_channel = cmyk_uint8[:, :, 2]
        k_channel = cmyk_uint8[:, :, 3]
        
        # Create PIL images for each channel and merge
        c_img = Image.fromarray(c_channel, mode='L')
        m_img = Image.fromarray(m_channel, mode='L')
        y_img = Image.fromarray(y_channel, mode='L')
        k_img = Image.fromarray(k_channel, mode='L')
        
        # Merge into CMYK image
        cmyk_pil = Image.merge('CMYK', (c_img, m_img, y_img, k_img))
        
        # Save as TIFF with CMYK color space at 300 DPI for professional printing
        cmyk_pil.save('cmyk_pattern.tiff', format='TIFF', dpi=(300, 300))
        
        # Verify the TIFF was saved by checking file size
        import os
        if os.path.exists('cmyk_pattern.tiff'):
            file_size = os.path.getsize('cmyk_pattern.tiff')
            print(f"CMYK TIFF saved successfully: {file_size} bytes at 300 DPI")
        
    except Exception as e:
        print(f"Error saving CMYK TIFF: {e}")
    
    print("\nImages saved successfully!")
    print("CMYK-generated RGB: cmyk_generated_rgb.png")
    print("Grayscale verification: grayscale_from_cmyk.png")
    print("CMYK data: cmyk_from_grayscale.npy")
    print("CMYK TIFF (300 DPI): cmyk_pattern.tiff (professional printing format)")

if __name__ == '__main__':
    main()