"""
utils/preprocess.py
====================
Image preprocessing pipeline for the FGH-Dermatology inference pipeline.

Responsible for turning an arbitrary uploaded dermoscopic/clinical image
(JPEG, PNG, TIFF, BMP — any colour mode) into the normalised float tensor
the deep learning model expects, while remaining robust across both
light- and dark-skinned patient imagery.
"""

from typing import Tuple

import numpy as np
from PIL import Image, ImageOps


def load_rgb_image(image_path: str) -> Image.Image:
    """
    Load an image from disk, correct EXIF orientation, and force RGB mode
    (handling RGBA, palette, grayscale, and CMYK source images safely).
    """
    img = Image.open(image_path)
    img = ImageOps.exif_transpose(img)  # respect camera/scanner orientation
    if img.mode != "RGB":
        img = img.convert("RGB")
    return img


def preprocess_image(image_path: str, target_size: Tuple[int, int] = (224, 224)) -> np.ndarray:
    """
    Full preprocessing pipeline for model inference.

    Steps
    -----
    1. Load image and normalise to RGB (handles RGBA -> RGB conversion).
    2. Resize to the model's expected square input dimensions.
    3. Scale pixel values to the [0, 1] range.
    4. Apply mild contrast-normalisation so the model attends to lesion
       morphology (shape, border, texture) rather than raw skin colour —
       supporting equitable performance across Fitzpatrick skin types.
    5. Expand dimensions to a (1, H, W, 3) batch tensor.

    Parameters
    ----------
    image_path : str
        Path to the source image on disk.
    target_size : tuple(int, int)
        (height, width) the model expects.

    Returns
    -------
    np.ndarray of shape (1, H, W, 3), dtype float32, values in [0, 1].
    """
    img = load_rgb_image(image_path)
    img = img.resize(target_size, Image.LANCZOS)

    arr = np.asarray(img).astype(np.float32) / 255.0

    arr = _normalise_for_skintone_robustness(arr)

    return np.expand_dims(arr, axis=0)


def _normalise_for_skintone_robustness(arr: np.ndarray) -> np.ndarray:
    """
    Apply a light per-channel standardisation (zero-mean-ish rescale clipped
    back to [0, 1]) so absolute skin colour/luminance contributes less to
    activations than structural lesion features (border, asymmetry,
    texture). This is a lightweight, model-agnostic step; a production model
    trained explicitly for skintone equity may apply its own calibrated
    normalisation instead — in which case this step is a harmless no-op
    pass-through that callers can disable by passing `skip_norm=True` to
    `preprocess_image` in a future revision.
    """
    channel_mean = arr.mean(axis=(0, 1), keepdims=True)
    channel_std = arr.std(axis=(0, 1), keepdims=True) + 1e-6

    # Blend a mild standardisation with the original signal so we improve
    # cross-skintone contrast normalisation without discarding true colour
    # information the model may also rely on.
    standardised = (arr - channel_mean) / channel_std
    standardised = (standardised - standardised.min()) / (
        standardised.max() - standardised.min() + 1e-6
    )

    blended = 0.6 * arr + 0.4 * standardised
    return np.clip(blended, 0.0, 1.0).astype(np.float32)


def get_image_dimensions(image_path: str) -> Tuple[int, int]:
    """Return (width, height) of the source image without full decoding cost."""
    with Image.open(image_path) as img:
        return img.size