# src/preprocess.py
# Image preprocessing utilities.

import cv2
import numpy as np


def apply_clahe(
    img_rgb: np.ndarray,
    clip_limit: float = 4.0,
    grid_size: int = 6,
) -> np.ndarray:
    """
    Apply CLAHE on the L channel of the LAB colour space.

    Flow:
    RGB → LAB → enhance L channel → RGB
    """

    clahe = cv2.createCLAHE(
        clipLimit=clip_limit,
        tileGridSize=(grid_size, grid_size),
    )

    lab = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2LAB)
    lab[:, :, 0] = clahe.apply(lab[:, :, 0])

    return cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)


def preprocess(
    img_rgb: np.ndarray,
    use_clahe: bool = True,
    clip: float = 4.0,
    grid: int = 6,
) -> np.ndarray:
    """
    Run preprocessing on an RGB uint8 array.
    """

    out = img_rgb.copy()

    if use_clahe:
        out = apply_clahe(
            out,
            clip_limit=clip,
            grid_size=grid,
        )

    return out