4_Matrix_Figure_Generator.py

#!/usr/bin/env python3
"""
HPA ROI Matrix Figure Generator
Creates a matrix figure showing cropped ROIs with various processing methods.
"""

import csv
from pathlib import Path
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.backends.backend_pdf import PdfPages
import matplotlib.gridspec as gridspec
import tifffile


def load_roi_examples(roi_file_path):
    """
    Load ROI examples from CSV file.
    
    Args:
        roi_file_path: Path to the ROI_examples.txt file
    
    Returns:
        list: List of dictionaries containing ROI data for examples only
    """
    examples = []
    
    if not roi_file_path.exists():
        print(f"Error: {roi_file_path} not found!")
        return examples
    
    with open(roi_file_path, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        for row in reader:
            if row['Example'] == 'T':
                examples.append(row)
    
    return examples


def create_crops_folder(script_dir):
    """
    Create Crops folder if it doesn't exist.
    
    Args:
        script_dir: Script directory
    
    Returns:
        Path: Path to Crops folder
    """
    crops_dir = script_dir / "Crops"
    crops_dir.mkdir(exist_ok=True)
    return crops_dir


def get_file_paths(entry, script_dir):
    """
    Get all file paths for an entry.
    
    Args:
        entry: ROI entry dictionary
        script_dir: Script directory
    
    Returns:
        dict: Dictionary of file paths
    """
    folder_path = script_dir / entry['FolderPath']
    prefix = entry['ImagePrefix']
    
    # For TIF files, try to find them with or without hash suffix
    paths = {
        'blue_tif': find_tif_file(folder_path, prefix, 'blue'),
        'red_tif': find_tif_file(folder_path, prefix, 'red'),
        'green_tif': find_tif_file(folder_path, prefix, 'green'),
        'yellow_tif': find_tif_file(folder_path, prefix, 'yellow'),
        'blue_red_green_yellow_jpg': folder_path / f"{prefix}_blue_red_green_yellow.jpg",
        'blue_red_green_jpg': folder_path / f"{prefix}_blue_red_green.jpg",
        'blue_jpg': folder_path / f"{prefix}_blue.jpg",
        'red_jpg': folder_path / f"{prefix}_red.jpg",
        'green_jpg': folder_path / f"{prefix}_green.jpg",
        'yellow_jpg': folder_path / f"{prefix}_yellow.jpg",
    }
    
    return paths

def find_tif_file(folder_path, prefix, channel):
    """
    Find TIF file with or without hash suffix.
    
    Args:
        folder_path: Path to the folder containing images
        prefix: Image prefix (may contain hash)
        channel: Channel name (blue, red, green, yellow)
    
    Returns:
        Path or None: Path to the TIF file if found, None otherwise
    """
    # First try with simplified prefix (first 3 parts)
    parts = prefix.split('_')
    if len(parts) >= 3:
        simplified_prefix = '_'.join(parts[:3])
        simple_path = folder_path / f"{simplified_prefix}_{channel}.tif"
        if simple_path.exists():
            return simple_path
    
    # Then try with full prefix (including hash)
    full_path = folder_path / f"{prefix}_{channel}.tif"
    if full_path.exists():
        return full_path
    
    # Not found
    return None

def crop_image(image_path, roi_x, roi_y, roi_size=640):
    """
    Crop an image to the ROI.
    
    Args:
        image_path: Path to the image file
        roi_x: X coordinate of ROI top-left
        roi_y: Y coordinate of ROI top-left
        roi_size: Size of ROI square
    
    Returns:
        numpy array: Cropped image as numpy array (preserving original dtype and range)
    """
    if image_path is None or not image_path.exists():
        if image_path:
            print(f"Warning: {image_path} not found")
        return None
    
    try:
        # Check if it's a TIF file - use tifffile library
        if str(image_path).lower().endswith('.tif') or str(image_path).lower().endswith('.tiff'):
            import tifffile
            # Read TIF ignoring OME metadata - just read the raw image data
            img_array = tifffile.imread(image_path, is_ome=False)
            # Crop directly on numpy array - PRESERVE ORIGINAL DTYPE
            cropped = img_array[roi_y:roi_y + roi_size, roi_x:roi_x + roi_size]
            return cropped
        else:
            # For JPG and other formats, use PIL
            img = Image.open(image_path)
            cropped = img.crop((roi_x, roi_y, roi_x + roi_size, roi_y + roi_size))
            return np.array(cropped)
            
    except Exception as e:
        print(f"Warning: Could not open {image_path}: {e}")
        print(f"  File will NOT be deleted - may need special handling")
        return None


def normalize_to_full_range(img_array):
    """
    Normalize image to full range based on bit depth, then convert to 8-bit.
    
    Args:
        img_array: Input image as numpy array
    
    Returns:
        numpy array: 8-bit image using full range of input data
    """
    if img_array is None:
        return None
    
    # Get actual min and max values in the image
    img_min = img_array.min()
    img_max = img_array.max()
    
    # Avoid division by zero
    if img_max == img_min:
        return np.zeros_like(img_array, dtype=np.uint8)
    
    # Normalize to 0-255 based on actual data range
    normalized = ((img_array.astype(np.float32) - img_min) / (img_max - img_min) * 255).astype(np.uint8)
    
    return normalized

def adjust_intensity(img_array, vmin, vmax, is_tif=False):
    """
    Adjust image intensity to specified range and convert to grayscale.
    Automatically scales vmin/vmax for 8-bit TIFs.
    
    Args:
        img_array: Input image as numpy array
        vmin: Minimum intensity value
        vmax: Maximum intensity value
        is_tif: If True, applies bit-depth scaling for TIF files
    
    Returns:
        numpy array: Adjusted grayscale image (0-255)
    """
    if img_array is None:
        return None
    
    # Only scale for TIF files with 8-bit depth
    if is_tif and img_array.dtype == np.uint8:
        # 8-bit TIF: scale down the range from 16-bit to 8-bit
        # 16-bit range is 0-65535, 8-bit is 0-255
        # So divide by 257 (65535/255)
        scale_factor = 255 / 65535
        vmin_scaled = int(vmin * scale_factor)
        vmax_scaled = int(vmax * scale_factor)
    else:
        # JPGs or 16-bit TIFs: use the range as-is
        vmin_scaled = vmin
        vmax_scaled = vmax
    
    # Clip values to the specified range
    img_clipped = np.clip(img_array, vmin_scaled, vmax_scaled)
    
    # Normalize to 0-255
    if vmax_scaled > vmin_scaled:
        img_normalized = ((img_clipped - vmin_scaled) / (vmax_scaled - vmin_scaled) * 255).astype(np.uint8)
    else:
        img_normalized = np.zeros_like(img_clipped, dtype=np.uint8)
    
    return img_normalized

def extract_channel_from_rgb(rgb_image, channel_idx):
    """
    Extract a single channel from RGB image.
    
    Args:
        rgb_image: RGB image as numpy array
        channel_idx: Channel index (0=R, 1=G, 2=B)
    
    Returns:
        numpy array: Single channel as 2D array
    """
    if len(rgb_image.shape) == 3:
        return rgb_image[:, :, channel_idx]
    else:
        return rgb_image


def compute_intensity_profile_correlation(ch1, ch2):
    """
    Compute correlation using intensity profiles along image center.
    This shows channel crosstalk much more clearly than scatter plots.
    
    Args:
        ch1, ch2: Two different channel arrays (e.g., blue vs red)
    
    Returns:
        dict: Contains profiles and correlation coefficient
    """
    if ch1 is None or ch2 is None:
        return None
    
    # Get image dimensions
    h, w = ch1.shape
    
    # Extract horizontal profile through center
    center_row = h // 2
    profile1_h = ch1[center_row, :]
    profile2_h = ch2[center_row, :]
    
    # Extract vertical profile through center
    center_col = w // 2
    profile1_v = ch1[:, center_col]
    profile2_v = ch2[:, center_col]
    
    # Combine both profiles for overall correlation
    combined1 = np.concatenate([profile1_h, profile1_v])
    combined2 = np.concatenate([profile2_h, profile2_v])
    
    # Compute correlation
    if len(combined1) > 10 and np.std(combined1) > 0 and np.std(combined2) > 0:
        pearson_r = np.corrcoef(combined1, combined2)[0, 1]
    else:
        pearson_r = 0.0
    
    return {
        'profile1_h': profile1_h,
        'profile2_h': profile2_h,
        'profile1_v': profile1_v,
        'profile2_v': profile2_v,
        'pearson_r': pearson_r,
        'x_coords': np.arange(len(profile1_h)),
        'y_coords': np.arange(len(profile1_v))
    }

def add_profile_lines_to_image(img_array):
    """
    Add dashed white horizontal and vertical lines to show where profiles are extracted.
    
    Args:
        img_array: Image as numpy array (grayscale or RGB)
    
    Returns:
        numpy array: Image with profile lines drawn
    """
    if img_array is None:
        return None
    
    # Create a copy to avoid modifying original
    img_with_lines = img_array.copy()
    
    # Get dimensions
    if len(img_array.shape) == 3:
        h, w, _ = img_array.shape
    else:
        h, w = img_array.shape
    
    # Calculate center positions
    center_row = h // 2
    center_col = w // 2
    
    # Line parameters - FINE DASHED WHITE LINE
    line_thickness = 1  # Fine line
    dash_length = 8
    gap_length = 4
    
    if len(img_array.shape) == 3:
        # RGB image - white lines
        line_color = [255, 255, 255]
        
        # Draw dashed horizontal line
        for x in range(0, w, dash_length + gap_length):
            for offset in range(-line_thickness//2, line_thickness//2 + 1):
                if 0 <= center_row + offset < h:
                    x_end = min(x + dash_length, w)
                    img_with_lines[center_row + offset, x:x_end, :] = line_color
        
        # Draw dashed vertical line
        for y in range(0, h, dash_length + gap_length):
            for offset in range(-line_thickness//2, line_thickness//2 + 1):
                if 0 <= center_col + offset < w:
                    y_end = min(y + dash_length, h)
                    img_with_lines[y:y_end, center_col + offset, :] = line_color
    else:
        # Grayscale image - white lines
        line_color = 255
        
        # Draw dashed horizontal line
        for x in range(0, w, dash_length + gap_length):
            for offset in range(-line_thickness//2, line_thickness//2 + 1):
                if 0 <= center_row + offset < h:
                    x_end = min(x + dash_length, w)
                    img_with_lines[center_row + offset, x:x_end] = line_color
        
        # Draw dashed vertical line
        for y in range(0, h, dash_length + gap_length):
            for offset in range(-line_thickness//2, line_thickness//2 + 1):
                if 0 <= center_col + offset < w:
                    y_end = min(y + dash_length, h)
                    img_with_lines[y:y_end, center_col + offset] = line_color
    
    return img_with_lines

def save_crop(img_array, output_path):
    """
    Save cropped image as PNG.
    
    Args:
        img_array: Image as numpy array
        output_path: Path to save the image
    """
    if img_array is None:
        return
    
    # Convert to 8-bit if needed
    if img_array.dtype != np.uint8:
        # This is a higher bit-depth image (16-bit, 32-bit, etc.)
        # Normalize to full range
        img_array = normalize_to_full_range(img_array)
    
    # Convert to PIL Image and save
    if len(img_array.shape) == 2:
        # Grayscale
        img = Image.fromarray(img_array, mode='L')
    else:
        # RGB
        img = Image.fromarray(img_array, mode='RGB')
    
    img.save(output_path)


def process_entry(entry, script_dir, crops_dir):
    """
    Process a single ROI entry and create all crops.
    
    Args:
        entry: ROI entry dictionary
        script_dir: Script directory
        crops_dir: Crops directory
    
    Returns:
        dict: Dictionary containing all processed images
    """
    gene = entry['Gene']
    antibody = entry['Antibody']
    cell_line = entry['CellLine']
    roi_x = int(entry['Roi_X'])
    roi_y = int(entry['Roi_Y'])
    
    prefix = entry['ImagePrefix']
    base_name = f"{gene}_{antibody}_{cell_line}_{prefix}"
    
    print(f"Processing: {base_name}")
    
    # Get file paths
    paths = get_file_paths(entry, script_dir)
    
    results = {}
    
    
    # 1. Blue-Red-Green-Yellow JPG (4-channel composite) - FIRST
    crop = crop_image(paths['blue_red_green_yellow_jpg'], roi_x, roi_y)
    if crop is not None:
        save_crop(crop, crops_dir / f"{base_name}_blue_red_green_yellow_jpg.png")
        results['blue_red_green_yellow_jpg'] = crop

    # 2-5. Single TIFs unadjusted - normalize to full range
    crop = crop_image(paths['blue_tif'], roi_x, roi_y)
    if crop is not None:
        crop_normalized = normalize_to_full_range(crop)
        save_crop(crop_normalized, crops_dir / f"{base_name}_blue_tif_unadj.png")
        results['blue_tif_unadj'] = crop_normalized

    crop = crop_image(paths['red_tif'], roi_x, roi_y)
    if crop is not None:
        crop_normalized = normalize_to_full_range(crop)
        save_crop(crop_normalized, crops_dir / f"{base_name}_red_tif_unadj.png")
        results['red_tif_unadj'] = crop_normalized

    crop = crop_image(paths['green_tif'], roi_x, roi_y)
    if crop is not None:
        crop_normalized = normalize_to_full_range(crop)
        save_crop(crop_normalized, crops_dir / f"{base_name}_green_tif_unadj.png")
        results['green_tif_unadj'] = crop_normalized

    crop = crop_image(paths['yellow_tif'], roi_x, roi_y)
    if crop is not None:
        crop_normalized = normalize_to_full_range(crop)
        save_crop(crop_normalized, crops_dir / f"{base_name}_yellow_tif_unadj.png")
        results['yellow_tif_unadj'] = crop_normalized

    # 6-9. Single TIFs adjusted 0-25700
    crop = crop_image(paths['blue_tif'], roi_x, roi_y)
    if crop is not None:
        blue_adj = adjust_intensity(crop, 0, 25700, is_tif=True)
        save_crop(blue_adj, crops_dir / f"{base_name}_blue_tif_adj.png")
        results['blue_tif_adj'] = blue_adj

    crop = crop_image(paths['red_tif'], roi_x, roi_y)
    if crop is not None:
        red_adj = adjust_intensity(crop, 0, 25700, is_tif=True)
        save_crop(red_adj, crops_dir / f"{base_name}_red_tif_adj.png")
        results['red_tif_adj'] = red_adj

    crop = crop_image(paths['green_tif'], roi_x, roi_y)
    if crop is not None:
        green_adj = adjust_intensity(crop, 0, 25700, is_tif=True)
        save_crop(green_adj, crops_dir / f"{base_name}_green_tif_adj.png")
        results['green_tif_adj'] = green_adj

    crop = crop_image(paths['yellow_tif'], roi_x, roi_y)
    if crop is not None:
        yellow_adj = adjust_intensity(crop, 0, 25700, is_tif=True)
        save_crop(yellow_adj, crops_dir / f"{base_name}_yellow_tif_adj.png")
        results['yellow_tif_adj'] = yellow_adj

    # 10-13. Single JPGs adjusted 0-100
    crop = crop_image(paths['blue_jpg'], roi_x, roi_y)
    if crop is not None:
        # Extract only blue channel if RGB
        if len(crop.shape) == 3:
            crop = crop[:, :, 2]  # Blue channel
        blue_adj = adjust_intensity(crop, 0, 100)
        save_crop(blue_adj, crops_dir / f"{base_name}_blue_jpg_adj.png")
        results['blue_jpg_adj'] = blue_adj

    crop = crop_image(paths['red_jpg'], roi_x, roi_y)
    if crop is not None:
        # Extract only red channel if RGB
        if len(crop.shape) == 3:
            crop = crop[:, :, 0]  # Red channel
        red_adj = adjust_intensity(crop, 0, 100)
        save_crop(red_adj, crops_dir / f"{base_name}_red_jpg_adj.png")
        results['red_jpg_adj'] = red_adj

    crop = crop_image(paths['green_jpg'], roi_x, roi_y)
    if crop is not None:
        # Extract only green channel if RGB
        if len(crop.shape) == 3:
            crop = crop[:, :, 1]  # Green channel
        green_adj = adjust_intensity(crop, 0, 100)
        save_crop(green_adj, crops_dir / f"{base_name}_green_jpg_adj.png")
        results['green_jpg_adj'] = green_adj

    crop = crop_image(paths['yellow_jpg'], roi_x, roi_y)
    if crop is not None:
        # Extract yellow channel if RGB
        if len(crop.shape) == 3:
            crop = crop[:, :, 1]  # Green channel (yellow is typically stored here)
        yellow_adj = adjust_intensity(crop, 0, 100)
        save_crop(yellow_adj, crops_dir / f"{base_name}_yellow_jpg_adj.png")
        results['yellow_jpg_adj'] = yellow_adj

    # 14-17. RGB composite and extracted channels (SHOWS BLEED-THROUGH) - LAST
    crop = crop_image(paths['blue_red_green_jpg'], roi_x, roi_y)
    if crop is not None:
        save_crop(crop, crops_dir / f"{base_name}_blue_red_green_jpg.png")
        results['blue_red_green_jpg'] = crop
        
        # Extract and adjust channels from blue_red_green.jpg
        # Note: JPG is RGB, so indices are [R=0, G=1, B=2]
        # HPA composite is typically [Blue=2, Red=0, Green=1]
        
        # Blue channel (B) adjusted 0-100
        blue_ch = extract_channel_from_rgb(crop, 2)
        blue_adj = adjust_intensity(blue_ch, 0, 100)
        save_crop(blue_adj, crops_dir / f"{base_name}_blue_ch_from_composite_adj.png")
        results['blue_ch_composite'] = blue_adj
        
        # Red channel (R) adjusted 0-100
        red_ch = extract_channel_from_rgb(crop, 0)
        red_adj = adjust_intensity(red_ch, 0, 100)
        save_crop(red_adj, crops_dir / f"{base_name}_red_ch_from_composite_adj.png")
        results['red_ch_composite'] = red_adj
        
        # Green channel (G) adjusted 0-100
        green_ch = extract_channel_from_rgb(crop, 1)
        green_adj = adjust_intensity(green_ch, 0, 100)
        save_crop(green_adj, crops_dir / f"{base_name}_green_ch_from_composite_adj.png")
        results['green_ch_composite'] = green_adj
    
    # === RGB COMPOSITE CHANNELS (expect HIGH correlation = bleed-through) ===
    if 'blue_ch_composite' in results and 'red_ch_composite' in results:
        results['corr_rgb_blue_vs_red'] = compute_intensity_profile_correlation(
            results['blue_ch_composite'], results['red_ch_composite']
        )

    if 'blue_ch_composite' in results and 'green_ch_composite' in results:
        results['corr_rgb_blue_vs_green'] = compute_intensity_profile_correlation(
            results['blue_ch_composite'], results['green_ch_composite']
        )

    if 'red_ch_composite' in results and 'green_ch_composite' in results:
        results['corr_rgb_red_vs_green'] = compute_intensity_profile_correlation(
            results['red_ch_composite'], results['green_ch_composite']
        )

    # === SINGLE-CHANNEL JPGs (expect LOW correlation = clean separation) ===
    if 'blue_jpg_adj' in results and 'red_jpg_adj' in results:
        results['corr_jpg_blue_vs_red'] = compute_intensity_profile_correlation(
            results['blue_jpg_adj'], results['red_jpg_adj']
        )

    if 'blue_jpg_adj' in results and 'green_jpg_adj' in results:
        results['corr_jpg_blue_vs_green'] = compute_intensity_profile_correlation(
            results['blue_jpg_adj'], results['green_jpg_adj']
        )

    if 'red_jpg_adj' in results and 'green_jpg_adj' in results:
        results['corr_jpg_red_vs_green'] = compute_intensity_profile_correlation(
            results['red_jpg_adj'], results['green_jpg_adj']
        )

    # Store metadata
    results['gene'] = gene
    results['antibody'] = antibody
    results['cell_line'] = cell_line
    results['annotation'] = entry.get('Annotation', 'N/A')
    
    return results


def create_matrix_figure(processed_entries, output_pdf, output_svg, output_png):
    """
    Create matrix figure with all examples.
    """
    if not processed_entries:
        print("No entries to process!")
        return
    
    num_rows = len(processed_entries)
    
    # Column definitions with headers - MODIFIED
    # Column definitions with headers - MODIFIED
    columns = [
        ('info', 'Information', 2),
        ('sep0', '', 0.2),
        ('blue_red_green_yellow_jpg', 'All Channels\nMerged JPG', 1),
        ('sep1', '', 0.2),
        ('blue_tif_unadj', 'Blue TIF\nFull Range', 1),
        ('red_tif_unadj', 'Red TIF\nFull Range', 1),
        ('green_tif_unadj', 'Green TIF\nFull Range', 1),
        ('yellow_tif_unadj', 'Yellow TIF\nFull Range', 1),
        ('sep2', '', 0.2),
        ('blue_tif_adj', 'Blue TIF\nAdj: 0-25700', 1),
        ('red_tif_adj', 'Red TIF\nAdj: 0-25700', 1),
        ('green_tif_adj', 'Green TIF\nAdj: 0-25700', 1),
        ('yellow_tif_adj', 'Yellow TIF\nAdj: 0-25700', 1),
        ('sep3', '', 0.2),
        ('blue_jpg_adj', 'Blue JPG\nSingle\nAdj: 0-100', 1),
        ('red_jpg_adj', 'Red JPG\nSingle\nAdj: 0-100', 1),
        ('green_jpg_adj', 'Green JPG\nSingle\nAdj: 0-100', 1),
        ('yellow_jpg_adj', 'Yellow JPG\nSingle\nAdj: 0-100', 1),
        ('sep4', '', 0.2),
        ('blue_red_green_jpg', 'RGB JPG\nComposite', 1),
        ('blue_ch_composite', 'Blue from\nRGB JPG\nAdj: 0-100', 1),
        ('red_ch_composite', 'Red from\nRGB JPG\nAdj: 0-100', 1),
        ('green_ch_composite', 'Green from\nRGB JPG\nAdj: 0-100', 1),
        ('sep5', '', 0.3),
        # RGB Composite - Horizontal and Vertical profiles (NO correlation values column)
        ('corr_rgb_plot_h', 'RGB Composite\nHorizontal Profile', 1.8),
        ('corr_rgb_plot_v', 'RGB Composite\nVertical Profile', 1.8),
        ('sep6', '', 0.3),
        # Single JPG - Horizontal and Vertical profiles (NO correlation values column)
        ('corr_jpg_plot_h', 'Single JPG\nHorizontal Profile', 1.8),
        ('corr_jpg_plot_v', 'Single JPG\nVertical Profile', 1.8),
    ]
    
    num_cols = len(columns)
    col_widths = [c[2] for c in columns]
    
    # Create figure
    fig_width = sum(col_widths) * 1.5
    fig_height = (num_rows + 1) * 1.5
    
    fig = plt.figure(figsize=(fig_width, fig_height))
    
    gs = gridspec.GridSpec(num_rows + 2, num_cols, 
                            figure=fig,
                            width_ratios=col_widths,
                            height_ratios=[0.8] + [1] * num_rows + [0.3],
                            hspace=0.08, wspace=0.08)  # Increased spacing to prevent overlap

    # Add header row
    for col_idx, (col_key, col_header, _) in enumerate(columns):
        ax = fig.add_subplot(gs[0, col_idx])
        ax.axis('off')
        
        if col_key.startswith('sep'):
            ax.axvline(x=0.5, color='black', linewidth=3)
        else:
            ax.text(0.5, 0.5, col_header, 
                   ha='center', va='center',
                   fontsize=8, fontweight='bold',
                   wrap=True)
            ax.set_xlim(0, 1)
            ax.set_ylim(0, 1)
    
    # Add data rows
    for row_idx, entry in enumerate(processed_entries):
        actual_row = row_idx + 1
        
        for col_idx, (col_key, col_header, _) in enumerate(columns):
            ax = fig.add_subplot(gs[actual_row, col_idx])
            ax.axis('off')
            
            if col_key.startswith('sep'):
                ax.axvline(x=0.5, color='black', linewidth=3)
            
            elif col_key == 'info':
                info_text = f"Gene: {entry['gene']}\n"
                info_text += f"HPA: {entry['antibody']}\n"
                info_text += f"Cell: {entry['cell_line']}\n"
                info_text += f"Anno: {entry['annotation']}"
                
                ax.text(0.1, 0.5, info_text,
                       ha='left', va='center',
                       fontsize=7,
                       family='monospace')
                ax.set_xlim(0, 1)
                ax.set_ylim(0, 1)
            
            # === RGB HORIZONTAL PROFILE ===
            elif col_key == 'corr_rgb_plot_h':
                ax.axis('on')
                
                corr_br = entry.get('corr_rgb_blue_vs_red')
                corr_bg = entry.get('corr_rgb_blue_vs_green')
                corr_rg = entry.get('corr_rgb_red_vs_green')
                
                if corr_br and corr_bg and corr_rg:
                    x_coords = corr_br['x_coords']
                    
                    # Scale from 0-255 to 0-100 (RGB composite was adjusted 0-100)
                    scale_factor = 100 / 255
                    
                    ax.plot(x_coords, corr_br['profile1_h'] * scale_factor, 'b-', linewidth=0.4, alpha=0.8, label='B')
                    ax.plot(x_coords, corr_br['profile2_h'] * scale_factor, 'r-', linewidth=0.4, alpha=0.8, label='R')
                    ax.plot(x_coords, corr_bg['profile2_h'] * scale_factor, 'g-', linewidth=0.4, alpha=0.8, label='G')
                    
                    # Get correlations
                    r_br = corr_br['pearson_r']
                    r_bg = corr_bg['pearson_r']
                    r_rg = corr_rg['pearson_r']
                    
                    ax.set_xlabel('X', fontsize=4)
                    ax.set_ylabel('Int', fontsize=4)
                    ax.set_ylim(0, 100)
                    ax.legend(fontsize=3.5, loc='upper right', framealpha=0.6, handlelength=1)
                    ax.grid(True, alpha=0.2, linewidth=0.2)
                    ax.tick_params(labelsize=3.5, pad=0.5, length=2)
                    
                    # Tighter layout
                    ax.spines['top'].set_visible(False)
                    ax.spines['right'].set_visible(False)
                    ax.spines['left'].set_linewidth(0.5)
                    ax.spines['bottom'].set_linewidth(0.5)
                    
                    #title = f'H: BR:{r_br:.2f} BG:{r_bg:.2f} RG:{r_rg:.2f}'
                    title = ""
                    ax.set_title(title, fontsize=4.5, color='red', fontweight='bold', pad=0.5)
                else:
                    ax.text(0.5, 0.5, 'N/A', ha='center', va='center', fontsize=8, color='red')
                    ax.set_xlim(0, 1)
                    ax.set_ylim(0, 1)
            
            # === RGB VERTICAL PROFILE ===
            elif col_key == 'corr_rgb_plot_v':
                ax.axis('on')
                
                corr_br = entry.get('corr_rgb_blue_vs_red')
                corr_bg = entry.get('corr_rgb_blue_vs_green')
                corr_rg = entry.get('corr_rgb_red_vs_green')
                
                if corr_br and corr_bg and corr_rg:
                    y_coords = corr_br['y_coords']
                    
                    # Scale from 0-255 to 0-100 (RGB composite was adjusted 0-100)
                    scale_factor = 100 / 255
                    
                    ax.plot(y_coords, corr_br['profile1_v'] * scale_factor, 'b-', linewidth=0.4, alpha=0.8, label='B')
                    ax.plot(y_coords, corr_br['profile2_v'] * scale_factor, 'r-', linewidth=0.4, alpha=0.8, label='R')
                    ax.plot(y_coords, corr_bg['profile2_v'] * scale_factor, 'g-', linewidth=0.4, alpha=0.8, label='G')
                    
                    # Get correlations
                    r_br = corr_br['pearson_r']
                    r_bg = corr_bg['pearson_r']
                    r_rg = corr_rg['pearson_r']
                    
                    ax.set_xlabel('Y', fontsize=4)
                    ax.set_ylabel('Int', fontsize=4)
                    ax.set_ylim(0, 100)
                    ax.legend(fontsize=3.5, loc='upper right', framealpha=0.6, handlelength=1)
                    ax.grid(True, alpha=0.2, linewidth=0.2)
                    ax.tick_params(labelsize=3.5, pad=0.5, length=2)
                    
                    # Tighter layout
                    ax.spines['top'].set_visible(False)
                    ax.spines['right'].set_visible(False)
                    ax.spines['left'].set_linewidth(0.5)
                    ax.spines['bottom'].set_linewidth(0.5)
                    
                    #title = f'V: BR:{r_br:.2f} BG:{r_bg:.2f} RG:{r_rg:.2f}'
                    title = ""
                    ax.set_title(title, fontsize=4.5, color='red', fontweight='bold', pad=0.5)
                else:
                    ax.text(0.5, 0.5, 'N/A', ha='center', va='center', fontsize=8, color='red')
                    ax.set_xlim(0, 1)
                    ax.set_ylim(0, 1)
            
            # === JPG HORIZONTAL PROFILE ===
            elif col_key == 'corr_jpg_plot_h':
                ax.axis('on')
                
                corr_br = entry.get('corr_jpg_blue_vs_red')
                corr_bg = entry.get('corr_jpg_blue_vs_green')
                corr_rg = entry.get('corr_jpg_red_vs_green')
                
                if corr_br and corr_bg and corr_rg:
                    x_coords = corr_br['x_coords']
                    
                    # Scale from 0-255 to 0-100 (JPG was adjusted 0-100)
                    scale_factor = 100 / 255
                    
                    ax.plot(x_coords, corr_br['profile1_h'] * scale_factor, 'b-', linewidth=0.4, alpha=0.8, label='B')
                    ax.plot(x_coords, corr_br['profile2_h'] * scale_factor, 'r-', linewidth=0.4, alpha=0.8, label='R')
                    ax.plot(x_coords, corr_bg['profile2_h'] * scale_factor, 'g-', linewidth=0.4, alpha=0.8, label='G')
                    
                    # Get correlations
                    r_br = corr_br['pearson_r']
                    r_bg = corr_bg['pearson_r']
                    r_rg = corr_rg['pearson_r']
                    
                    ax.set_xlabel('X', fontsize=4)
                    ax.set_ylabel('Int', fontsize=4)
                    ax.set_ylim(0, 100)
                    ax.legend(fontsize=3.5, loc='upper right', framealpha=0.6, handlelength=1)
                    ax.grid(True, alpha=0.2, linewidth=0.2)
                    ax.tick_params(labelsize=3.5, pad=0.5, length=2)
                    
                    # Tighter layout
                    ax.spines['top'].set_visible(False)
                    ax.spines['right'].set_visible(False)
                    ax.spines['left'].set_linewidth(0.5)
                    ax.spines['bottom'].set_linewidth(0.5)
                    
                    #title = f'H: BR:{r_br:.2f} BG:{r_bg:.2f} RG:{r_rg:.2f}'
                    title = ""
                    ax.set_title(title, fontsize=4.5, color='green', fontweight='bold', pad=0.5)
                else:
                    ax.text(0.5, 0.5, 'N/A', ha='center', va='center', fontsize=8, color='red')
                    ax.set_xlim(0, 1)
                    ax.set_ylim(0, 1)
            
            # === JPG VERTICAL PROFILE ===
            elif col_key == 'corr_jpg_plot_v':
                ax.axis('on')
                
                corr_br = entry.get('corr_jpg_blue_vs_red')
                corr_bg = entry.get('corr_jpg_blue_vs_green')
                corr_rg = entry.get('corr_jpg_red_vs_green')
                
                if corr_br and corr_bg and corr_rg:
                    y_coords = corr_br['y_coords']
                    
                    # Scale from 0-255 to 0-100 (JPG was adjusted 0-100)
                    scale_factor = 100 / 255
                    
                    ax.plot(y_coords, corr_br['profile1_v'] * scale_factor, 'b-', linewidth=0.4, alpha=0.8, label='B')
                    ax.plot(y_coords, corr_br['profile2_v'] * scale_factor, 'r-', linewidth=0.4, alpha=0.8, label='R')
                    ax.plot(y_coords, corr_bg['profile2_v'] * scale_factor, 'g-', linewidth=0.4, alpha=0.8, label='G')
                    
                    # Get correlations
                    r_br = corr_br['pearson_r']
                    r_bg = corr_bg['pearson_r']
                    r_rg = corr_rg['pearson_r']
                    
                    ax.set_xlabel('Y', fontsize=4)
                    ax.set_ylabel('Int', fontsize=4)
                    ax.set_ylim(0, 100)
                    ax.legend(fontsize=3.5, loc='upper right', framealpha=0.6, handlelength=1)
                    ax.grid(True, alpha=0.2, linewidth=0.2)
                    ax.tick_params(labelsize=3.5, pad=0.5, length=2)
                    
                    # Tighter layout
                    ax.spines['top'].set_visible(False)
                    ax.spines['right'].set_visible(False)
            
            # === KEEP EXISTING IMAGE DISPLAY CODE ===
            elif col_key in entry:
                if not col_key.startswith('corr_'):
                    img = entry[col_key]
                    
                    # Show profile lines on composite images used for correlation
                    show_profile_lines = col_key in [
                        'blue_red_green_jpg',      # RGB composite
                        'blue_ch_composite',       # Blue from RGB
                        'red_ch_composite',        # Red from RGB
                        'green_ch_composite',      # Green from RGB
                        'blue_jpg_adj',            # Blue single JPG
                        'red_jpg_adj',             # Red single JPG
                        'green_jpg_adj'            # Green single JPG
                    ]
                    
                    # Add profile lines if applicable
                    if show_profile_lines:
                        img_display = add_profile_lines_to_image(img)
                    else:
                        img_display = img
                    
                    if len(img_display.shape) == 2:
                        ax.imshow(img_display, cmap='turbo', vmin=0, vmax=255)
                    else:
                        ax.imshow(img_display)
                    
                    ax.set_aspect('equal')
            else:
                ax.text(0.5, 0.5, 'N/A', ha='center', va='center', fontsize=8, color='red')
                ax.set_xlim(0, 1)
                ax.set_ylim(0, 1)
    
    # Add title
    fig.suptitle('HPA Image Processing Comparison Matrix', 
                fontsize=14, fontweight='bold', y=0.995)
    
    # Add colorbar
    cbar_ax = fig.add_subplot(gs[num_rows + 1, :])
    cbar_ax.axis('off')
    
    import matplotlib.colors as mcolors
    norm = mcolors.Normalize(vmin=0, vmax=255)
    cmap = plt.cm.turbo
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    
    cbar = plt.colorbar(sm, ax=cbar_ax, orientation='horizontal', 
                       fraction=0.8, pad=0.1, aspect=50)
    cbar.set_label('Intensity 0 - max: turbo colormap', fontsize=10)
    cbar.ax.tick_params(labelsize=8)

    plt.savefig(output_pdf, format='pdf', bbox_inches='tight', dpi=150)
    print(f"Saved PDF: {output_pdf}")
    
    # Save as SVG
    plt.savefig(output_svg, format='svg', bbox_inches='tight')
    print(f"Saved SVG: {output_svg}")
    
    # Save as PNG
    plt.savefig(output_png, format='png', bbox_inches='tight')
    print(f"Saved PNG: {output_png}")
    
    plt.close()



def main():
    """
    Main function to orchestrate the matrix figure generation.
    """
    # Setup paths
    script_dir = Path(__file__).parent.resolve()
    crops_dir = create_crops_folder(script_dir)
    
    print("=" * 60)
    print("HPA ROI Matrix Figure Generator")
    print("=" * 60)
    
    # Ask user which file to process
    print("\nWhich ROI file would you like to process?")
    print("  1. ROI_examples.txt (all examples)")
    print("  2. ROI_examples_Selected.txt (selected subset)")
    
    while True:
        choice = input("\nEnter your choice (1 or 2): ").strip()
        
        if choice == "1":
            roi_file = script_dir / "ROI_examples.txt"
            output_suffix = ""
            break
        elif choice == "2":
            roi_file = script_dir / "ROI_examples_Selected.txt"
            output_suffix = "_Selected"
            break
        else:
            print("Invalid choice. Please enter 1 or 2.")
    
    # Set output file names based on choice
    output_pdf = script_dir / f"HPA_comparison_matrix{output_suffix}.pdf"
    output_svg = script_dir / f"HPA_comparison_matrix{output_suffix}.svg"
    output_png = script_dir / f"HPA_comparison_matrix{output_suffix}.png"
    
    # Check if file exists
    if not roi_file.exists():
        print(f"\nError: {roi_file} not found!")
        print("Please make sure the file exists in the script directory.")
        return
    
    # Load ROI examples
    print(f"\nLoading ROI examples from: {roi_file}")
    examples = load_roi_examples(roi_file)
    
    if not examples:
        print("No examples found with Example='T'")
        return
    
    print(f"Found {len(examples)} examples to process")
    
    # Process each entry
    print("\n" + "=" * 60)
    print("Processing entries and creating crops...")
    print("=" * 60 + "\n")
    
    processed_entries = []
    
    for idx, entry in enumerate(examples, 1):
        print(f"\n[{idx}/{len(examples)}]")
        result = process_entry(entry, script_dir, crops_dir)
        processed_entries.append(result)
    
    # Create matrix figure
    print("\n" + "=" * 60)
    print("Creating matrix figure...")
    print("=" * 60 + "\n")
    
    create_matrix_figure(processed_entries, output_pdf, output_svg, output_png)
    
    print("\n" + "=" * 60)
    print("Processing complete!")
    print("=" * 60)
    print(f"\nOutputs:")
    print(f"  - Crops folder: {crops_dir}")
    print(f"  - PDF figure: {output_pdf}")
    print(f"  - SVG figure: {output_svg}")
    print(f"  - PNG figure: {output_png}")
    print(f"\nProcessed {len(processed_entries)} examples")


if __name__ == "__main__":
    main()