app/preprocessing.py

"""Preprocessing pipeline for scanned film negatives."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional, Tuple

import numpy as np
import torch
import warnings
from PIL import Image, ImageOps

# WP-11 Fix A: oversized-upload cap (50 MP) — enforced locally in
# _guard_intake_size. PIL's own decompression-bomb limit (~178 MP error) is
# left at its default: mutating Image.MAX_IMAGE_PIXELS at import would be
# process-global state, and the removed line actually *raised* PIL's
# threshold. Uploads beyond PIL's limit raise DecompressionBombError, which
# process_negative's try/except turns into a clean error message.
INTAKE_MAX_MEGAPIXELS: float = 50.0


@dataclass
class PreprocessedNegative:
    """Container for a preprocessed negative scan."""

    # --- Existing fields (do not remove or rename) ---
    rgb: np.ndarray          # float32, shape (H, W, 3), range [0, 1]
    luminance: np.ndarray    # float32, shape (H, W), range [0, 1]
    log_exposure: torch.Tensor  # shape (1, 1, H, W)
    was_inverted: bool
    original_size: Tuple[int, int]

    # --- WP-2 densitometry fields (None when densitometry unavailable) ---
    # NOTE: outputs are RELATIVE — see densitometry.py module docstring.
    density: Optional[np.ndarray] = field(default=None)          # (H, W) float32, D_physical
    h_total: Optional[np.ndarray] = field(default=None)          # (H, W) float32, linear exposure
    confidence_mask: Optional[np.ndarray] = field(default=None)  # (H, W) uint8, TOE=0/VALID=1/SHOULDER=2

    # --- WP-8 color fields (additive) ---
    density_rgb: Optional[np.ndarray] = field(default=None)  # (H, W, 3)
    h_total_rgb: Optional[np.ndarray] = field(default=None)  # (H, W, 3)
    confidence_mask_rgb: Optional[np.ndarray] = field(default=None)  # (H, W, 3)
    is_color: bool = field(default=False)

    # --- WP-11 post-review: fraction bbox of the auto_trim crop, relative to
    # the pre-trim working image (top, bottom, left, right in [0,1]). None when
    # auto_trim is off or trimmed nothing. Consumers that re-derive geometry
    # from the ORIGINAL upload (WP-12 full-res export) must apply this crop,
    # or their guide image will include the border the working images lost.
    trim_bbox_frac: Optional[Tuple[float, float, float, float]] = field(default=None)

    # --- WP-11.1: densitometry anchor actually used (None in linear mode).
    # Consumed by full-res export so work-res and full-res share the same anchor.
    d_min_override_used: Optional[float] = field(default=None)

    # --- WP-11.1 post-review: physics polarity of the upload. True ONLY when
    # the user explicitly declared scan_type="positive" — never set by the
    # display heuristic (was_inverted is a separate, display-only concern).
    # Consumed by full-res export so all densitometry entry points agree.
    physics_is_positive: bool = field(default=False)


_HIGH_BIT_MODES = frozenset({"I;16", "I;16B", "I;16L", "I;16N", "I", "F"})


def _to_float_rgb(image: Image.Image) -> np.ndarray:
    """Load PIL image to float32 RGB in [0, 1], preserving 16-bit/float precision.

    8-bit path (RGB/L/etc.) stays byte-identical to convert("RGB")/255.
    16-bit int modes scale by 65535; float mode normalizes by its own max.
    """
    mode = image.mode
    if mode in _HIGH_BIT_MODES:
        arr = np.asarray(image)
        if arr.ndim == 2:
            # Single-channel high-bit → replicate to RGB
            if mode == "F":
                mx = float(np.max(arr)) if arr.size else 1.0
                mx = mx if mx > 1e-12 else 1.0
                ch = np.clip(arr.astype(np.float32) / mx, 0.0, 1.0)
            else:
                # 16-bit / 32-bit int modes
                ch = np.clip(arr.astype(np.float32) / 65535.0, 0.0, 1.0)
            return np.stack([ch, ch, ch], axis=-1).astype(np.float32)
        # Multi-channel rare high-bit path
        if mode == "F":
            mx = float(np.max(arr)) if arr.size else 1.0
            mx = mx if mx > 1e-12 else 1.0
            out = np.clip(arr.astype(np.float32) / mx, 0.0, 1.0)
        else:
            out = np.clip(arr.astype(np.float32) / 65535.0, 0.0, 1.0)
        if out.ndim == 2:
            out = np.stack([out, out, out], axis=-1)
        elif out.shape[-1] == 1:
            out = np.repeat(out, 3, axis=-1)
        elif out.shape[-1] > 3:
            out = out[..., :3]
        return out.astype(np.float32)

    # 8-bit and ordinary modes — byte-identical to previous convert("RGB")/255
    arr = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
    return arr


def _guard_intake_size(image: Image.Image) -> Image.Image:
    """Downscale uploads beyond INTAKE_MAX_MEGAPIXELS (never OOM). Warns once."""
    w, h = image.size
    mp = (w * h) / 1_000_000.0
    if mp <= INTAKE_MAX_MEGAPIXELS:
        return image
    scale = (INTAKE_MAX_MEGAPIXELS * 1_000_000.0 / (w * h)) ** 0.5
    nw, nh = max(1, int(w * scale)), max(1, int(h * scale))
    warnings.warn(
        f"Upload {w}×{h} ({mp:.1f} MP) exceeds intake cap "
        f"{INTAKE_MAX_MEGAPIXELS:.0f} MP; downscaling to {nw}×{nh}."
    )
    return image.resize((nw, nh), Image.Resampling.LANCZOS)


def _rgb_to_luminance(rgb: np.ndarray) -> np.ndarray:
    # Rec. 709 luma
    return (
        0.2126 * rgb[..., 0] + 0.7152 * rgb[..., 1] + 0.0722 * rgb[..., 2]
    ).astype(np.float32)


def trim_uniform_border(
    rgb: np.ndarray,
    tol: float = 1e-3,
    max_frac: float = 0.25,
) -> Tuple[np.ndarray, Tuple[int, int, int, int]]:
    """Drop near-uniform outer rows/cols (rebate / letterbox), capped by max_frac.

    A row/col is "uniform" if its per-channel std is below ``tol``. Never removes
    more than ``max_frac`` of each edge. Returns (cropped_rgb, (top, bottom, left, right)
    where bottom/right are exclusive indices into the original).

    Conservative and opt-in (WP-11 Fix D). Default preprocess path never calls this.
    """
    h, w = rgb.shape[:2]
    max_t = int(h * max_frac)
    max_b = int(h * max_frac)
    max_l = int(w * max_frac)
    max_r = int(w * max_frac)

    def _row_uniform(i: int) -> bool:
        return float(np.std(rgb[i])) < tol

    def _col_uniform(j: int) -> bool:
        return float(np.std(rgb[:, j])) < tol

    top = 0
    while top < max_t and top < h - 1 and _row_uniform(top):
        top += 1
    bottom = h
    while (h - bottom) < max_b and bottom > top + 1 and _row_uniform(bottom - 1):
        bottom -= 1
    left = 0
    while left < max_l and left < w - 1 and _col_uniform(left):
        left += 1
    right = w
    while (w - right) < max_r and right > left + 1 and _col_uniform(right - 1):
        right -= 1

    # Never return empty
    if bottom <= top or right <= left:
        return rgb, (0, h, 0, w)

    return rgb[top:bottom, left:right].copy(), (top, bottom, left, right)


def _detect_negative_inversion(luminance: np.ndarray) -> bool:
    """Two-statistic negative-scan heuristic (WP-11 Fix B).

    Treat as negative if both:
      - mean(lum) > 0.55  (bright overall — orange mask / inverted image)
      - median(lum) > 0.50 (mass in the bright half, not just a bright outlier)

    Pure function of the luminance array. Manual ``scan_type`` override is
    preferred when the user knows; this heuristic only runs for ``"auto"``.
    """
    mean_l = float(luminance.mean())
    med_l = float(np.median(luminance))
    return mean_l > 0.55 and med_l > 0.50


def _normalize01(arr: np.ndarray, percentile: float = 99.5) -> np.ndarray:
    lo = float(np.percentile(arr, 0.5))
    hi = float(np.percentile(arr, percentile))
    if hi - lo < 1e-6:
        return np.clip(arr, 0.0, 1.0)
    out = (arr - lo) / (hi - lo)
    return np.clip(out, 0.0, 1.0).astype(np.float32)


def luminance_to_log_exposure(luminance: np.ndarray) -> torch.Tensor:
    """
    Map normalized luminance (proxy for transmitted light) to log-exposure.

    For a negative, darker areas = more exposure on film. After inversion
    during preprocessing, higher luminance ≈ more original scene exposure.
    """
    exposure = np.clip(luminance, 1e-4, 1.0)
    log_h = np.log10(exposure)
    tensor = torch.from_numpy(log_h).float().unsqueeze(0).unsqueeze(0)
    return tensor


def preprocess_negative(
    image: Image.Image,
    max_side: int = 1536,
    stock: str = "Generic",
    scan_type: str = "auto",
    scan_calibration: str = "linear",
    auto_trim: bool = False,
    mask_mode: str = "density_margin",
) -> PreprocessedNegative:
    """
    Load, optionally resize, detect inversion, normalize, and build exposure map.

    Also runs the WP-2 densitometry pipeline to populate ``density``,
    ``h_total``, and ``confidence_mask`` on the returned object.  These
    fields are set to None on any error so the UI keeps working.

    Args:
        image:    PIL image from user upload.
        max_side: Longest edge after resize (preserves aspect ratio).
        stock:    Film stock preset name (used for D_min fallback).
        scan_type: ``"auto"`` (heuristic), ``"positive"`` (never invert),
                   ``"negative"`` (always invert). Default ``"auto"``.
        scan_calibration: ``"linear"`` (default — synthetic fixtures, white-point
            estimate cancels) or ``"auto_exposed"`` (pass stock D_min as
            ``d_min_override`` for real auto-exposed scanners).
        auto_trim: If True, trim uniform border (default False — never silent crop).
        mask_mode: confidence-mask criterion, forwarded to densitometry —
            "density_margin" (default, legacy WP-2 contract) or "slope" (VALID
            reflects actual H reliability; the app path opts in).

    Returns:
        PreprocessedNegative ready for physics scoring and API calls.
    """
    # WP-11 Fix A: honor EXIF orientation first, then size-guard oversized uploads
    image = ImageOps.exif_transpose(image) or image
    image = _guard_intake_size(image)

    original_size = image.size
    w, h = image.size
    scale = min(1.0, max_side / max(w, h))
    if scale < 1.0:
        image = image.resize(
            (int(w * scale), int(h * scale)), Image.Resampling.LANCZOS
        )

    rgb = _to_float_rgb(image)

    # WP-11 Fix D: opt-in uniform border trim (default OFF)
    trim_bbox_frac: Optional[Tuple[float, float, float, float]] = None
    if auto_trim:
        pre_h, pre_w = rgb.shape[:2]
        rgb, bbox = trim_uniform_border(rgb)
        if bbox != (0, pre_h, 0, pre_w):
            top, bottom, left, right = bbox
            trim_bbox_frac = (
                top / pre_h, bottom / pre_h, left / pre_w, right / pre_w
            )

    # Pristine scan for densitometry: density must be computed from the raw
    # sRGB scan, BEFORE inversion/normalization (MASTERPLAN Part I.3 / red
    # flag (d) — percentile stretching destroys density information).
    raw_scan_rgb = rgb
    luminance = _rgb_to_luminance(rgb)

    # WP-11 Fix B: scan_type override wins over heuristic
    st = (scan_type or "auto").lower()
    # WP-11.1: physics polarity comes ONLY from the explicit declaration
    physics_is_positive = st == "positive"
    if st == "positive":
        was_inverted = False
    elif st == "negative":
        was_inverted = True
    else:
        was_inverted = _detect_negative_inversion(luminance)

    if was_inverted:
        rgb = 1.0 - rgb
        luminance = 1.0 - luminance

    rgb = _normalize01(rgb)
    luminance = _normalize01(luminance)
    log_exposure = luminance_to_log_exposure(luminance)

    # --- WP-2: densitometry (soft — falls back gracefully) ---
    density_map: Optional[np.ndarray] = None
    h_total_map: Optional[np.ndarray] = None
    conf_mask: Optional[np.ndarray] = None
    density_rgb: Optional[np.ndarray] = None
    h_total_rgb: Optional[np.ndarray] = None
    conf_mask_rgb: Optional[np.ndarray] = None
    is_c: bool = False
    d_min_override_used: Optional[float] = None
    try:
        from densitometry import (
            scan_to_density,
            density_to_h_total,
            scan_to_density_rgb,
            density_to_h_total_rgb,
            combine_confidence_rgb,
            prepare_densitometry_input,
        )
        from film_physics import get_film_curve, get_color_curves, COLOR_STOCK_PRESETS

        is_c = stock in COLOR_STOCK_PRESETS
        # WP-11 Fix C: auto_exposed pins D_min to stock preset (real scanners)
        d_min_override: Optional[float] = None
        cal = (scan_calibration or "linear").lower()
        if cal == "auto_exposed":
            if is_c:
                d_min_override = float(get_color_curves(stock).g.d_min.item())
            else:
                d_min_override = float(get_film_curve(stock).d_min.item())
        d_min_override_used = d_min_override

        # WP-11.1: canonical physics-polarity policy (one place: densitometry).
        # auto/negative return raw_scan_rgb unchanged — byte-identical to today.
        dens_scan_rgb, dens_curves = prepare_densitometry_input(
            raw_scan_rgb, stock, positive_source=physics_is_positive
        )

        if is_c:
            density_map_rgb = scan_to_density_rgb(
                dens_scan_rgb, stock, d_min_override=d_min_override, curves=dens_curves
            )
            h_total_map_rgb, conf_mask_rgb = density_to_h_total_rgb(
                density_map_rgb, dens_curves, mask_mode=mask_mode
            )
            conf_mask = combine_confidence_rgb(conf_mask_rgb)
            density_map = density_map_rgb[..., 1]
            h_total_map = h_total_map_rgb[..., 1]
            density_rgb = density_map_rgb
            h_total_rgb = h_total_map_rgb
            conf_mask_rgb = conf_mask_rgb
        else:
            curve = get_film_curve(stock)
            density_map, _ = scan_to_density(
                dens_scan_rgb, stock=stock, d_min_override=d_min_override
            )
            h_total_map, conf_mask = density_to_h_total(density_map, curve, mask_mode=mask_mode)
    except Exception as exc:
        warnings.warn(f"densitometry failed: {exc}")  # Fix 3: surface degraded mode
        # Keep existing fields working; densitometry fields stay None (no dead pass)

    return PreprocessedNegative(
        rgb=rgb,
        luminance=luminance,
        log_exposure=log_exposure,
        was_inverted=was_inverted,
        original_size=original_size,
        density=density_map,
        h_total=h_total_map,
        confidence_mask=conf_mask,
        density_rgb=density_rgb,
        h_total_rgb=h_total_rgb,
        confidence_mask_rgb=conf_mask_rgb,
        is_color=is_c,
        trim_bbox_frac=trim_bbox_frac,
        d_min_override_used=d_min_override_used,
        physics_is_positive=physics_is_positive,
    )


def to_pil(rgb: np.ndarray) -> Image.Image:
    """Convert float RGB [0,1] to PIL Image."""
    arr = (np.clip(rgb, 0.0, 1.0) * 255.0).astype(np.uint8)
    return Image.fromarray(arr, mode="RGB")