noaa-goes-visualization/tests/test_model.py

import numpy as np
import pytest

torch = pytest.importorskip("torch")

from suvi import model  # noqa: E402

BANDS = model.BANDS


def conditioning(states, satellites=None, dts=None):
    """Build a (1, S, COND_DIM) conditioning tensor."""
    count = len(states)
    satellites = satellites or [True] * count
    dts = dts if dts is not None else [(i - count // 2) * 240.0 for i in range(count)]
    rows = [
        model.frame_conditioning(states[i], satellites[i], dts[i])
        for i in range(count)
    ]
    return torch.stack(rows)[None]


def expand(net, frames, condition, present=None):
    """The stack with the model's trend candidates appended, as the blend sees it."""
    trends, condition, present = net._trends(frames, condition, present)
    return torch.cat([frames, trends], dim=1), condition, present


def stack(count, size=64, seed=0):
    generator = torch.Generator().manual_seed(seed)
    return torch.randn(1, count, BANDS, size, size, generator=generator)


@pytest.fixture(scope="module")
def net():
    torch.manual_seed(0)
    return model.build(base=16, depth=3).eval()


# ------------------------------------------------------------------- conditioning


def test_conditioning_is_the_declared_width():
    assert model.frame_conditioning("available", True, 0.0).shape == (model.COND_DIM,)


def test_the_layout_names_every_column_once():
    assert model.COND_DIM == len(model.COND_LAYOUT) == len(set(model.COND_LAYOUT))
    assert [model.AVAILABLE, model.MISSING, model.TARGET_SATELLITE, model.DT_SIGN,
            model.DT_LOG, model.AGREE_MEAN, model.AGREE_MAX, model.DERIVED] == list(
        range(model.COND_DIM))


def test_conditioning_separates_the_two_states():
    available = model.frame_conditioning("available", True, 240.0)
    missing = model.frame_conditioning("missing", True, 240.0)
    assert available[model.AVAILABLE] == 1.0 and available[model.MISSING] == 0.0
    assert missing[model.AVAILABLE] == 0.0 and missing[model.MISSING] == 1.0


@pytest.mark.parametrize("state", ["suspect", "probably-fine"])
def test_conditioning_rejects_an_unknown_state(state):
    """v9 has no `suspect`: a judged-bad band is withheld, never shown."""
    with pytest.raises(KeyError):
        model.frame_conditioning(state, True, 0.0)


def test_conditioning_distinguishes_direction_and_distance():
    back = model.frame_conditioning("available", True, -7200.0)
    forward = model.frame_conditioning("available", True, 7200.0)
    far = model.frame_conditioning("available", True, 72000.0)
    assert back[model.DT_SIGN] == -1.0 and forward[model.DT_SIGN] == 1.0
    assert back[model.DT_LOG] == pytest.approx(forward[model.DT_LOG])  # magnitude only
    assert far[model.DT_LOG] > forward[model.DT_LOG]


def test_conditioning_leaves_agreement_to_the_model():
    """The agreement columns are computed from the stack inside forward(),
    identically at training and inference -- the fix for detector scores that were
    always zero in training and populated on the bench."""
    vector = model.frame_conditioning("available", True, 240.0)
    assert vector[model.AGREE_MEAN] == 0.0 and vector[model.AGREE_MAX] == 0.0


# ------------------------------------------------------------------- blend prior


def test_prior_favours_the_nearest_frame_in_time():
    condition = conditioning(["available"] * 3, satellites=[True] * 3,
                             dts=[-240.0, -3840.0, -240.0 * 300])
    prior = model.blend_prior(condition)[0]
    assert prior[0] > prior[1] > prior[2]


def test_prior_no_longer_penalises_the_counterpart():
    """The old prior down-weighted cross-satellite frames because their gain was
    wrong at initialisation.  Candidates now arrive gain-matched by suvi.align, so
    a simultaneous counterpart is as trustworthy as a local frame."""
    same = conditioning(["available"], satellites=[True], dts=[0.0])
    cross = conditioning(["available"], satellites=[False], dts=[0.0])
    assert model.blend_prior(same)[0, 0] == model.blend_prior(cross)[0, 0]


def test_prior_disfavours_derived_candidates():
    real = conditioning(["available"], dts=[-240.0])
    derived = real.clone()
    derived[..., model.DERIVED] = 1.0
    gap = model.blend_prior(real)[0, 0] - model.blend_prior(derived)[0, 0]
    assert float(gap) == pytest.approx(model.PRIOR_DERIVED_PENALTY)


# ------------------------------------------------------------------------- shapes


def test_forward_returns_the_source_shape(net):
    out = net(stack(5), conditioning(["available"] * 5))
    assert out.shape == (1, BANDS, 64, 64)
    assert torch.isfinite(out).all()


@pytest.mark.parametrize("size", [160, 320, 640])
def test_one_model_serves_every_source_resolution(net, size):
    """Blend and polish are fields: they upsample, so 1280 needs no retraining."""
    sources = torch.randn(1, 3, BANDS, size, size)
    out = net(sources, conditioning(["available"] * 3))
    assert out.shape == (1, BANDS, size, size)
    assert torch.isfinite(out).all()


@pytest.mark.parametrize("count", [1, 2, 7, 23])
def test_stack_depth_is_not_fixed(net, count):
    """Availability varies per slot, so the stack cannot be a compile-time constant."""
    out = net(stack(count), conditioning(["available"] * count))
    assert out.shape == (1, BANDS, 64, 64)


# --------------------------------------------------------------- identity behaviour


def expected_blend(sources, condition):
    """What an untrained model must output: the prior-weighted blend of its sources.

    Both heads are zero-initialised, so the polish gain is exactly one and the blend
    logits are exactly the prior."""
    valid = condition[..., model.AVAILABLE]
    logits = model.blend_prior(condition).masked_fill(valid < 0.5, float("-inf"))
    weights = torch.nan_to_num(torch.softmax(logits, dim=1), nan=0.0)
    return (sources * weights[..., None, None, None]).sum(dim=1)


def test_untrained_model_is_the_prior_weighted_blend(net):
    """The wiring gate: any scaling or orientation bug in the composition shows up
    here, and is invisible once training starts moving the loss.  The stack is
    expanded through the trend builder first: derived candidates are inputs to the
    blend like any other, and their prior penalty is part of the wiring."""
    frames = stack(5)
    condition = conditioning(["available"] * 5)
    out = net(frames, condition)
    expanded_frames, expanded_condition, _ = expand(net, frames, condition)
    assert torch.allclose(out, expected_blend(expanded_frames, expanded_condition),
                          atol=1e-4)


def test_identity_holds_when_fields_are_upsampled(net):
    """A 160-resolution field applied to larger sources must still be the identity."""
    sources = torch.randn(1, 3, BANDS, 320, 320)
    condition = conditioning(["available"] * 3)
    out = net(sources, condition)
    assert torch.allclose(out, expected_blend(sources, condition), atol=1e-3)


def test_missing_frames_take_no_weight(net):
    """A frame with no pixels must not contribute, however confident the head is."""
    frames = stack(4)
    frames[:, 0] = 999.0            # would dominate any blend that included it
    condition = conditioning(["missing", "available", "available", "available"])
    out = net(frames, condition)
    assert torch.allclose(out, expected_blend(frames, condition), atol=1e-4)
    assert out.abs().max() < 100, "the masked frame leaked into the blend"


def test_the_targets_own_bands_win_only_where_it_holds_them(net):
    """The target slot is an entry at dt=0 holding its undamaged bands: the prior
    trusts it most where it has pixels, and it takes no weight in a lost band."""
    frames = torch.stack([torch.full((BANDS, 32, 32), v) for v in (0.9, 0.2, 0.3)])[None]
    condition = conditioning(["available"] * 3, satellites=[True, True, True],
                             dts=[0.0, -240.0, 240.0])
    present = torch.ones(1, 3, BANDS)
    present[0, 0, 2] = 0.0                              # the target lost band 2
    out = net(frames, condition, present)
    assert torch.allclose(out, expected_band_blend(frames, condition, present),
                          atol=1e-4)
    assert float(out[0, 2].detach().max()) < 0.31       # nothing of 0.9 in band 2
    assert float(out[0, 0].detach().min()) > 0.5        # mostly the target's own


def test_a_stack_with_nothing_usable_returns_zeros_not_nan(net):
    """The sampler declines such targets; the network itself must still not NaN."""
    out = net(stack(4), conditioning(["missing"] * 4))
    assert torch.isfinite(out).all()
    assert out.abs().max() == 0.0


def test_heads_start_at_zero():
    """Training must start exactly at the hand-written policy; every step afterwards
    is a departure the data paid for."""
    torch.manual_seed(0)
    fresh = model.build(base=8, depth=2)
    for head in (fresh.weight, fresh.polish):
        with torch.no_grad():
            assert float(head.weight.abs().sum()) == 0.0
            assert float(head.bias.abs().sum()) == 0.0


# --------------------------------------------------------- structural non-collapse


def test_output_is_a_bounded_combination_whatever_the_weights():
    """The invariant that makes the old failure mode unreachable: with *arbitrary*
    parameters, the output of a single-frame stack stays within the polish bound of
    that frame.  No parameter setting can emit an image that is not built from the
    inputs -- which is exactly what the collapsed run did."""
    torch.manual_seed(3)
    wild = model.build(base=8, depth=2)
    with torch.no_grad():
        for parameter in wild.parameters():
            parameter.normal_(0.0, 2.0)
    frames = torch.full((1, 1, BANDS, 48, 48), 0.5)
    with torch.no_grad():
        out = wild(frames, conditioning(["available"], dts=[240.0]))
    bound = 0.5 * float(np.exp(model.POLISH_RANGE))
    lower = 0.5 * float(np.exp(-model.POLISH_RANGE))
    assert float(out.max()) <= bound + 1e-4
    assert float(out.min()) >= lower - 1e-4


def test_polish_gain_saturates_at_its_bound():
    torch.manual_seed(0)
    fresh = model.build(base=8, depth=2).eval()
    with torch.no_grad():
        fresh.polish.bias.fill_(1000.0)
    frames = torch.ones(1, 1, BANDS, 48, 48)
    out = fresh(frames, conditioning(["available"], dts=[240.0]))
    assert float(out.mean()) == pytest.approx(np.exp(model.POLISH_RANGE), rel=1e-3)


# --------------------------------------------------------------------- agreement


def test_agreement_flags_the_frame_that_disagrees():
    """A frozen or torn frame announces itself as deviation from the stack median --
    the deterministic replacement for both cross-attention and detector scores."""
    torch.manual_seed(0)
    fresh = model.build(base=8, depth=2)
    frames = torch.zeros(1, 4, BANDS, 32, 32)
    frames[:, 3] = 5.0                                  # the odd one out
    condition = conditioning(["available"] * 4)
    valid = condition[..., model.AVAILABLE]
    got = fresh._with_agreement(frames.reshape(4, BANDS, 32, 32), condition, valid,
                                batch=1, stack=4)
    assert float(got[0, 3, model.AGREE_MEAN]) > float(got[0, 0, model.AGREE_MEAN]) + 0.5
    assert float(got[0, 0, model.AGREE_MEAN]) < 0.1


# ------------------------------------------------------------------------- flow


def test_flow_is_zero_for_a_simultaneous_frame():
    """dt=0 means a real observation of the target instant: the flow allowance is
    zero there, however loud the head, so the most trusted candidate stays put."""
    torch.manual_seed(0)
    fresh = model.build(base=8, depth=2).eval()
    with torch.no_grad():
        fresh.flow.bias.fill_(1000.0)
    frames = stack(1, size=48)
    out = fresh(frames, conditioning(["available"], dts=[0.0]))
    assert torch.allclose(out, frames[:, 0], atol=1e-4)


def test_flow_displaces_by_the_dt_scaled_bound_at_saturation():
    """A saturated flow head moves a ramp by exactly the allowance for its dt --
    and no further.  Displacement is a fraction of frame width, so the same model
    moves features the same physical distance at 640 and 1280."""
    torch.manual_seed(0)
    fresh = model.build(base=8, depth=2).eval()
    with torch.no_grad():
        fresh.flow.bias.fill_(1000.0)   # tanh saturates: full allowance, +x and +y
    size = 64
    ramp = torch.linspace(-1.0, 1.0, size).view(1, 1, 1, size)
    frames = ramp.expand(1, 1, BANDS, size, size).clone()
    condition = conditioning(["available"], dts=[240.0])
    out = fresh(frames, condition)
    allowance = model.MAX_FLOW_FRACTION * float(condition[0, 0, model.DT_LOG])
    interior = out[..., 8:-8, 8:-8] - frames[:, 0, :, 8:-8, 8:-8]
    assert float(interior.mean()) == pytest.approx(2.0 * allowance, rel=0.05)


def test_untrained_flow_is_the_identity():
    frames = stack(3, size=32)
    warped = model._displace(frames[0], torch.zeros(3, 2, 32, 32))
    assert torch.allclose(warped, frames[0], atol=1e-5)


# ----------------------------------------------------------------------- trends


def test_dt_encoding_round_trips_exactly():
    """The trend builder recovers real seconds from the conditioning encoding."""
    dts = [-72000.0, -960.0, -240.0, 0.0, 240.0, 61440.0]
    condition = conditioning(["available"] * len(dts), dts=dts)
    recovered = model.dt_from_conditioning(condition)[0]
    assert torch.allclose(recovered, torch.tensor(dts), rtol=1e-5, atol=1e-3)


def trend_setup(dts, values, states=None):
    """A stack of uniform frames at given dts/values, expanded through the builder."""
    count = len(dts)
    frames = torch.zeros(1, count, BANDS, 16, 16)
    for i, value in enumerate(values):
        frames[:, i] = value
    condition = conditioning(states or ["available"] * count, dts=dts)
    net = model.build(base=8, depth=2)
    return expand(net, frames, condition)[:2]


def test_trend_candidates_recover_a_linear_ramp():
    """Brightness rising linearly into the gap: the extrapolation must land exactly
    on the line's value at the target instant -- the reach a convex blend of the
    observations alone cannot have."""
    expanded, condition = trend_setup(
        dts=[-480.0, -240.0, 240.0, 480.0], values=[0.2, 0.3, 0.5, 0.6])
    past, future = expanded[0, -2], expanded[0, -1]
    assert torch.allclose(past, torch.full_like(past, 0.4), atol=1e-6)
    assert torch.allclose(future, torch.full_like(future, 0.4), atol=1e-6)
    assert bool(condition[0, -2, model.DERIVED]) and bool(condition[0, -1, model.DERIVED])
    assert float(condition[0, -2, model.AVAILABLE]) == 1.0    # available, not missing


def test_trend_reach_is_capped():
    """A short-baseline slope far from the target must not extrapolate freely: at
    6x its baseline the line would say 0.95, but reach caps at 4x -> 0.75."""
    expanded, _ = trend_setup(dts=[-1680.0, -1440.0], values=[0.25, 0.35])
    past = expanded[0, -2]
    assert torch.allclose(past, torch.full_like(past, 0.75), atol=1e-6)


def test_trend_output_is_clamped_to_model_range():
    expanded, _ = trend_setup(dts=[-480.0, -240.0], values=[0.2, 0.9])
    assert float(expanded[0, -2].max()) <= 1.0


def test_trends_missing_when_a_side_lacks_two_frames():
    expanded, condition = trend_setup(dts=[-240.0, 240.0, 480.0], values=[0.1, 0.2, 0.3])
    assert float(condition[0, -2, model.MISSING]) == 1.0    # only one past frame
    assert float(condition[0, -1, model.AVAILABLE]) == 1.0  # future trend exists
    assert float(expanded[0, -2].abs().max()) == 0.0


def test_trends_ignore_invalid_and_cross_satellite_frames():
    """Extrapolation slopes must come from clean target-satellite data only: a
    missing frame has no pixels, and a counterpart's residual gain error would be
    amplified by the lever arm."""
    frames = torch.full((1, 4, BANDS, 16, 16), 0.5)
    condition = conditioning(["missing", "available", "available", "available"],
                             satellites=[True, False, True, True],
                             dts=[-480.0, -240.0, 240.0, 480.0])
    net = model.build(base=8, depth=2)
    _, expanded, _ = net._trends(frames, condition)
    assert float(expanded[0, -2, model.MISSING]) == 1.0    # one missing, one cross-sat
    assert float(expanded[0, -1, model.AVAILABLE]) == 1.0  # two clean same-sat frames


def test_agreement_is_zero_for_a_unanimous_stack():
    torch.manual_seed(0)
    fresh = model.build(base=8, depth=2)
    frames = torch.full((1, 3, BANDS, 32, 32), 0.25)
    condition = conditioning(["available"] * 3)
    valid = condition[..., model.AVAILABLE]
    got = fresh._with_agreement(frames.reshape(3, BANDS, 32, 32), condition, valid,
                                batch=1, stack=3)
    assert float(got[..., model.AGREE_MEAN:model.AGREE_MAX + 1].abs().max()) < 1e-6


def test_agreement_ignores_missing_frames():
    """A zeroed missing frame must not drag the median toward black."""
    torch.manual_seed(0)
    fresh = model.build(base=8, depth=2)
    frames = torch.full((1, 3, BANDS, 32, 32), 0.5)
    frames[:, 0] = 0.0                                  # missing, zero-filled
    condition = conditioning(["missing", "available", "available"])
    valid = condition[..., model.AVAILABLE]
    got = fresh._with_agreement(frames.reshape(3, BANDS, 32, 32), condition, valid,
                                batch=1, stack=3)
    assert float(got[0, 1, model.AGREE_MEAN]) < 1e-6
    assert float(got[0, 2, model.AGREE_MEAN]) < 1e-6


# ------------------------------------------------------------- per-band presence


def expected_band_blend(sources, condition, present):
    """The untrained model's output when some frames lack some bands: per band,
    the prior blend over the valid frames that hold that band."""
    valid = condition[..., model.AVAILABLE]
    logits = model.blend_prior(condition)[..., None].expand(-1, -1, BANDS)
    logits = logits.masked_fill((valid[..., None] < 0.5) | (present < 0.5),
                                float("-inf"))
    weights = torch.nan_to_num(torch.softmax(logits, dim=1), nan=0.0)
    return (sources * weights[..., None, None]).sum(dim=1)


def partial_setup():
    """Three frames, none a trend source pair, the first missing bands 0 and 4."""
    frames = torch.stack([torch.full((BANDS, 32, 32), v) for v in (0.2, 0.4, 0.6)])[None]
    condition = conditioning(["available"] * 3, satellites=[False, True, False],
                             dts=[0.0, -240.0, 480.0])
    present = torch.ones(1, 3, BANDS)
    present[0, 0, [0, 4]] = 0.0
    return frames, condition, present


def test_all_bands_present_is_the_default(net):
    """`present` of all ones is exactly what None means, bit for bit."""
    frames = stack(5)
    condition = conditioning(["available", "available", "available", "missing",
                              "available"])
    assert torch.equal(net(frames, condition),
                       net(frames, condition, torch.ones(1, 5, BANDS)))


def test_a_frame_missing_a_band_still_contributes_its_other_bands(net):
    frames, condition, present = partial_setup()
    out = net(frames, condition, present)
    assert torch.allclose(out, expected_band_blend(frames, condition, present),
                          atol=1e-4)
    without_first = expected_band_blend(frames, condition,
                                        torch.tensor([[[0.0] * BANDS, [1.0] * BANDS,
                                                       [1.0] * BANDS]]))
    # Where the first frame holds the band, it pulls the blend toward its 0.2...
    assert float(out[0, 1].mean()) < float(without_first[0, 1].mean()) - 0.01
    # ...and where it does not, the band is the blend of the other two alone.
    assert torch.allclose(out[0, 0], without_first[0, 0], atol=1e-4)
    assert torch.allclose(out[0, 4], without_first[0, 4], atol=1e-4)


def test_a_missing_bands_pixels_never_reach_the_output(net):
    frames, condition, present = partial_setup()
    reference = net(frames, condition, present)
    frames[0, 0, [0, 4]] = 999.0            # whatever the stack carries there
    out = net(frames, condition, present)
    assert torch.allclose(out, reference, atol=1e-4)


def test_a_band_no_frame_holds_comes_back_zero_not_nan(net):
    frames, condition, present = partial_setup()
    present[0, :, 2] = 0.0
    out = net(frames, condition, present)
    assert torch.isfinite(out).all()
    assert float(out[0, 2].abs().max()) == 0.0
    assert float(out[0, 3].abs().max()) > 0.1


def test_trends_come_only_from_frames_holding_every_band():
    """A frame missing a band has no line to continue in it."""
    frames = torch.full((1, 4, BANDS, 16, 16), 0.5)
    condition = conditioning(["available"] * 4, dts=[-480.0, -240.0, 240.0, 480.0])
    present = torch.ones(1, 4, BANDS)
    present[0, 1, 3] = 0.0                      # one of the past pair is partial
    net = model.build(base=8, depth=2)
    _, expanded, held = net._trends(frames, condition, present)
    assert float(expanded[0, -2, model.MISSING]) == 1.0     # past trend: missing
    assert float(expanded[0, -1, model.AVAILABLE]) == 1.0   # future trend: built
    assert held.shape == (1, 6, BANDS) and bool((held[0, 4:] == 1.0).all())


def test_agreement_ignores_the_bands_a_frame_lacks():
    """A partial frame is judged on the bands it holds, not on whatever stands in
    for the rest."""
    torch.manual_seed(0)
    fresh = model.build(base=8, depth=2)
    frames = torch.full((1, 3, BANDS, 32, 32), 0.5)
    frames[0, 0, 0] = 0.9                               # absent band, stand-in
    condition = conditioning(["available"] * 3)
    valid = condition[..., model.AVAILABLE]
    present = torch.ones(1, 3, BANDS)
    present[0, 0, 0] = 0.0
    got = fresh._with_agreement(frames.reshape(3, BANDS, 32, 32), condition, valid,
                                batch=1, stack=3, present=present)
    assert float(got[0, :, model.AGREE_MEAN:model.AGREE_MAX + 1].abs().max()) < 1e-6


def test_the_median_skips_frames_without_the_band():
    torch.manual_seed(0)
    fresh = model.build(base=8, depth=2)
    frames = torch.full((1, 3, BANDS, 8, 8), 0.5)
    frames[0, 0, 2] = -1.0
    valid = torch.ones(1, 3)
    present = torch.ones(1, 3, BANDS)
    present[0, 0, 2] = 0.0
    median = fresh._stack_median(frames.reshape(3, BANDS, 8, 8), valid, 1, 3, present)
    assert float(median[0, 0, 2].min()) == pytest.approx(0.5)


# ------------------------------------------------------------------------ training


def test_gradients_reach_every_head():
    """A head with no gradient path never learns, and zero-init hides it.  Checked
    after three optimiser steps, because the zero initialisation is nested two deep:
    at step 1 only the output heads move, at step 2 the FiLM output layers behind
    them, and only then does gradient reach everything.  Asserting earlier would
    fail on a model that is perfectly healthy."""
    torch.manual_seed(0)
    fresh = model.build(base=8, depth=2)
    frames = stack(3, size=48)
    condition = conditioning(["available", "missing", "available"])
    optimiser = torch.optim.Adam(fresh.parameters(), lr=1e-3)

    for _ in range(3):
        optimiser.zero_grad()
        out = fresh(frames, condition)
        out.square().mean().backward()
        optimiser.step()

    for name, parameter in fresh.named_parameters():
        assert parameter.grad is not None, f"{name} has no gradient"
        assert parameter.grad.abs().sum() > 0, f"{name} gradient is identically zero"


def smooth_stack(count, size=64, seed=1):
    """Spatially correlated frames; white noise tests nothing the real task needs."""
    generator = torch.Generator().manual_seed(seed)
    coarse = torch.randn(1, count * BANDS, size // 8, size // 8, generator=generator)
    smooth = torch.nn.functional.interpolate(coarse, size=(size, size), mode="bicubic",
                                             align_corners=False)
    return smooth.view(1, count, BANDS, size, size)


def train_a_little(net, frames, condition, target, steps=120, lr=3e-3):
    optimiser = torch.optim.Adam(net.parameters(), lr=lr)
    first = None
    for step in range(steps):
        optimiser.zero_grad()
        out = net(frames, condition)
        loss = (out - target).abs().mean()
        loss.backward()
        optimiser.step()
        if step == 0:
            first = loss.item()
    return first, loss.item()


def test_model_learns_to_select_one_frame_from_the_stack():
    """The weight head's core job: when one candidate is the answer, put the weight
    there.  Fails in seconds if the softmax, masking or composition is miswired."""
    torch.manual_seed(0)
    fresh = model.build(base=16, depth=2)
    frames = smooth_stack(3)
    first, last = train_a_little(
        fresh, frames, conditioning(["available"] * 3), frames[:, 1].clone()
    )
    assert last < first * 0.2, f"loss barely moved: {first:.4f} -> {last:.4f}"


def test_model_learns_a_per_band_gain_within_the_polish_bound():
    """The polish head's core job: the residual drift the daily fit leaves behind."""
    torch.manual_seed(0)
    fresh = model.build(base=16, depth=2)
    frames = smooth_stack(1)
    per_band = torch.tensor([1.2, 0.85, 1.1, 0.9, 1.25, 0.8]).view(1, BANDS, 1, 1)
    target = frames[:, 0] * per_band
    first, last = train_a_little(fresh, frames, conditioning(["available"],
                                                             dts=[240.0]), target)
    assert last < first * 0.3, f"loss barely moved: {first:.4f} -> {last:.4f}"


# ---------------------------------------------------------------- variant arms


@pytest.mark.parametrize("flags", [
    {"relative": True}, {"corr_flow": True}, {"attention": True},
    {"relative": True, "corr_flow": True, "attention": True},
])
def test_every_variant_is_the_prior_blend_untrained(flags):
    """Each screened cross-frame pathway must start silent: zero-init heads and
    gates mean an untrained variant is exactly the same sane policy as the
    baseline, and the v1-v3 collapse mode stays structurally unreachable."""
    torch.manual_seed(0)
    net = model.build(base=8, depth=2, **flags).eval()
    frames = stack(5)
    condition = conditioning(["available"] * 5)
    out = net(frames, condition)
    expanded_frames, expanded_condition, _ = expand(net, frames, condition)
    assert torch.allclose(out, expected_blend(expanded_frames, expanded_condition),
                          atol=1e-4)


@pytest.mark.parametrize("flags", [
    {"relative": True}, {"corr_flow": True}, {"attention": True},
])
def test_variant_forward_and_backward_are_finite(flags):
    torch.manual_seed(1)
    net = model.build(base=8, depth=2, **flags)
    frames = stack(6, size=32)
    condition = conditioning(["available", "available", "available", "missing",
                              "available", "available"])
    out = net(frames, condition)
    assert torch.isfinite(out).all()
    out.square().mean().backward()
    for name, parameter in net.named_parameters():
        if parameter.grad is not None:
            assert torch.isfinite(parameter.grad).all(), name


def test_variant_checkpoints_round_trip_through_load_learned(tmp_path, monkeypatch):
    """A variant's flags ride in the checkpoint; loading must rebuild the same
    shapes or state_dict loading fails loudly."""
    import torch as torch_module

    from suvi import fillers

    net = model.build(base=8, depth=2, relative=True, corr_flow=True)
    path = str(tmp_path / "variant.pt")
    torch_module.save({"model": net.state_dict(),
                       "args": {"base": 8, "depth": 2,
                                "cond_layout": list(model.COND_LAYOUT),
                                "relative": True, "corr_flow": True,
                                "offsets": [1, 4, 16, 64]}}, path)
    monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, path)
    fillers._LEARNED.clear()
    loaded, _, _ = fillers.load_learned()
    assert loaded.relative and loaded.corr_flow and loaded.attend is None
    assert loaded.offsets == (1, 4, 16, 64)
    fillers._LEARNED.clear()


@pytest.mark.parametrize("args", [
    {"base": 8, "depth": 2, "cond_dim": 9},                   # v8: a suspect column
    {"base": 8, "depth": 2, "cond_layout": ["available", "missing", "suspect"]},
])
def test_a_checkpoint_under_another_conditioning_layout_is_refused(
        tmp_path, monkeypatch, args):
    """Its weights would read every column shifted: refuse by name, not by a
    shape error -- or worse, a v5 checkpoint whose shapes happen to fit."""
    import torch as torch_module

    from suvi import fillers

    path = str(tmp_path / "old.pt")
    torch_module.save({"model": model.build(base=8, depth=2).state_dict(),
                       "args": args}, path)
    monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, path)
    fillers._LEARNED.clear()
    with pytest.raises(SystemExit, match="conditioning layout"):
        fillers.load_learned()
    fillers._LEARNED.clear()


def test_correlation_peaks_at_the_true_shift():
    """The cost volume must actually measure correspondence: a frame shifted by
    +2 px correlates best with the median at the matching displacement."""
    torch.manual_seed(0)
    base_img = torch.randn(1, 6, 32, 32)
    shifted = torch.roll(base_img, shifts=2, dims=-1)
    volume = model._correlation(shifted, base_img, radius=3)
    centre = volume.shape[1] // 2
    interior = volume[..., 8:-8, 8:-8].mean(dim=(-2, -1))[0]
    assert int(interior.argmax()) != centre
    best_dx = int(interior.argmax()) % 7 - 3
    assert best_dx == -2 or best_dx == 2


# ------------------------------------------------- inference path memory refactor


def _trained_looking(attention):
    """A model whose heads are non-zero, so the blend, gain and flow all act."""
    torch.manual_seed(1)
    built = model.build(base=8, depth=2, attention=attention)
    with torch.no_grad():
        for head in (built.weight, built.polish, built.flow):
            head.weight.normal_(0.0, 0.2)
            head.bias.normal_(0.0, 0.2)
        if built.attn_gate is not None:
            built.attn_gate.fill_(0.5)
    return built.eval()


@pytest.mark.parametrize("attention", [True, False])
def test_band_at_a_time_inference_matches_the_batched_path_exactly(
        attention, monkeypatch):
    """Inference blends one band at a time to bound memory (a 24-candidate stack
    at 1280px peaked at 9.4 GB all at once); it must not change one bit.  Both
    sides run without gradients: attention itself takes a different fused kernel
    when gradients are off, so a grad-vs-no-grad comparison would test torch."""
    net_ = _trained_looking(attention)
    sources = stack(6, size=96, seed=3).clamp(-1, 1)
    cond = conditioning(["available", "available", "available", "missing",
                         "available", "available"],
                        satellites=[True, True, True, True, False, True],
                        dts=[-240.0, 240.0, -960.0, 960.0, 0.0, -3840.0])
    with torch.no_grad():
        monkeypatch.setattr(model, "BAND_AT_A_TIME", False)
        batched = net_(sources, cond)
        monkeypatch.setattr(model, "BAND_AT_A_TIME", True)
        banded = net_(sources, cond)
    assert banded.shape == batched.shape == (1, BANDS, 96, 96)
    assert torch.equal(banded, batched)


def test_inference_matches_the_gradient_path(monkeypatch):
    """Without gradients the encoder and the blend run a few frames at a time and the
    softmax is streamed; the same computation, so equal up to float rounding."""
    net_ = _trained_looking(False)
    sources = stack(6, size=96, seed=4).clamp(-1, 1)
    cond = conditioning(["available", "available", "available", "missing",
                         "available", "available"],
                        satellites=[True, True, True, True, False, True],
                        dts=[-240.0, 240.0, -960.0, 960.0, 0.0, -3840.0])
    present = torch.ones(1, 6, BANDS)
    present[0, 4, 1] = 0.0
    monkeypatch.setattr(model, "BAND_AT_A_TIME", False)
    with torch.enable_grad():
        reference = net_(sources, cond, present).detach()
    with torch.no_grad():
        assert torch.allclose(net_(sources, cond, present), reference,
                              rtol=0, atol=5e-5)


@pytest.mark.parametrize("band_at_a_time", [True, False])
def test_chunked_inference_handles_batches_ragged_chunks_and_unheld_bands(
        monkeypatch, band_at_a_time):
    """Two stacks of seven frames -- not a multiple of the chunk -- with one band no
    frame of the second stack holds: that band blends to zero, not NaN."""
    monkeypatch.setattr(model, "BAND_AT_A_TIME", band_at_a_time)
    net_ = _trained_looking(False)
    sources = torch.cat([stack(7, size=48, seed=6), stack(7, size=48, seed=7)]).clamp(-1, 1)
    states = ["available", "available", "missing", "available", "available",
              "available", "available"]
    dts = [-240.0, 240.0, -960.0, 960.0, 0.0, -3840.0, 3840.0]
    cond = torch.cat([conditioning(states, dts=dts),
                      conditioning(states, satellites=[True] * 6 + [False], dts=dts)])
    present = torch.ones(2, 7, BANDS)
    present[1, :, 3] = 0.0
    with torch.enable_grad():
        reference = net_(sources, cond, present).detach()
    with torch.no_grad():
        out = net_(sources, cond, present)
    assert torch.isfinite(out).all()
    assert float(out[1, 3].abs().max()) == 0.0
    assert torch.allclose(out, reference, rtol=0, atol=5e-5)


def test_trends_are_returned_apart_from_the_stack():
    frames = torch.full((1, 4, BANDS, 16, 16), 0.5)
    condition = conditioning(["available"] * 4, dts=[-480.0, -240.0, 240.0, 480.0])
    trends, expanded, present = model.build(base=8, depth=2)._trends(frames, condition)
    assert trends.shape == (1, 2, BANDS, 16, 16)
    assert expanded.shape[1] == present.shape[1] == 6


def test_a_one_frame_stack_has_no_trends_and_still_blends(net):
    frames = stack(1, size=48)
    trends, condition, _ = net._trends(frames, conditioning(["available"]))
    assert trends.shape[1] == 0 and condition.shape[1] == 1
    out = net(frames, conditioning(["available"], dts=[240.0]))
    assert torch.isfinite(out).all()


def test_training_keeps_the_batched_path_and_its_gradients():
    net_ = _trained_looking(False).train()
    sources = stack(4, size=48, seed=5).clamp(-1, 1).requires_grad_(True)
    cond = conditioning(["available"] * 4)
    net_(sources, cond).square().mean().backward()
    assert sources.grad is not None and float(sources.grad.abs().sum()) > 0
    assert net_.weight.weight.grad is not None