tests/test_latent_optimizer.py

"""WP-0: Tests for LatentSpaceOptimizer — pixel fallback path and VAE path (slow)."""

from __future__ import annotations

import numpy as np
import pytest
import torch

from film_physics import get_film_curve
from hybrid_loss import HybridFilmLoss
from latent_optimizer import LatentSpaceOptimizer, RefinementResult


def _make_small_pair(size: int = 32) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Return (img_a, img_b, obs_rgb) all as float32 (H, W, 3) in [0.2, 0.8]."""
    rng = np.random.default_rng(42)
    img_a = rng.uniform(0.2, 0.8, (size, size, 3)).astype(np.float32)
    img_b = rng.uniform(0.2, 0.8, (size, size, 3)).astype(np.float32)
    obs_rgb = rng.uniform(0.2, 0.8, (size, size, 3)).astype(np.float32)
    return img_a, img_b, obs_rgb


def _build_obs_log_h(size: int = 32) -> torch.Tensor:
    return torch.full((1, 1, size, size), -0.3)


def _build_loss_fn() -> HybridFilmLoss:
    curve = get_film_curve("Generic")
    return HybridFilmLoss(
        film_curve=curve,
        physics_weight=1.0,
        perceptual_weight=0.0,  # skip LPIPS so pixel-fallback tests stay fast
    )


class TestPixelFallback:
    """Pixel-space optimization runs offline without a VAE."""

    def test_returns_refinement_result(self, monkeypatch):
        monkeypatch.setattr(
            LatentSpaceOptimizer, "_try_load_vae", lambda self: None
        )
        img_a, img_b, obs_rgb = _make_small_pair()
        obs_log_h = _build_obs_log_h()
        loss_fn = _build_loss_fn()

        opt = LatentSpaceOptimizer(hybrid_loss=loss_fn, steps=3, lr=0.01)
        result = opt.refine(img_a, img_b, obs_log_h, obs_rgb)

        assert isinstance(result, RefinementResult)
        assert not result.used_vae
        assert result.steps_run == 3

    def test_output_shape_matches_input(self, monkeypatch):
        monkeypatch.setattr(
            LatentSpaceOptimizer, "_try_load_vae", lambda self: None
        )
        img_a, img_b, obs_rgb = _make_small_pair(size=48)
        obs_log_h = _build_obs_log_h(size=48)
        loss_fn = _build_loss_fn()

        opt = LatentSpaceOptimizer(hybrid_loss=loss_fn, steps=2, lr=0.01)
        result = opt.refine(img_a, img_b, obs_log_h, obs_rgb)

        assert result.refined_a.shape == img_a.shape
        assert result.refined_b.shape == img_b.shape

    def test_output_values_in_range(self, monkeypatch):
        monkeypatch.setattr(
            LatentSpaceOptimizer, "_try_load_vae", lambda self: None
        )
        img_a, img_b, obs_rgb = _make_small_pair()
        obs_log_h = _build_obs_log_h()
        loss_fn = _build_loss_fn()

        opt = LatentSpaceOptimizer(hybrid_loss=loss_fn, steps=3, lr=0.01)
        result = opt.refine(img_a, img_b, obs_log_h, obs_rgb)

        assert float(result.refined_a.min()) >= 0.0 - 1e-5
        assert float(result.refined_a.max()) <= 1.0 + 1e-5
        assert float(result.refined_b.min()) >= 0.0 - 1e-5
        assert float(result.refined_b.max()) <= 1.0 + 1e-5

    def test_improved_flag(self, monkeypatch):
        monkeypatch.setattr(
            LatentSpaceOptimizer, "_try_load_vae", lambda self: None
        )
        img_a, img_b, obs_rgb = _make_small_pair()
        obs_log_h = _build_obs_log_h()
        loss_fn = _build_loss_fn()

        opt = LatentSpaceOptimizer(hybrid_loss=loss_fn, steps=10, lr=0.05)
        result = opt.refine(img_a, img_b, obs_log_h, obs_rgb)

        # The 'improved' property is just final < initial — verify it's a bool
        assert isinstance(result.improved, bool)

    def test_progress_callback_called(self, monkeypatch):
        monkeypatch.setattr(
            LatentSpaceOptimizer, "_try_load_vae", lambda self: None
        )
        img_a, img_b, obs_rgb = _make_small_pair()
        obs_log_h = _build_obs_log_h()
        loss_fn = _build_loss_fn()

        calls = []
        opt = LatentSpaceOptimizer(hybrid_loss=loss_fn, steps=4, lr=0.01)
        result = opt.refine(
            img_a, img_b, obs_log_h, obs_rgb,
            progress_callback=lambda step, total, loss: calls.append((step, total, loss)),
        )

        assert len(calls) == 4
        assert all(isinstance(l, float) for _, _, l in calls)

    def test_refine_large_nonmultiple_with_density_no_crash(self, monkeypatch):
        """Fix 1 regression: 700x900 (>max_side, not %8) with density+mask must not crash."""
        monkeypatch.setattr(
            LatentSpaceOptimizer, "_try_load_vae", lambda self: None
        )
        h, w = 700, 900
        rng = np.random.default_rng(42)
        img_a = rng.uniform(0.2, 0.8, (h, w, 3)).astype(np.float32)
        img_b = rng.uniform(0.2, 0.8, (h, w, 3)).astype(np.float32)
        obs_rgb = rng.uniform(0.2, 0.8, (h, w, 3)).astype(np.float32)
        obs_log_h = torch.full((1, 1, h, w), -0.3)
        density = np.full((h, w), 0.5, dtype=np.float32)
        conf_mask = np.full((h, w), 1, dtype=np.uint8)  # VALID

        loss_fn = _build_loss_fn()  # perceptual=0

        opt = LatentSpaceOptimizer(hybrid_loss=loss_fn, steps=1, lr=0.05, max_side=512, vae_id="nonexistent")
        result = opt.refine(
            img_a, img_b, obs_log_h, obs_rgb,
            observed_density=density, confidence_mask=conf_mask,
        )

        assert isinstance(result, RefinementResult)
        assert np.isfinite(result.initial_loss)
        assert np.isfinite(result.final_loss)
        assert result.refined_a.shape == (h, w, 3)


@pytest.mark.slow
class TestVAEPath:
    """VAE path — requires downloading model weights, marked slow/offline-skip."""

    def test_vae_path_runs(self):
        img_a, img_b, obs_rgb = _make_small_pair(size=64)
        obs_log_h = _build_obs_log_h(size=64)
        curve = get_film_curve("Generic")
        loss_fn = HybridFilmLoss(film_curve=curve, physics_weight=1.0, perceptual_weight=0.0)

        opt = LatentSpaceOptimizer(hybrid_loss=loss_fn, steps=2, lr=0.01)
        result = opt.refine(img_a, img_b, obs_log_h, obs_rgb)

        # If VAE was loaded, used_vae=True; if not available, falls back gracefully
        assert isinstance(result, RefinementResult)
        assert result.refined_a.shape == img_a.shape


class TestLatentOptimizerWP7Hardening:
    """WP-7: cosine LR, best-snapshot, early-stop on physics, degeneracy guard, g_final, new flags."""

    def test_new_refinement_result_fields_present(self, monkeypatch):
        monkeypatch.setattr(LatentSpaceOptimizer, "_try_load_vae", lambda self: None)
        img_a, img_b, obs_rgb = _make_small_pair()
        obs_log_h = _build_obs_log_h()
        loss_fn = _build_loss_fn()
        opt = LatentSpaceOptimizer(hybrid_loss=loss_fn, steps=2, lr=0.01)
        result = opt.refine(img_a, img_b, obs_log_h, obs_rgb)
        assert hasattr(result, "g_final")
        assert isinstance(result.g_final, float)
        assert hasattr(result, "early_stopped")
        assert hasattr(result, "degeneracy_aborted")
        assert result.used_density in (True, False)

    def test_early_stop_with_lr_zero(self, monkeypatch):
        """With lr=0 loss never improves => early stop after ~patience steps, flag set."""
        monkeypatch.setattr(LatentSpaceOptimizer, "_try_load_vae", lambda self: None)
        img_a, img_b, obs_rgb = _make_small_pair()
        obs_log_h = _build_obs_log_h()
        loss_fn = _build_loss_fn()
        opt = LatentSpaceOptimizer(hybrid_loss=loss_fn, steps=200, lr=0.0)
        result = opt.refine(img_a, img_b, obs_log_h, obs_rgb)
        assert result.early_stopped is True
        assert result.steps_run <= 30 + 35  # initial + patience
        assert result.steps_run < 200

    def test_g_final_reported_and_finite(self, monkeypatch):
        monkeypatch.setattr(LatentSpaceOptimizer, "_try_load_vae", lambda self: None)
        img_a, img_b, obs_rgb = _make_small_pair()
        obs_log_h = _build_obs_log_h()
        loss_fn = _build_loss_fn()
        opt = LatentSpaceOptimizer(hybrid_loss=loss_fn, steps=3, lr=0.01)
        result = opt.refine(img_a, img_b, obs_log_h, obs_rgb)
        assert np.isfinite(result.g_final)
        assert 0.01 < result.g_final < 100.0  # within reasonable after clamp

    def test_degeneracy_guard_monkey_k_and_fields(self, monkeypatch):
        """Monkey k_selection to K=2 high; guard fields exercised (trigger force demonstrated via teeth)."""
        monkeypatch.setattr(LatentSpaceOptimizer, "_try_load_vae", lambda self: None)
        img_a, img_b, obs_rgb = _make_small_pair()
        obs_log_h = _build_obs_log_h()
        loss_fn = _build_loss_fn()

        class FakeBD:
            k_selection_score = 0.9
        monkeypatch.setattr(loss_fn, "evaluate", lambda *a, **k: FakeBD())

        opt = LatentSpaceOptimizer(hybrid_loss=loss_fn, steps=4, lr=0.05)
        result = opt.refine(img_a, img_b, obs_log_h, obs_rgb)
        assert result.degeneracy_aborted in (True, False)
        assert isinstance(result.g_final, float)


    def test_best_snapshot_matches_evaluated_loss(self, monkeypatch):
        """Regression (WP-7 review bugs 1+3): the returned state must be one whose loss
        was actually evaluated. With a divergent lr the first steps explode, so a correct
        best-snapshot returns ~the initial state (final_loss <= initial_loss). The buggy
        ordering snapshotted post-step latents under the pre-step loss, returning an
        exploded state labeled with a good loss."""
        monkeypatch.setattr(LatentSpaceOptimizer, "_try_load_vae", lambda self: None)
        img_a, img_b, obs_rgb = _make_small_pair()
        obs_log_h = _build_obs_log_h()
        loss_fn = _build_loss_fn()
        opt = LatentSpaceOptimizer(hybrid_loss=loss_fn, steps=5, lr=50.0)
        result = opt.refine(img_a, img_b, obs_log_h, obs_rgb)
        assert result.final_loss <= result.initial_loss * 1.05 + 1e-6, (
            f"returned state's loss {result.final_loss:.6f} is worse than initial "
            f"{result.initial_loss:.6f} — best-snapshot returned a state that never earned its loss"
        )