latent_optimizer.py

"""Gradient-based refinement of separation candidates in VAE latent space."""

from __future__ import annotations

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

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim.lr_scheduler import CosineAnnealingLR

from hybrid_loss import HybridFilmLoss, GAIN_GRID_RANGE
from densitometry import VALID, srgb_to_linear, luminance_from_linear

DEFAULT_VAE_ID = "stabilityai/sd-vae-ft-mse"

_VAE_CACHE: dict[str, nn.Module] = {}


def _get_shared_vae(model_id: str) -> nn.Module:
    """Return a cached, frozen ``AutoencoderKL`` (lazy-loaded once per id)."""
    if model_id not in _VAE_CACHE:
        from diffusers import AutoencoderKL

        vae = AutoencoderKL.from_pretrained(model_id)
        vae.eval()
        vae.requires_grad_(False)
        _VAE_CACHE[model_id] = vae
    return _VAE_CACHE[model_id]


def _pick_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def _rgb_numpy_to_tensor(rgb: np.ndarray) -> torch.Tensor:
    """Float RGB (H, W, 3) in [0, 1] -> (1, 3, H, W) tensor."""
    return torch.from_numpy(rgb).float().permute(2, 0, 1).unsqueeze(0)


def _tensor_to_rgb_numpy(t: torch.Tensor) -> np.ndarray:
    """(1, 3, H, W) tensor in [0, 1] -> float RGB (H, W, 3)."""
    return t.squeeze(0).permute(1, 2, 0).clamp(0.0, 1.0).cpu().numpy().astype(np.float32)


def _downscale(t: torch.Tensor, max_side: int) -> torch.Tensor:
    h, w = t.shape[-2:]
    scale = min(1.0, max_side / max(h, w))
    if scale >= 1.0:
        return t
    size = (max(1, round(h * scale)), max(1, round(w * scale)))
    return F.interpolate(t, size=size, mode="bilinear", align_corners=False, antialias=True)


def _pad_to_multiple(t: torch.Tensor, multiple: int = 8, *, value: int = 0, mode: str = "replicate") -> Tuple[torch.Tensor, Tuple[int, int]]:
    """Right/bottom pad so H and W are multiples of ``multiple``.
    mode: 'replicate' (default) or 'constant' (uses value for pad, e.g. TOE=0 for mask).
    """
    h, w = t.shape[-2:]
    ph = (multiple - h % multiple) % multiple
    pw = (multiple - w % multiple) % multiple
    if ph or pw:
        if mode == "constant":
            t = F.pad(t, (0, pw, 0, ph), mode="constant", value=value)
        else:
            t = F.pad(t, (0, pw, 0, ph), mode="replicate")
    return t, (h, w)


@dataclass
class RefinementResult:
    """Output of a latent-space refinement run."""

    refined_a: np.ndarray  # float RGB [0, 1], original resolution
    refined_b: np.ndarray
    initial_loss: float
    final_loss: float
    steps_run: int
    used_vae: bool
    loss_history: List[float] = field(default_factory=list)
    used_density: bool = False  # Fix 3
    g_final: float = 1.0
    early_stopped: bool = False
    degeneracy_aborted: bool = False

    @property
    def improved(self) -> bool:
        return self.final_loss < self.initial_loss


class LatentSpaceOptimizer:
    """
    Refines a candidate scene pair by gradient descent on the hybrid loss.

    Both reconstructions are encoded into the latent space of a pretrained
    VAE (from ``diffusers``); Adam then optimizes the latents so that the
    decoded pair better explains the observed double exposure under the
    differentiable film-curve forward model plus LPIPS.

    If the VAE cannot be loaded (no ``diffusers`` install or no network for
    the initial weight download), the optimizer transparently falls back to
    pixel-space optimization with a sigmoid parameterization, so demo mode
    keeps working offline.
    """

    def __init__(
        self,
        hybrid_loss: HybridFilmLoss,
        vae_id: str = DEFAULT_VAE_ID,
        steps: int = 40,
        lr: float = 0.05,
        max_side: int = 512,
        device: Optional[torch.device] = None,
    ):
        self.hybrid_loss = hybrid_loss
        self.vae_id = vae_id
        self.steps = max(1, min(int(steps), 300))
        self.lr = lr
        self.max_side = max_side
        self.device = device or _pick_device()

    # ------------------------------------------------------------------
    # Latent <-> image codecs
    # ------------------------------------------------------------------

    def _try_load_vae(self) -> Optional[nn.Module]:
        try:
            return _get_shared_vae(self.vae_id).to(self.device)
        except Exception:
            return None

    @staticmethod
    def _vae_encode(vae: nn.Module, image01: torch.Tensor) -> torch.Tensor:
        with torch.no_grad():
            return vae.encode(image01 * 2.0 - 1.0).latent_dist.mean

    @staticmethod
    def _vae_decode(vae: nn.Module, latents: torch.Tensor) -> torch.Tensor:
        return ((vae.decode(latents).sample + 1.0) * 0.5).clamp(0.0, 1.0)

    @staticmethod
    def _pixel_encode(image01: torch.Tensor) -> torch.Tensor:
        return torch.logit(image01.clamp(1e-3, 1.0 - 1e-3))

    @staticmethod
    def _pixel_decode(params: torch.Tensor) -> torch.Tensor:
        return torch.sigmoid(params)

    # ------------------------------------------------------------------
    # Refinement
    # ------------------------------------------------------------------

    def refine(
        self,
        recon_a: np.ndarray,
        recon_b: np.ndarray,
        observed_log_exposure: torch.Tensor,
        observed_rgb: np.ndarray,
        progress_callback: Optional[Callable[[int, int, float], None]] = None,
        observed_density: Optional[np.ndarray] = None,
        confidence_mask: Optional[np.ndarray] = None,
    ) -> RefinementResult:
        """
        Refine a scene pair to minimize the hybrid physics + perceptual loss.

        Args:
            recon_a: Candidate scene A, float RGB (H, W, 3) in [0, 1].
            recon_b: Candidate scene B, same format.
            observed_log_exposure: (1, 1, H, W) observed log-exposure map
                (``PreprocessedNegative.log_exposure``).
            observed_rgb: Observed positive scan, float RGB (H, W, 3) in [0, 1].
            progress_callback: Optional ``(step, total_steps, loss)`` hook.
            observed_density: Optional measured density map from densitometry (WP-3).
            confidence_mask: Optional confidence mask (VALID etc).

        Returns:
            ``RefinementResult`` with refined images at the input resolution.
        """
        device = self.device
        orig_h, orig_w = recon_a.shape[:2]

        # Work at reduced resolution for tractable LPIPS + VAE gradients.
        a = _downscale(_rgb_numpy_to_tensor(recon_a).to(device), self.max_side)
        b = _downscale(_rgb_numpy_to_tensor(recon_b).to(device), self.max_side)
        obs_rgb = _downscale(_rgb_numpy_to_tensor(observed_rgb).to(device), self.max_side)
        obs_log_h = _downscale(observed_log_exposure.float().to(device), self.max_side)

        # Pre-pad working size for density/mask downscale (Fix 1)
        pre_h, pre_w = a.shape[-2:]

        dens_t = None
        conf_t = None
        if observed_density is not None:
            d_t = torch.from_numpy(observed_density.astype(np.float32)).unsqueeze(0).unsqueeze(0).to(device)
            dens_t = F.interpolate(d_t, size=(pre_h, pre_w), mode="bilinear", antialias=True)
            if confidence_mask is not None:
                m_t = torch.from_numpy(confidence_mask.astype(np.float32)).unsqueeze(0).unsqueeze(0).to(device)
                m_down = F.interpolate(m_t, size=(pre_h, pre_w), mode="nearest")
                conf_t = m_down.round().to(torch.long)
            else:
                # synthesize all-VALID so pad region can be excluded
                conf_t = torch.full((1, 1, pre_h, pre_w), VALID, dtype=torch.long, device=device)
        elif confidence_mask is not None:
            # unlikely but handle
            m_t = torch.from_numpy(confidence_mask.astype(np.float32)).unsqueeze(0).unsqueeze(0).to(device)
            conf_t = F.interpolate(m_t, size=(pre_h, pre_w), mode="nearest").round().to(torch.long)

        a, unpadded = _pad_to_multiple(a)
        b, _ = _pad_to_multiple(b)
        obs_rgb, _ = _pad_to_multiple(obs_rgb)
        obs_log_h, _ = _pad_to_multiple(obs_log_h)

        # Pad dens (replicate) and mask with constant TOE=0 (using extended _pad_to_multiple)
        if dens_t is not None:
            dens_t, _ = _pad_to_multiple(dens_t)
        if conf_t is not None:
            conf_t, _ = _pad_to_multiple(conf_t, value=0, mode="constant")

        # K-sel + initial share for degeneracy guard (I.6, on input pair) — do on CPU before moving loss/curve
        with torch.no_grad():
            init_bd = self.hybrid_loss.evaluate(
                observed_log_exposure,
                observed_rgb,
                recon_a,
                recon_b,
                density=observed_density,
                confidence_mask=confidence_mask,
            )
            k_sel = float(init_bd.k_selection_score)
            la0 = luminance_from_linear(srgb_to_linear(recon_a))
            lb0 = luminance_from_linear(srgb_to_linear(recon_b))
            ea0 = float(np.mean(la0) + 1e-12)
            eb0 = float(np.mean(lb0) + 1e-12)
            s0 = min(ea0, eb0) / (ea0 + eb0)

        loss_fn = self.hybrid_loss.to(device)

        vae = self._try_load_vae()
        used_vae = vae is not None
        if used_vae:
            lat_a = self._vae_encode(vae, a).clone().requires_grad_(True)
            lat_b = self._vae_encode(vae, b).clone().requires_grad_(True)
            decode = lambda lat: self._vae_decode(vae, lat)  # noqa: E731
        else:
            lat_a = self._pixel_encode(a).clone().requires_grad_(True)
            lat_b = self._pixel_encode(b).clone().requires_grad_(True)
            decode = self._pixel_decode

        # --- WP-7 g (learnable nuisance gain) ---
        log10_g = torch.zeros((), device=device, requires_grad=True)

        def compute_loss(g_t: torch.Tensor) -> torch.Tensor:
            dec_a = decode(lat_a)
            dec_b = decode(lat_b)
            return loss_fn.forward_tensor(obs_log_h, obs_rgb, dec_a, dec_b, density=dens_t, confidence_mask=conf_t, g=g_t)

        with torch.no_grad():
            g_init = torch.pow(10.0, log10_g)
            initial_loss = float(compute_loss(g_init).item())

        # Adam over latents + log10_g (so g is optimized jointly)
        optimizer = torch.optim.Adam([lat_a, lat_b, log10_g], lr=self.lr)
        scheduler = CosineAnnealingLR(optimizer, T_max=self.steps)

        history: List[float] = []
        best_loss = initial_loss
        best_state = (lat_a.detach().clone(), lat_b.detach().clone())
        best_g_log = log10_g.detach().clone()

        # Early stop on *masked physics residual*
        patience = 30
        min_delta = 1e-5
        best_phys = float("inf")
        no_improve_phys = 0
        early_stopped = False
        degeneracy_aborted = False

        # local torch linearize + lum for current share (no extra import)
        def _lin_lum_t(rgb01: torch.Tensor) -> torch.Tensor:
            t = rgb01.clamp(0.0, 1.0)
            lin = torch.where(
                t <= 0.04045,
                t / 12.92,
                ((t + 0.055) / 1.055) ** 2.4,
            )
            return 0.2126 * lin[:, 0:1] + 0.7152 * lin[:, 1:2] + 0.0722 * lin[:, 2:3]

        for step in range(self.steps):
            optimizer.zero_grad()
            g = torch.pow(10.0, log10_g)
            loss = compute_loss(g)

            # Snapshot BEFORE stepping: this loss belongs to the CURRENT latents/g,
            # so the stored best state is exactly the state that earned the loss.
            # (Snapshotting after step() captured a different — possibly collapsed —
            # state under a stale good loss, defeating the guard's rollback.)
            loss_val = float(loss.item())
            history.append(loss_val)
            if loss_val < best_loss:
                best_loss = loss_val
                best_state = (lat_a.detach().clone(), lat_b.detach().clone())
                best_g_log = log10_g.detach().clone()

            loss.backward()
            optimizer.step()
            scheduler.step()

            # clamp log10_g after step (keeps g in documented range)
            with torch.no_grad():
                log10_g.data.clamp_(GAIN_GRID_RANGE[0], GAIN_GRID_RANGE[1])

            # physics residual for early stop: the RAW masked physics term of the
            # post-step state, independent of the caller's term weights
            with torch.no_grad():
                g_now = torch.pow(10.0, log10_g)
                dec_a = decode(lat_a)
                dec_b = decode(lat_b)
                weights = ("physics_weight", "perceptual_weight", "exclusivity_weight", "balance_weight", "naturalness_weight")
                saved = {w: getattr(loss_fn, w) for w in weights}
                try:
                    for w in weights:
                        setattr(loss_fn, w, 0.0)
                    loss_fn.physics_weight = 1.0
                    phys_val = float(compute_loss(g_now).item())
                finally:
                    for w, v in saved.items():
                        setattr(loss_fn, w, v)

            if phys_val + min_delta < best_phys:
                best_phys = phys_val
                no_improve_phys = 0
            else:
                no_improve_phys += 1
            if no_improve_phys >= patience:
                early_stopped = True
                break

            # Degeneracy guard on the post-step state with the post-step g
            # (only when K=2 hypothesis and started balanced)
            if k_sel >= 0.5 and s0 >= 0.10:
                with torch.no_grad():
                    la = _lin_lum_t(dec_a)
                    lb = _lin_lum_t(dec_b)
                    ea = float((g_now * la).mean().item() + 1e-12)
                    eb = float(((1.0 / g_now) * lb).mean().item() + 1e-12)
                    share = min(ea, eb) / (ea + eb)
                if share < 0.05:
                    degeneracy_aborted = True
                    break

            if progress_callback is not None:
                progress_callback(step + 1, self.steps, loss_val)

        steps_run = (step + 1) if (early_stopped or degeneracy_aborted) else self.steps

        with torch.no_grad():
            # Last-chance snapshot: the post-loop state was stepped but never scored.
            # Skipped on abort so the rollback contract (return pre-collapse best) holds.
            if not degeneracy_aborted:
                tail_val = float(compute_loss(torch.pow(10.0, log10_g)).item())
                if tail_val < best_loss:
                    best_loss = tail_val
                    best_state = (lat_a.detach().clone(), lat_b.detach().clone())
                    best_g_log = log10_g.detach().clone()

            # Restore the best state into the latents so compute_loss (whose closure
            # decodes lat_a/lat_b) scores exactly the images we return.
            lat_a.data.copy_(best_state[0])
            lat_b.data.copy_(best_state[1])
            log10_g.data.copy_(best_g_log)
            g_final = float(10.0 ** best_g_log)
            dec_a = decode(lat_a)
            dec_b = decode(lat_b)
            final_loss = float(
                compute_loss(torch.tensor(g_final, device=device)).item()
            )
            h, w = unpadded
            dec_a = dec_a[..., :h, :w]
            dec_b = dec_b[..., :h, :w]
            if (h, w) != (orig_h, orig_w):
                dec_a = F.interpolate(
                    dec_a, size=(orig_h, orig_w), mode="bilinear", align_corners=False
                )
                dec_b = F.interpolate(
                    dec_b, size=(orig_h, orig_w), mode="bilinear", align_corners=False
                )

        return RefinementResult(
            refined_a=_tensor_to_rgb_numpy(dec_a),
            refined_b=_tensor_to_rgb_numpy(dec_b),
            initial_loss=initial_loss,
            final_loss=final_loss,
            steps_run=steps_run,
            used_vae=used_vae,
            loss_history=history,
            used_density=(observed_density is not None),
            g_final=g_final,
            early_stopped=early_stopped,
            degeneracy_aborted=degeneracy_aborted,
        )


# Public aliases for the working-resolution helpers so other optimization-style
# modules (baselines/double_dip.py) can reuse the exact downscale/pad/device
# contract without depending on private names.
pick_device = _pick_device
rgb_numpy_to_tensor = _rgb_numpy_to_tensor
tensor_to_rgb_numpy = _tensor_to_rgb_numpy
downscale_to_max_side = _downscale
pad_to_multiple = _pad_to_multiple