"""Ben Graham Preprocessing Pipeline for Retinal Fundus Images.

This pipeline performs the canonical preprocessing used in competitive Diabetic
Retinopathy grading:
1. Circular ROI bounding detection and black-border cropping.
2. Aspect-ratio preserving / isotropic resize to target resolution (e.g., 380x380).
3. Contrast Limited Adaptive Histogram Equalization (CLAHE) on L-channel in LAB space.
4. Local mean subtraction via Gaussian blurring and cv2.addWeighted to eliminate
   vignetting and inter-camera illumination gradients.
"""

from typing import Dict, Tuple, Union
import cv2
import numpy as np
from PIL import Image


def crop_circular_roi(image: np.ndarray, tolerance: int = 7) -> np.ndarray:
    """Crop black borders and isolate circular retinal fundus ROI.

    Args:
        image: Input image as RGB numpy array of shape (H, W, 3).
        tolerance: Pixel intensity threshold to distinguish foreground from border.

    Returns:
        Cropped numpy array containing the circular fundus region.
    """
    if image.ndim == 2:
        gray = image
    else:
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    mask = gray > tolerance
    if not np.any(mask):
        return image

    # Find bounding box of foreground mask
    row_sums = np.sum(mask, axis=1)
    col_sums = np.sum(mask, axis=0)

    rows = np.where(row_sums > 0)[0]
    cols = np.where(col_sums > 0)[0]

    if len(rows) == 0 or len(cols) == 0:
        return image

    y_min, y_max = rows[0], rows[-1]
    x_min, x_max = cols[0], cols[-1]

    # Ensure box is square or preserve aspect
    return image[y_min : y_max + 1, x_min : x_max + 1]


def apply_clahe(
    image: np.ndarray,
    clip_limit: float = 2.0,
    tile_grid_size: Tuple[int, int] = (8, 8),
) -> np.ndarray:
    """Apply CLAHE to the L channel of LAB color space.

    Args:
        image: RGB numpy array of shape (H, W, 3).
        clip_limit: Threshold for contrast limiting.
        tile_grid_size: Grid size for histogram equalization.

    Returns:
        Contrast-enhanced RGB numpy array.
    """
    lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)
    l_channel, a_channel, b_channel = cv2.split(lab)

    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
    cl = clahe.apply(l_channel)

    merged = cv2.merge((cl, a_channel, b_channel))
    return cv2.cvtColor(merged, cv2.COLOR_LAB2RGB)


def subtract_local_mean(
    image: np.ndarray,
    sigma: float = 10.0,
    alpha: float = 4.0,
    beta: float = -4.0,
    gamma: int = 128,
) -> np.ndarray:
    """Subtract local mean using Gaussian blur (Ben Graham's method).

    Calculates: output = alpha * image + beta * GaussianBlur(image) + gamma.

    Args:
        image: RGB numpy array of shape (H, W, 3).
        sigma: Standard deviation for Gaussian kernel.
        alpha: Weight of original image.
        beta: Weight of blurred image.
        gamma: Scalar added to shift midtones to 128.

    Returns:
        Local-mean subtracted RGB numpy array.
    """
    blurred = cv2.GaussianBlur(image, (0, 0), sigma)
    # cv2.addWeighted expects uint8 or float32. We use float32 for precision then clip.
    output = cv2.addWeighted(image, alpha, blurred, beta, gamma)
    return np.clip(output, 0, 255).astype(np.uint8)


def mask_circle(image: np.ndarray, radius_ratio: float = 0.95) -> np.ndarray:
    """Apply a smooth circular mask to remove edge artifacts.

    Args:
        image: RGB numpy array of shape (H, W, 3).
        radius_ratio: Ratio of circle radius to half min dimension.

    Returns:
        Circled-masked RGB numpy array.
    """
    h, w = image.shape[:2]
    center = (w // 2, h // 2)
    radius = int(min(center) * radius_ratio)

    mask = np.zeros((h, w), dtype=np.uint8)
    cv2.circle(mask, center, radius, 255, -1)

    result = cv2.bitwise_and(image, image, mask=mask)
    return result


def ben_graham_pipeline(
    image: Union[np.ndarray, Image.Image],
    image_size: int = 380,
    return_intermediates: bool = False,
) -> Union[np.ndarray, Dict[str, np.ndarray]]:
    """Execute the full Ben Graham preprocessing pipeline.

    Args:
        image: Input image as RGB numpy array or PIL Image.
        image_size: Target square resolution (e.g., 380).
        return_intermediates: If True, returns dict with intermediate stages:
            'raw', 'cropped', 'clahe_and_mean' (for Gradio 3-panel UI).

    Returns:
        Processed RGB numpy array, or dictionary of intermediate RGB arrays.
    """
    if isinstance(image, Image.Image):
        raw_rgb = np.array(image.convert("RGB"))
    else:
        raw_rgb = image.copy()

    # Step 1: Circular ROI crop
    cropped = crop_circular_roi(raw_rgb)
    if cropped.size == 0 or cropped.shape[0] < 10 or cropped.shape[1] < 10:
        cropped = raw_rgb

    # Step 2: Resize to target dimension
    resized = cv2.resize(cropped, (image_size, image_size), interpolation=cv2.INTER_AREA)

    # Step 3: CLAHE contrast enhancement
    clahe_img = apply_clahe(resized)

    # Step 4: Subtract local mean
    normalized = subtract_local_mean(clahe_img, sigma=image_size / 30.0)

    # Step 5: Clean circular boundary mask
    final_img = mask_circle(normalized)

    if return_intermediates:
        return {
            "raw": raw_rgb,
            "cropped": resized,
            "clahe_normalized": final_img,
        }
    return final_img