"""The learned filler's wiring: the stack it is handed, and how it is loaded.

The model itself is covered in test_model.py.  What matters here is that `bench.py`
hands it the same stack shape `suvi.samples` builds during training, that a band judged
bad is withheld from it, and that damage the judgment missed arrives as *corrupted
pixels* rather than as the pristine original.
"""

import numpy as np
import pytest

import bench
from conftest import solar_disc
from suvi import cases, fillers, paths, samples


WAVELENGTHS = paths.WAVELENGTHS


def overlay_for(satellites=(16, 18), count=40, deleted=()):
    """An Overlay whose slots resolve through in-memory lookups rather than files."""
    base = 1715400000 // paths.CADENCE * paths.CADENCE
    times = [base + i * paths.CADENCE for i in range(count)]
    archive = {(s, w, t): f"truth/{s}/{w}/{t}"
               for s in satellites for w in WAVELENGTHS for t in times}
    return cases.Overlay(archive=archive, deleted=frozenset(deleted)), times, base


def reader(size=16, missing=()):
    """A six-band reader.  `missing` names (satellite, time) pairs with no frame."""
    frame = np.stack([solar_disc(size=size, radius=size // 3, peak=1.0 + i)
                      for i in range(6)]).astype(np.float32)

    def read(satellite, when):
        if (satellite, when) in missing:
            return None
        return frame.copy()

    return read, frame


def judged_reader(monkeypatch, overlay, bad, size=16):
    """The real `bench._band_reader` over `overlay_for`'s in-memory archive: every
    path decodes to the matching band of one six-band frame."""
    _, frame = reader(size)
    index = {wavelength: i for i, wavelength in enumerate(WAVELENGTHS)}

    def read_image(path):
        wavelength = int(path.split("/")[2])
        return frame[index[wavelength]].copy(), None

    monkeypatch.setattr(bench.fitsio, "read_image", read_image)
    return bench._band_reader(overlay, WAVELENGTHS, set(bad)), frame


def build(overlay, target, read, bad=frozenset(), satellites=(16, 18)):
    """`bench._build_stack` for GOES-16, with anchors from the judgment `read` uses."""
    band_times = bench._band_clean_times(overlay, satellites, WAVELENGTHS, bad)
    return bench._build_stack(overlay, 16, target, satellites, WAVELENGTHS, read,
                              band_times)


# --------------------------------------------------------------------------- stack


def test_stack_matches_the_training_layout():
    """Trained on one stack shape, evaluated on another, would be a silent mismatch.

    The window spans 560 slots so even the +/-256 rungs of the exponential ladder
    exist; a real window edge simply omits the rungs it cannot reach.
    """
    overlay, times, _ = overlay_for(count=560)
    read, _ = reader()
    stack = build(overlay, times[280], read)

    offsets = {(entry["same_satellite"], entry["dt"] / paths.CADENCE) for entry in stack}
    for source, offset in samples.stack_layout((16, 18), 16):
        assert (source == 16, float(offset)) in offsets, f"missing {(source, offset)}"


def test_stack_carries_all_six_bands():
    """The model is joint across bands; a per-band stack would not fit it."""
    overlay, times, _ = overlay_for()
    read, _ = reader()
    stack = build(overlay, times[20], read)
    for entry in stack:
        assert entry["image"].shape[0] == len(WAVELENGTHS)


def test_stack_holds_the_targets_undamaged_bands_at_dt_zero(monkeypatch):
    """As in training: the renderer fills only the bands a slot lacks, so the bands
    it kept are the model's best evidence -- and the lost ones are withheld."""
    overlay, times, _ = overlay_for()
    target = times[20]
    bad = {(16, 171, target), (16, 304, target)}
    read, frame = judged_reader(monkeypatch, overlay, bad)
    stack = build(overlay, target, read, bad)

    own = [e for e in stack if e["same_satellite"] and e["dt"] == 0.0]
    assert len(own) == 1 and own[0]["slot"] == (16, target)
    assert own[0]["state"] == "available"
    assert own[0]["bands"] == (True, True, False, True, True, False)
    assert np.isnan(own[0]["image"][[2, 5]]).all()
    np.testing.assert_allclose(own[0]["image"][[0, 1, 3, 4]], frame[[0, 1, 3, 4]])


def test_a_target_with_no_usable_band_is_a_missing_entry(monkeypatch):
    overlay, times, _ = overlay_for()
    target = times[20]
    bad = {(16, w, target) for w in WAVELENGTHS}
    read, _ = judged_reader(monkeypatch, overlay, bad)
    stack = build(overlay, target, read, bad)
    own = next(e for e in stack if e["same_satellite"] and e["dt"] == 0.0)
    assert own["state"] == "missing" and own["image"] is None


def test_a_target_absent_from_the_archive_is_no_entry():
    overlay, times, _ = overlay_for()
    for wavelength in WAVELENGTHS:
        del overlay.archive[(16, wavelength, times[20])]
    read, _ = reader()
    stack = build(overlay, times[20], read)
    assert not any(e["same_satellite"] and e["dt"] == 0.0 for e in stack)


def test_stack_includes_the_counterpart_at_the_target_instant():
    overlay, times, _ = overlay_for()
    read, _ = reader()
    stack = build(overlay, times[20], read)
    assert any(entry["dt"] == 0.0 and not entry["same_satellite"] for entry in stack)


def test_judged_bad_neighbours_are_withheld(monkeypatch):
    """2024-04-17: shown to v8, N_LONG=0 static still got a small blend weight and
    speckled every fill's far field.  A judged-bad frame is a hole."""
    overlay, times, _ = overlay_for()
    target = times[20]
    neighbour_time = target - paths.CADENCE
    bad = {(16, w, neighbour_time) for w in WAVELENGTHS}
    read, _ = judged_reader(monkeypatch, overlay, bad)
    stack = build(overlay, target, read, bad)

    entry = next(e for e in stack if e["slot"] == (16, neighbour_time))
    assert entry["state"] == "missing"
    assert entry["image"] is None
    assert {e["state"] for e in stack} <= {"available", "missing"}


def test_a_neighbour_bad_in_one_band_keeps_its_other_bands_available(monkeypatch):
    """Per band: a fault in 195 withholds 195, not the frame's five clean bands."""
    overlay, times, _ = overlay_for()
    target = times[20]
    neighbour_time = target - paths.CADENCE
    bad = {(16, 195, neighbour_time)}
    read, frame = judged_reader(monkeypatch, overlay, bad)
    stack = build(overlay, target, read, bad)

    entry = next(e for e in stack if e["slot"] == (16, neighbour_time))
    assert entry["state"] == "available"
    assert entry["bands"] == (True, True, True, False, True, True)
    assert np.isnan(entry["image"][3]).all()
    np.testing.assert_allclose(entry["image"][[0, 1, 2, 4, 5]], frame[[0, 1, 2, 4, 5]])


def test_unreadable_neighbours_arrive_missing_with_no_pixels():
    overlay, times, _ = overlay_for()
    target = times[20]
    gone = (16, target - paths.CADENCE)
    read, _ = reader(missing={gone})
    stack = build(overlay, target, read)
    entry = next(e for e in stack if e["slot"] == gone)
    assert entry["state"] == "missing"
    assert entry["image"] is None


def test_undamaged_neighbours_arrive_available():
    overlay, times, _ = overlay_for()
    read, frame = reader()
    stack = build(overlay, times[20], read)
    assert {e["state"] for e in stack} == {"available"}
    for entry in stack:
        np.testing.assert_allclose(entry["image"], frame)


def test_stack_adds_anchors_beyond_the_fixed_offsets(monkeypatch):
    """A long outage must still reach a real frame."""
    overlay, times, _ = overlay_for(count=80)
    target = times[40]
    bad = {(16, w, target + n * paths.CADENCE)
           for w in WAVELENGTHS for n in range(-30, 31)}
    read, _ = judged_reader(monkeypatch, overlay, bad)
    stack = build(overlay, target, read, bad)

    reach = [abs(e["dt"]) for e in stack
             if e["same_satellite"] and e["state"] == "available"]
    assert reach, "no usable same-satellite frame found at all"
    assert max(reach) >= 31 * paths.CADENCE


def test_stack_skips_slots_the_window_does_not_contain():
    """Near a window edge the offsets run off the end; those must not appear as data."""
    overlay, times, _ = overlay_for(count=40)
    read, _ = reader()
    stack = build(overlay, times[0], read)
    for entry in stack:
        satellite, when = entry["slot"]
        assert any((satellite, w, when) in overlay.archive for w in WAVELENGTHS)


# ------------------------------------------------------------------- band reader


def test_band_reader_returns_damaged_pixels_the_judgment_missed(tmp_path):
    """A filler handed the pristine frame for a corrupted slot reads the answer key."""
    from conftest import write_fits

    truth_frame = solar_disc(size=64, radius=20, peak=1.0)
    damaged_frame = solar_disc(size=64, radius=20, peak=9.0)
    when = 1715400000
    archive_map, overrides = {}, {}
    for wavelength in WAVELENGTHS:
        good = str(tmp_path / f"t{wavelength}.fits")
        harmed = str(tmp_path / f"d{wavelength}.fits")
        write_fits(good, truth_frame)
        write_fits(harmed, damaged_frame)
        archive_map[(16, wavelength, when)] = good
        overrides[(16, wavelength, when)] = harmed

    overlay = cases.Overlay(archive=archive_map, overrides=overrides)
    read = bench._band_reader(overlay, WAVELENGTHS, set())
    got = read(16, when)
    assert got.shape == (6, 64, 64)
    # Compared loosely: the FITS tile compression is lossy at the 1e-3 level, and what
    # is under test is which *file* was opened, not the codec.
    assert got[0].max() == pytest.approx(damaged_frame.max(), rel=1e-2)
    assert got[0].max() > truth_frame.max() * 5


def test_band_reader_caches_but_stays_bounded(tmp_path):
    """Six bands of 1280x1280 float32 is 39 MB; an unbounded cache exhausts the VM."""
    from conftest import write_fits

    frame = solar_disc(size=32, radius=10, peak=1.0)
    archive_map = {}
    times = [1715400000 + i * paths.CADENCE for i in range(10)]
    for when in times:
        for wavelength in WAVELENGTHS:
            path = str(tmp_path / f"{when}_{wavelength}.fits")
            write_fits(path, frame)
            archive_map[(16, wavelength, when)] = path

    overlay = cases.Overlay(archive=archive_map)
    reads = []
    real = bench.fitsio.read_image

    def counting(path):
        reads.append(path)
        return real(path)

    slot_bytes = 6 * 32 * 32 * 4
    read = bench._band_reader(overlay, WAVELENGTHS, set(), budget=3 * slot_bytes)
    original = bench.fitsio.read_image
    bench.fitsio.read_image = counting
    try:
        read(16, times[0])
        first = len(reads)
        read(16, times[0])
        assert len(reads) == first, "a repeat read was not served from cache"
        for when in times[1:5]:
            read(16, when)
        before_evicted = len(reads)
        read(16, times[0])                      # evicted by the 3-slot budget
        assert len(reads) > before_evicted, "cache grew past its budget"

        # Least recently used goes first: touching a slot keeps it.
        for when in times[5:8]:
            read(16, when)
        read(16, times[5])                      # refresh the oldest
        read(16, times[8])                      # evicts times[6], not times[5]
        before = len(reads)
        read(16, times[5])
        assert len(reads) == before, "a recently used slot was evicted"
        read(16, times[6])
        assert len(reads) > before, "the least recently used slot was kept"
    finally:
        bench.fitsio.read_image = original


def test_band_reader_builds_one_float32_array_per_slot(monkeypatch):
    overlay, times, _ = overlay_for()
    read, frame = judged_reader(monkeypatch, overlay, {(16, 171, times[3])})
    image = read(16, times[3])
    assert image.dtype == np.float32 and image.shape == (6, 16, 16)
    assert np.isnan(image[2]).all()
    np.testing.assert_array_equal(image[[0, 1, 3, 4, 5]], frame[[0, 1, 3, 4, 5]])


def _write_bands(tmp_path, wavelengths, when=1715400000, satellite=16):
    from conftest import write_fits

    frame = solar_disc(size=32, radius=10, peak=1.0)
    archive_map = {}
    for wavelength in wavelengths:
        path = str(tmp_path / f"{satellite}_{wavelength}.fits")
        write_fits(path, frame)
        archive_map[(satellite, wavelength, when)] = path
    return cases.Overlay(archive=archive_map), when


def test_band_reader_keeps_the_bands_a_slot_has(tmp_path):
    """Per band, not all or nothing: a slot without 304 still carries five real
    bands, and the missing one is NaN so nothing mistakes it for dark sky."""
    overlay, when = _write_bands(tmp_path, WAVELENGTHS[:-1])
    image = bench._band_reader(overlay, WAVELENGTHS, set())(16, when)
    assert image.shape == (6, 32, 32)
    assert np.isnan(image[5]).all()
    assert np.isfinite(image[:5]).all()
    assert bench.band_presence(image) == (True,) * 5 + (False,)
    assert not bench._complete(image)


def test_band_reader_withholds_bands_judged_bad(tmp_path):
    """A judged-bad band reads as absent, like a missing one -- its pixels never
    reach the model -- and a slot bad in every band reads as nothing at all."""
    overlay, when = _write_bands(tmp_path, WAVELENGTHS)
    image = bench._band_reader(overlay, WAVELENGTHS, {(16, 131, when)})(16, when)
    assert np.isnan(image[1]).all()
    assert bench.band_presence(image) == (True, False, True, True, True, True)
    other = bench._band_reader(overlay, WAVELENGTHS, {(18, 131, when)})(16, when)
    assert bench._complete(other)            # the other satellite's verdict
    every = {(16, wavelength, when) for wavelength in WAVELENGTHS}
    assert bench._band_reader(overlay, WAVELENGTHS, every)(16, when) is None


def test_band_reader_returns_none_only_when_no_band_reads(tmp_path):
    overlay, when = _write_bands(tmp_path, ())
    assert bench._band_reader(overlay, WAVELENGTHS, set())(16, when) is None
    overlay, when = _write_bands(tmp_path, WAVELENGTHS)
    assert bench._complete(bench._band_reader(overlay, WAVELENGTHS, set())(16, when))
    assert not bench._complete(None)


# ------------------------------------------------------------ per-band stacks


def partial_reader(size=16, absent=()):
    """A reader whose slots lack the bands named in `absent`:
    {(satellite, time): (band index, ...)}."""
    read, frame = reader(size)
    absent = dict(absent)

    def read_partial(satellite, when):
        image = read(satellite, when)
        for band in absent.get((satellite, when), ()):
            image[band] = np.nan
        return image

    return read_partial, frame


def test_entries_say_which_bands_they_hold():
    overlay, times, _ = overlay_for()
    read, _ = reader()
    stack = build(overlay, times[20], read)
    assert all(entry["bands"] == (True,) * 6 for entry in stack)


def test_a_counterpart_missing_a_band_still_joins_the_stack():
    """2024-03-03: GOES-18 had no 304 all day.  Its 94 and 131 were exactly what
    GOES-16 lacked, and dropping the whole frame left the model nothing."""
    overlay, times, _ = overlay_for()
    target = times[20]
    read, _ = partial_reader(absent={(18, target): (5,)})
    stack = build(overlay, target, read)
    counterpart = next(e for e in stack if e["slot"] == (18, target))
    assert counterpart["state"] == "available"
    assert counterpart["bands"] == (True,) * 5 + (False,)
    assert np.isfinite(counterpart["image"][:5]).all()


def test_band_clean_times_are_per_satellite_and_band():
    overlay, times, _ = overlay_for(count=5)
    bad = {(16, 94, times[1]), (18, 304, times[3])}
    band_times = bench._band_clean_times(overlay, (16, 18), WAVELENGTHS, bad)
    assert band_times[16][94] == [times[i] for i in (0, 2, 3, 4)]
    assert band_times[16][304] == times
    assert band_times[18][304] == [times[i] for i in (0, 1, 2, 4)]


def test_anchors_reach_each_band_separately(monkeypatch):
    """A band bad for a long stretch gets its own anchor at its nearest clean
    slot, though every other band is clean right next to the target."""
    overlay, times, _ = overlay_for(count=140)
    target = times[70]
    step = paths.CADENCE
    bad = {(16, 94, target + n * step) for n in range(-45, 46)}
    read, _ = judged_reader(monkeypatch, overlay, bad)
    stack = build(overlay, target, read, bad)
    anchors = {e["slot"][1] - target: e for e in stack
               if e["same_satellite"] and e["slot"][1] in (target - 46 * step,
                                                           target + 46 * step)}
    assert sorted(anchors) == [-46 * step, 46 * step]
    assert all(e["bands"][0] and e["state"] == "available" for e in anchors.values())


def test_an_anchor_leaves_out_its_bands_judged_bad(monkeypatch):
    overlay, times, _ = overlay_for(count=140)
    target = times[70]
    step = paths.CADENCE
    bad = ({(16, 94, target + n * step) for n in range(-45, 46)}
           | {(16, 304, target - 46 * step)})
    read, _ = judged_reader(monkeypatch, overlay, bad)
    stack = build(overlay, target, read, bad)
    anchor = next(e for e in stack if e["slot"] == (16, target - 46 * step))
    assert anchor["bands"] == (True,) * 5 + (False,)
    assert anchor["state"] == "available"


# -------------------------------------------------------------------------- filler


def test_learned_returns_nothing_without_a_stack():
    assert fillers.learned(fillers.FillContext()) is None


def test_learned_returns_nothing_without_a_checkpoint(monkeypatch):
    monkeypatch.delenv(fillers.LEARNED_CHECKPOINT_ENV, raising=False)
    fillers._LEARNED.clear()
    context = fillers.FillContext(stack=[
        {"image": np.zeros((6, 16, 16), np.float32), "state": "available",
         "dt": -240.0, "same_satellite": True}
    ])
    assert fillers.learned(context) is None


def test_learned_is_registered_alongside_the_others():
    assert "learned" in fillers.FILLERS
    assert fillers.FILLERS["learned"] is fillers.learned


def test_existing_fillers_ignore_the_stack():
    """Adding `stack` must not perturb any measured baseline."""
    before = solar_disc(size=64, radius=20, peak=1.0)
    after = solar_disc(size=64, radius=20, peak=1.2)
    extra = np.stack([after * 5] * 6)
    plain = fillers.FillContext(before=before, after=after, dt_before=240, dt_after=240)
    with_stack = fillers.FillContext(
        before=before, after=after, dt_before=240, dt_after=240,
        stack=[{"image": extra, "state": "available", "dt": -240.0,
                "same_satellite": True}],
    )
    for name in ("hold_last", "linear_blend", "optical_flow", "solar_rotation"):
        np.testing.assert_allclose(fillers.FILLERS[name](plain),
                                   fillers.FILLERS[name](with_stack))


def test_learned_runs_end_to_end_against_a_saved_checkpoint(tmp_path, monkeypatch):
    """Checkpoint -> load -> fill, at the archive's native frame size."""
    torch = pytest.importorskip("torch")
    from suvi import model

    net = model.build(base=8, depth=2)
    path = tmp_path / "model.pt"
    torch.save({"model": net.state_dict(), "args": {"base": 8, "depth": 2, "cond_layout": list(model.COND_LAYOUT)},
                "epoch": 0}, path)
    monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, str(path))
    fillers._LEARNED.clear()

    def six(peak):
        return np.stack([solar_disc(size=128, radius=40, peak=peak + b * 0.2)
                         for b in range(6)]).astype(np.float32)

    when = 1715400000
    context = fillers.FillContext(
        stack=[
            {"image": six(1.0), "state": "available", "dt": -240.0,
             "same_satellite": True, "slot": (16, when - 240)},
            {"image": six(1.1), "state": "available", "dt": 240.0,
             "same_satellite": True, "slot": (16, when + 240)},
            {"image": six(0.9), "state": "available", "dt": 0.0,
             "same_satellite": False, "slot": (18, when)},
            {"image": None, "state": "missing", "dt": -960.0,
             "same_satellite": True, "slot": (16, when - 960)},
        ],
        calibration=(six(0.9), six(1.05)),
    )
    filled = fillers.learned(context)
    assert filled.shape == (6, 128, 128)
    assert np.isfinite(filled).all()
    fillers._LEARNED.clear()


def _save_tiny_model(tmp_path, monkeypatch):
    torch = pytest.importorskip("torch")
    from suvi import model

    net = model.build(base=8, depth=2)
    path = tmp_path / "model.pt"
    torch.save({"model": net.state_dict(),
                "args": {"base": 8, "depth": 2, "cond_layout": list(model.COND_LAYOUT)}}, path)
    monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, str(path))
    fillers._LEARNED.clear()


def _six(peak, size=64, absent=()):
    image = np.stack([solar_disc(size=size, radius=size // 3, peak=peak + b * 0.2)
                      for b in range(6)]).astype(np.float32)
    image[list(absent)] = np.nan
    return image


def _entry(image, dt, same, slot):
    return {"image": image, "bands": bench.band_presence(image), "state": "available",
            "dt": dt, "same_satellite": same, "slot": slot}


def test_assemble_stack_reports_which_bands_each_entry_holds(tmp_path, monkeypatch):
    torch = pytest.importorskip("torch")
    when = 1715400000
    context = fillers.FillContext(stack=[
        _entry(_six(1.0, absent=(0, 1)), -240.0, True, (16, when - 240)),
        _entry(_six(1.0, absent=(5,)), 0.0, False, (18, when)),
    ])
    aligned, condition, present = fillers.assemble_stack(context, torch,
                                                         torch.device("cpu"))
    assert present.tolist() == [[[0, 0, 1, 1, 1, 1], [1, 1, 1, 1, 1, 0]]]
    assert torch.isfinite(aligned).all()


def test_learned_fills_from_bands_split_across_satellites(tmp_path, monkeypatch):
    """The 2024-03-03 case: the target satellite lacks 94/131, the counterpart
    lacks 304.  Every band is observed somewhere, so every band is filled."""
    _save_tiny_model(tmp_path, monkeypatch)
    when = 1715400000
    context = fillers.FillContext(stack=[
        _entry(_six(1.0, absent=(0, 1)), -240.0, True, (16, when - 240)),
        _entry(_six(1.1, absent=(0, 1)), 240.0, True, (16, when + 240)),
        _entry(_six(0.9, absent=(5,)), 0.0, False, (18, when)),
    ])
    filled = fillers.learned(context)
    assert filled.shape == (6, 64, 64)
    assert np.isfinite(filled).all()
    fillers._LEARNED.clear()


def test_learned_releases_the_raw_stack_before_the_forward_pass(tmp_path, monkeypatch):
    """The raw six-band images are a second copy of everything the model blends;
    held through the forward pass they were ~0.8 GB of a 1280px fill's peak."""
    import weakref

    _save_tiny_model(tmp_path, monkeypatch)
    when = 1715400000
    images = [_six(1.0), _six(1.1), _six(0.9)]
    refs = [weakref.ref(image) for image in images]
    context = fillers.FillContext(
        stack=[_entry(images[0], -240.0, True, (16, when - 240)),
               _entry(images[1], 240.0, True, (16, when + 240)),
               _entry(images[2], 0.0, False, (18, when))],
        calibration=(_six(0.9), _six(1.0)),
    )
    del images
    net = fillers.load_learned()[0]
    alive_at_forward = []
    real_forward = type(net).forward

    def forward(self, *args, **kwargs):
        alive_at_forward.append(sum(ref() is not None for ref in refs))
        return real_forward(self, *args, **kwargs)

    monkeypatch.setattr(type(net), "forward", forward)
    filled = fillers.learned(context)
    assert filled.shape == (6, 64, 64)
    assert alive_at_forward == [0]
    assert context.stack == [] and context.calibration is None
    fillers._LEARNED.clear()


def test_a_band_no_frame_holds_comes_back_nan(tmp_path, monkeypatch):
    """Not a black band: the caller must be able to tell 'unfillable' from 'dark'."""
    _save_tiny_model(tmp_path, monkeypatch)
    when = 1715400000
    context = fillers.FillContext(stack=[
        _entry(_six(1.0, absent=(2,)), -240.0, True, (16, when - 240)),
        _entry(_six(0.9, absent=(2,)), 0.0, False, (18, when)),
    ])
    filled = fillers.learned(context)
    assert np.isnan(filled[2]).all()
    assert np.isfinite(np.delete(filled, 2, axis=0)).all()
    fillers._LEARNED.clear()


def test_learned_declines_a_stack_with_no_pixels_anywhere(tmp_path, monkeypatch):
    """Silently emitting a black frame here is how a fabricated fill would enter the
    archive; the contract is None."""
    torch = pytest.importorskip("torch")
    from suvi import model

    net = model.build(base=8, depth=2)
    path = tmp_path / "model.pt"
    torch.save({"model": net.state_dict(), "args": {"base": 8, "depth": 2, "cond_layout": list(model.COND_LAYOUT)}}, path)
    monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, str(path))
    fillers._LEARNED.clear()
    context = fillers.FillContext(stack=[
        {"image": None, "state": "missing", "dt": -240.0, "same_satellite": True,
         "slot": (16, 1715400000 - 240)},
    ])
    assert fillers.learned(context) is None
    fillers._LEARNED.clear()


def test_learned_checkpoint_is_loaded_once(tmp_path, monkeypatch):
    """Thousands of slots per bench run; reloading 14M parameters each time would
    dominate the wall clock."""
    torch = pytest.importorskip("torch")
    from suvi import model

    net = model.build(base=8, depth=2)
    path = tmp_path / "model.pt"
    torch.save({"model": net.state_dict(), "args": {"base": 8, "depth": 2, "cond_layout": list(model.COND_LAYOUT)}}, path)
    monkeypatch.setenv(fillers.LEARNED_CHECKPOINT_ENV, str(path))
    fillers._LEARNED.clear()

    loads = []
    real_load = torch.load
    monkeypatch.setattr(torch, "load", lambda *a, **k: (loads.append(1), real_load(*a, **k))[1])
    for _ in range(3):
        fillers.load_learned()
    assert len(loads) == 1
    fillers._LEARNED.clear()


def test_band_reader_caches_a_repeated_target(tmp_path, monkeypatch):
    """One target pulls one frame per band through the reader, and consecutive targets
    share most of their stack.  Without the cache each shared frame is re-read, and on
    this archive a read is an SMB round trip through two filesystems."""
    from conftest import write_fits

    frame = solar_disc(size=32, radius=10, peak=1.0)
    when = 1715400000
    archive_map = {}
    for wavelength in WAVELENGTHS:
        path = str(tmp_path / f"{wavelength}.fits")
        write_fits(path, frame)
        archive_map[(16, wavelength, when)] = path

    reads = []
    real = bench.fitsio.read_image
    monkeypatch.setattr(bench.fitsio, "read_image",
                        lambda p, *a, **k: reads.append(p) or real(p, *a, **k))

    overlay = cases.Overlay(archive=archive_map)
    read = bench._band_reader(overlay, WAVELENGTHS, set())
    read(16, when)
    assert len(reads) == len(WAVELENGTHS), "one read per band"
    read(16, when)
    assert len(reads) == len(WAVELENGTHS), "re-read a slot already in the cache"