densitometry.py

"""
Densitometry pipeline for scanned film negatives (MASTERPLAN WP-2 / Part III L0).

Converts an sRGB-encoded film scan into optical density D and from there into
linear total exposure H_total = 10^(f⁻¹(D)).

IMPORTANT — outputs are RELATIVE, not absolute:
  The scanner has an unknown tone curve and the Callier effect introduces a
  per-setup gain/gamma shift between collimated (densitometer) and diffuse
  (scanner) illumination.  Absolute calibration (nuisance gains g₁, g₂) is
  deferred to WP-3.  Every quantity returned here is accurate only up to a
  global scale factor.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Tuple

import numpy as np

from film_physics import (
    PiecewiseFilmCurve,
    get_film_curve,
    get_color_curves,
    ColorNegativeCurves,
    COLOR_STOCK_PRESETS,
)


# ---------------------------------------------------------------------------
# Mask label constants
# ---------------------------------------------------------------------------

TOE: int = 0       # Near D_min — noisy, H only a lower bound
VALID: int = 1     # Reliable region — curve well-constrained
SHOULDER: int = 2  # Near D_max — saturated, H only a lower bound


# ---------------------------------------------------------------------------
# sRGB transfer functions (canonical — synth/generate.py imports from here)
# ---------------------------------------------------------------------------

def srgb_to_linear(img: np.ndarray) -> np.ndarray:
    """Inverse sRGB EOTF (IEC 61966-2-1). Input/output in [0, 1]."""
    img = np.clip(img, 0.0, 1.0)
    return np.where(
        img <= 0.04045,
        img / 12.92,
        ((img + 0.055) / 1.055) ** 2.4,
    ).astype(np.float32)


def linear_to_srgb(img: np.ndarray) -> np.ndarray:
    """sRGB forward EOTF. Input/output in [0, 1]."""
    img = np.clip(img, 0.0, 1.0)
    return np.where(
        img <= 0.0031308,
        12.92 * img,
        1.055 * img ** (1.0 / 2.4) - 0.055,
    ).astype(np.float32)


def luminance_from_linear(lin_rgb: np.ndarray) -> np.ndarray:
    """Rec. 709 luminance from linear RGB (H, W, 3) → (H, W) float32."""
    return (
        0.2126 * lin_rgb[..., 0]
        + 0.7152 * lin_rgb[..., 1]
        + 0.0722 * lin_rgb[..., 2]
    ).astype(np.float32)


def phi_display(rgb: np.ndarray) -> np.ndarray:
    """φ = luminance_from_linear(srgb_to_linear(·)) — canonical display→exposure proxy.

    Single public φ used by demix, asymmetric recovery, and fullres (WP-14.1 P1).
    """
    return luminance_from_linear(srgb_to_linear(np.asarray(rgb, dtype=np.float32)))


# ---------------------------------------------------------------------------
# scan_to_density
# ---------------------------------------------------------------------------

def scan_to_density(
    scan_srgb: np.ndarray,
    stock: str = "Generic",
    white_level: float | None = None,
    d_min_override: float | None = None,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Convert an sRGB-encoded negative scan to optical density.

    Algorithm:
      1. Linearise via inverse sRGB EOTF.
      2. D_physical = −log₁₀(T_lum), assuming a calibrated linear scan where
         pixel value 1.0 = open-gate (no film) light.

    NOTE (found in Fable review): with the default estimated white point,
    D_obs + D_white = −log₁₀(lum/T_white) − log₁₀(T_white) = −log₁₀(lum) —
    the white-point term cancels by construction. That is correct for
    calibrated linear scans (incl. the WP-1 synthetic fixtures). Real
    auto-exposed scanners need a different anchor: pass ``d_min_override``
    (typically the stock's preset D_min) to pin the clearest film area,
    which does NOT cancel. WP-11 (real-scan intake) should choose the mode
    from scan metadata/heuristics.

    Args:
        scan_srgb:      (H, W, 3) float32 sRGB-encoded scan in [0, 1].
        stock:          Film stock name (for preset D_min fallback).
        white_level:    If provided, use as T_white directly (skip estimation).
        d_min_override: If provided, use as D_min instead of the white-point
                        estimate (useful for bench-calibrated scanners).

    Returns:
        d_physical:  (H, W) float32 — optical density in D units.
        scan_linear: (H, W, 3) float32 — linearised scan (for diagnostics).

    Note: outputs are RELATIVE (see module docstring).
    """
    scan_linear = srgb_to_linear(scan_srgb)  # (H, W, 3)
    lum = luminance_from_linear(scan_linear)  # (H, W)

    # --- White-point estimation ---
    if white_level is None:
        flat = lum[lum > 1e-6]
        if flat.size == 0:
            flat = lum.ravel()
        t_white = float(np.percentile(flat, 99.5))
        # Sanity check: if nearly all pixels are very dark, fall back to preset
        if t_white < 0.05:
            curve_preset = get_film_curve(stock)
            t_white = float(10.0 ** (-float(curve_preset.d_min)))
    else:
        t_white = float(white_level)

    t_white = max(t_white, 1e-6)

    # --- Density from white-point-normalised transmittance ---
    d_white: float
    if d_min_override is not None:
        d_white = float(d_min_override)
    else:
        d_white = float(-np.log10(t_white))  # estimated D_min from the scan

    t_norm = np.clip(lum / t_white, 1e-8, 1.0)
    d_obs = -np.log10(t_norm)               # relative density [0, ...]
    d_physical = (d_obs + d_white).astype(np.float32)

    return d_physical, scan_linear


# ---------------------------------------------------------------------------
# density_to_h_total
# ---------------------------------------------------------------------------

# Slope fraction below which the curve no longer meaningfully encodes exposure:
# where dD/dlogH < SLOPE_VALID_FRAC * gamma, a density step of one 8-bit JPEG code
# maps to a multi-stop H error, so the pixel belongs in TOE/SHOULDER, not VALID.
SLOPE_VALID_FRAC = 0.25


_SLOPE_BOUNDS_CACHE: dict = {}


def _slope_valid_bounds(curve: PiecewiseFilmCurve) -> Tuple[float, float]:
    """(D_lo, D_hi) between which the curve's local slope >= SLOPE_VALID_FRAC*gamma.

    Cached on the curve object. Fixes the density-margin mask defect (2026-07-16
    finding): 5%-of-density-range margins leave "VALID" spanning ~11 stops on
    Portra because the curve is asymptotically flat near d_max — D=1.25 -> H=19.7
    but D=1.2855 -> H=69.9. Slope is the honest reliability criterion.
    """
    # WP-18 D3b: get_film_curve constructs a FRESH curve per call, so a per-object
    # attribute cache never hits. Key the cache on the curve's parameter tuple —
    # shared across all instances of the same preset (and immune to threading:
    # worst case two threads compute the same value once).
    key = (
        float(curve.gamma), float(curve.d_min), float(curve.d_max),
        float(curve.toe_strength), float(curve.shoulder_strength),
        float(curve.toe_width), float(curve.shoulder_width),
    )
    cached = _SLOPE_BOUNDS_CACHE.get(key)
    if cached is not None:
        return cached
    import torch

    log_h = torch.linspace(-4.5, 3.5, 2048)
    with torch.no_grad():
        d = curve.forward(log_h).cpu().numpy().astype(np.float64)
    slope = np.gradient(d, log_h.numpy().astype(np.float64))
    ok = slope >= SLOPE_VALID_FRAC * float(curve.gamma)
    if ok.any():
        idx = np.nonzero(ok)[0]
        bounds = (float(d[idx[0]]), float(d[idx[-1]]))
    else:  # degenerate curve: fall back to the full density range
        bounds = (float(curve.d_min), float(curve.d_max))
    _SLOPE_BOUNDS_CACHE[key] = bounds
    return bounds


def density_to_h_total(
    d_physical: np.ndarray,
    curve: PiecewiseFilmCurve,
    mask_mode: str = "density_margin",
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Convert optical density to linear total exposure via the inverse H-D curve.

    Args:
        d_physical: (H, W) float32 optical density array.
        curve:      Film characteristic curve (must have ``inverse()`` method).
        mask_mode:  "density_margin" (default — WP-2 contract, 5% margins of the
                    density range, byte-identical legacy behavior) or "slope"
                    (VALID only where the curve's local slope >= 25% of gamma, so
                    the mask reflects actual H reliability; the app path opts in).

    Returns:
        h_total:         (H, W) float32 — linear exposure (relative, up to global scale).
        confidence_mask: (H, W) uint8 — TOE=0 / VALID=1 / SHOULDER=2.

    Note: outputs are RELATIVE (see module docstring).
    """
    d_min = float(curve.d_min)
    d_max = float(curve.d_max)

    # Invert the characteristic curve
    d_clamped = np.clip(d_physical, d_min, d_max).astype(np.float32)
    log_h = curve.inverse(d_clamped)  # (H, W) float32
    h_total = np.power(10.0, log_h).astype(np.float32)

    if mask_mode == "slope":
        d_lo, d_hi = _slope_valid_bounds(curve)
    else:
        # Confidence mask — 5% margins from d_min / d_max (legacy WP-2 contract)
        margin = 0.05 * (d_max - d_min)
        d_lo, d_hi = d_min + margin, d_max - margin
    mask = np.where(
        d_physical < d_lo, TOE,
        np.where(d_physical > d_hi, SHOULDER, VALID),
    ).astype(np.uint8)

    return h_total, mask


# ---------------------------------------------------------------------------
# Physics polarity — WP-11.1 post-review, the ONE place the policy lives
# ---------------------------------------------------------------------------

def prepare_densitometry_input(
    scan_srgb: np.ndarray,
    stock: str,
    positive_source: bool = False,
) -> Tuple[np.ndarray, ColorNegativeCurves | None]:
    """Canonical physics-polarity contract for every densitometry entry point.

    Densitometry assumes NEGATIVE polarity (dark = dense = high exposure). A
    lab-inverted positive is un-inverted in sRGB display space — the lab
    inversion is an involution, so ``1 − pos`` reconstructs the negative scan
    byte-exactly on uint8 (WP-11.1). Color positives already lack the orange
    mask (the lab removed it), so the returned curves carry zero mask offsets.

    For negative-polarity input the array is returned UNCHANGED (the same
    object), keeping the default synthetic path byte-identical.

    Returns:
        (dens_scan, curves) — ``curves`` is None for B&W stocks; for color
        stocks it is the object to pass as ``curves=`` to
        ``scan_to_density_rgb`` / ``density_to_h_total_rgb``.
    """
    dens_scan = (1.0 - scan_srgb) if positive_source else scan_srgb
    curves: ColorNegativeCurves | None = None
    if stock in COLOR_STOCK_PRESETS:
        curves = get_color_curves(stock)
        if positive_source:
            curves = ColorNegativeCurves(
                r=curves.r, g=curves.g, b=curves.b,
                mask_offset_rgb=(0.0, 0.0, 0.0),
            )
    return dens_scan, curves


# ---------------------------------------------------------------------------
# RGB (color negative) densitometry — WP-8, ADDITIVE ONLY
# Existing scalar functions and signatures untouched.
# RGB versions LOOP the scalar ones per channel; no reimplementation of math.
# ---------------------------------------------------------------------------

def scan_to_density_rgb(
    scan_srgb: np.ndarray,
    color_stock: str,
    white_level: float | None = None,
    curves: ColorNegativeCurves | None = None,
    d_min_override: float | None = None,
) -> np.ndarray:
    """Per-channel density for color negative.

    Default (``d_min_override is None``): absolute density path
      lin = srgb_to_linear(scan)
      d_abs_c = −log10(clip(lin_c, 1e-6, 1.0))
      d_phys_c = clip(d_abs_c − mask_offset_c, 0, None)

    When ``d_min_override`` is set (WP-11 auto_exposed): relative densitometry
    per channel with the stock D_min anchor (same cancel-safe path as scalar
    ``scan_to_density``), then subtract mask offsets. Signature-compatible;
    default path is byte-identical when override is None.
    """
    if curves is None:
        curves = get_color_curves(color_stock)
    offsets = curves.mask_offset_rgb
    d_rgb = np.zeros(scan_srgb.shape, dtype=np.float32)

    if d_min_override is not None:
        # Auto-exposed scanners: per-channel relative density + stock D_min
        for c in range(3):
            mono = np.stack([scan_srgb[..., c]] * 3, axis=-1)
            d_c, _ = scan_to_density(
                mono,
                stock=color_stock,
                white_level=white_level,
                d_min_override=float(d_min_override),
            )
            d_rgb[..., c] = np.clip(d_c - offsets[c], 0.0, None).astype(np.float32)
        return d_rgb

    # Default absolute path (unchanged for synthetic / linear calibration)
    scan_linear = srgb_to_linear(scan_srgb)  # (H, W, 3)
    for c in range(3):
        lin_c = scan_linear[..., c]
        d_abs_c = -np.log10(np.clip(lin_c, 1e-6, 1.0))
        d_phys_c = np.clip(d_abs_c - offsets[c], 0.0, None)
        d_rgb[..., c] = d_phys_c.astype(np.float32)

    return d_rgb


def density_to_h_total_rgb(
    d_rgb: np.ndarray,
    curves: ColorNegativeCurves,
    mask_mode: str = "density_margin",
) -> Tuple[np.ndarray, np.ndarray]:
    """Per-channel h_total and confidence (loops existing density_to_h_total)."""
    h_list = []
    conf_list = []
    for c, curve in enumerate([curves.r, curves.g, curves.b]):
        h_c, m_c = density_to_h_total(d_rgb[..., c], curve, mask_mode=mask_mode)
        h_list.append(h_c)
        conf_list.append(m_c)
    h_total_rgb = np.stack(h_list, axis=2).astype(np.float32)
    conf_rgb = np.stack(conf_list, axis=2).astype(np.uint8)
    return h_total_rgb, conf_rgb


def combine_confidence_rgb(mask_rgb: np.ndarray) -> np.ndarray:
    """Scalar confidence: TOE if any channel TOE; SHOULDER if any SHOULDER; else VALID."""
    has_toe = np.any(mask_rgb == TOE, axis=2)
    has_shoulder = np.any(mask_rgb == SHOULDER, axis=2)
    out = np.full(mask_rgb.shape[:2], VALID, dtype=np.uint8)
    out[has_toe] = TOE
    out[has_shoulder & ~has_toe] = SHOULDER
    return out