src/pipeline/visualization.py

"""
Visualization module for EL defect analysis.

Creates overlay images with color-coded defect masks:
- Crack → Blue (visible against bright cell regions)
- Dark → Red (contrast with bright areas)
- Cross → Cyan (distinguishable from regular cracks)
- Busbar → Green (feature, not defect)

All overlays use alpha blending so original image detail remains visible.
Handles resize alignment to prevent mask/image size mismatches.
"""

import cv2
import numpy as np
from typing import Dict, List, Tuple, Optional


# Color scheme (BGR for OpenCV)
DEFECT_COLORS_BGR = {
    "background": (0, 0, 0),       # Black (not drawn)
    "dark": (0, 0, 255),           # Red
    "crack": (255, 0, 0),          # Blue
    "cross": (255, 255, 0),        # Cyan
    "busbar": (0, 255, 0),         # Green
}

# RGB for matplotlib/PIL/Streamlit
DEFECT_COLORS_RGB = {
    "background": (0, 0, 0),
    "dark": (255, 0, 0),           # Red
    "crack": (0, 0, 255),          # Blue
    "cross": (0, 255, 255),        # Cyan
    "busbar": (0, 255, 0),         # Green
}

CLASS_NAMES = ["background", "dark", "crack", "cross", "busbar"]


def create_overlay(
    image: np.ndarray,
    mask: np.ndarray,
    alpha: float = 0.4,
    show_background: bool = False,
) -> np.ndarray:
    """
    Create colored overlay of segmentation mask on image.
    
    Args:
        image: Grayscale or BGR image (any size)
        mask: Class index mask (any size, will be resized to match image)
        alpha: Overlay transparency (0 = fully transparent, 1 = fully opaque)
        show_background: If True, also color background class
        
    Returns:
        BGR image with colored overlay
    """
    # Ensure image is BGR
    if image.ndim == 2:
        vis = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    else:
        vis = image.copy()
    
    # Ensure uint8
    if vis.dtype != np.uint8:
        if vis.max() <= 1.0:
            vis = (vis * 255).astype(np.uint8)
        else:
            vis = vis.astype(np.uint8)
    
    h, w = vis.shape[:2]
    
    # Resize mask to match image (CRITICAL: use NEAREST to preserve labels)
    if mask.shape[:2] != (h, w):
        mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
    
    # Create color overlay
    overlay = vis.copy()
    
    for class_idx, class_name in enumerate(CLASS_NAMES):
        if class_idx == 0 and not show_background:
            continue
        
        color = DEFECT_COLORS_BGR[class_name]
        class_mask = mask == class_idx
        
        if class_mask.any():
            overlay[class_mask] = color
    
    # Alpha blend
    result = cv2.addWeighted(vis, 1 - alpha, overlay, alpha, 0)
    
    return result


def create_class_overlay(
    image: np.ndarray,
    mask: np.ndarray,
    class_name: str,
    alpha: float = 0.5,
    color: Optional[Tuple[int, int, int]] = None,
) -> np.ndarray:
    """
    Create overlay for a single class.
    
    Args:
        image: Grayscale or BGR image
        mask: Binary mask for one class
        class_name: For color lookup
        alpha: Overlay transparency
        color: Override color (BGR)
    """
    if image.ndim == 2:
        vis = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    else:
        vis = image.copy()
    
    if vis.dtype != np.uint8:
        vis = (vis * 255).astype(np.uint8) if vis.max() <= 1 else vis.astype(np.uint8)
    
    h, w = vis.shape[:2]
    if mask.shape[:2] != (h, w):
        mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
    
    if color is None:
        color = DEFECT_COLORS_BGR.get(class_name, (255, 255, 255))
    
    overlay = vis.copy()
    overlay[mask > 0] = color
    
    return cv2.addWeighted(vis, 1 - alpha, overlay, alpha, 0)


def create_color_mask(
    mask: np.ndarray,
    include_background: bool = False,
) -> np.ndarray:
    """
    Convert class index mask to RGB color visualization.
    
    Returns:
        (H, W, 3) uint8 RGB image
    """
    h, w = mask.shape[:2]
    color_img = np.zeros((h, w, 3), dtype=np.uint8)
    
    for class_idx, class_name in enumerate(CLASS_NAMES):
        if class_idx == 0 and not include_background:
            continue
        
        color = DEFECT_COLORS_RGB[class_name]
        color_img[mask == class_idx] = color
    
    return color_img


def draw_cell_results(
    image: np.ndarray,
    cell_results: List[dict],
    cells: list,
) -> np.ndarray:
    """
    Draw cell analysis results on module image.
    
    Shows per-cell:
    - Bounding box (green = PASS, red = FAIL)
    - Cell ID
    - Defect score
    """
    if image.ndim == 2:
        vis = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    else:
        vis = image.copy()
    
    if vis.dtype != np.uint8:
        vis = (vis * 255).astype(np.uint8) if vis.max() <= 1 else vis.astype(np.uint8)
    
    for cell_info, result in zip(cells, cell_results):
        y1, x1, y2, x2 = cell_info.bbox
        score = result.get("defect_score", 0)
        
        # Color: green for good, yellow for moderate, red for bad
        if score < 25:
            color = (0, 255, 0)  # Green
        elif score < 50:
            color = (0, 255, 255)  # Yellow
        else:
            color = (0, 0, 255)  # Red
        
        cv2.rectangle(vis, (x1, y1), (x2, y2), color, 2)
        
        # Label
        label = f"C{cell_info.cell_id}: {score:.0f}"
        cv2.putText(
            vis, label, (x1 + 2, y1 + 15),
            cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1
        )
    
    return vis


def create_summary_image(
    original: np.ndarray,
    overlay: np.ndarray,
    mask_color: np.ndarray,
    decision: str,
    score: float,
) -> np.ndarray:
    """
    Create a summary image with original, overlay, and color mask side by side.
    
    Returns:
        (H, W*3, 3) BGR image with all three panels
    """
    # Ensure all are BGR
    panels = []
    for img in [original, overlay, mask_color]:
        if img.ndim == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        if img.dtype != np.uint8:
            img = (img * 255).astype(np.uint8) if img.max() <= 1 else img.astype(np.uint8)
        panels.append(img)
    
    # Resize to same height
    target_h = 400
    resized = []
    for p in panels:
        scale = target_h / p.shape[0]
        new_w = int(p.shape[1] * scale)
        resized.append(cv2.resize(p, (new_w, target_h)))
    
    # Concatenate horizontally
    summary = np.hstack(resized)
    
    # Add decision text
    color = (0, 255, 0) if decision == "PASS" else (0, 0, 255)
    text = f"{decision} (Score: {score:.1f})"
    cv2.putText(
        summary, text, (10, 30),
        cv2.FONT_HERSHEY_SIMPLEX, 1.0, color, 2
    )
    
    return summary


def create_legend(height: int = 400, width: int = 200) -> np.ndarray:
    """Create a color legend for defect classes."""
    legend = np.ones((height, width, 3), dtype=np.uint8) * 255
    
    y_offset = 30
    for class_name in CLASS_NAMES[1:]:  # Skip background
        color = DEFECT_COLORS_BGR[class_name]
        
        # Color swatch
        cv2.rectangle(
            legend, (10, y_offset), (40, y_offset + 20), color, -1
        )
        cv2.rectangle(
            legend, (10, y_offset), (40, y_offset + 20), (0, 0, 0), 1
        )
        
        # Label
        cv2.putText(
            legend, class_name.capitalize(), (50, y_offset + 15),
            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1
        )
        
        y_offset += 35
    
    return legend