"""ceiling.py's reference blends, which are per band: a frame may lack some bands."""

import numpy as np
import pytest

torch = pytest.importorskip("torch")

import ceiling  # noqa: E402
from suvi import model, samples  # noqa: E402


def conditioning(dts, states=None):
    states = states or ["available"] * len(dts)
    return torch.stack([model.frame_conditioning(state, True, dt)
                        for state, dt in zip(states, dts)])[None]


def test_prior_weights_normalise_per_band_over_the_frames_holding_it():
    cond = conditioning([-240.0, -2400.0, 240.0])[0].numpy()
    usable = np.ones((3, 6), dtype=bool)
    usable[0, 2] = False
    weights = ceiling.prior_weights(cond, usable, decay=1.5)
    assert weights.shape == (3, 6)
    np.testing.assert_allclose(weights.sum(axis=0), 1.0)
    assert weights[0, 2] == 0.0
    assert weights[0, 0] > weights[1, 0]                # recency still rules
    np.testing.assert_allclose(weights[:, 0], weights[:, 1])


def test_prior_weights_ignore_satellite_and_direction():
    """Recency alone: the counterpart arrives gain-matched, and past and future are
    equally good evidence."""
    cond = torch.stack([model.frame_conditioning("available", same, dt)
                        for same, dt in ((True, -240.0), (False, 240.0))]).numpy()
    weights = ceiling.prior_weights(cond, np.ones((2, 6), dtype=bool), decay=1.5)
    np.testing.assert_allclose(weights[0], weights[1])


def test_blend_psnr_blends_each_band_with_its_own_weights():
    """Truth is frame 1 in every band but 3, where it is frame 0 (slightly off,
    so every PSNR is finite): the per-band mix must beat taking frame 1 whole."""
    from suvi import metrics, paths
    radiance = (1.0, 2.0)
    candidates = np.stack([samples.encode_for_model(np.full((6, 8, 8), r, np.float32))
                           for r in radiance])
    truth = np.full((6, 8, 8), radiance[1] * 1.02, np.float32)
    truth[3] = radiance[0] * 1.02
    truth_display = [metrics.to_display(t, w) for t, w in zip(truth, paths.WAVELENGTHS)]
    whole = np.zeros((2, 6))
    whole[1] = 1.0
    mixed = whole.copy()
    mixed[:, 3] = (1.0, 0.0)
    assert (ceiling.blend_psnr(candidates, mixed, truth_display)
            > ceiling.blend_psnr(candidates, whole, truth_display) + 5.0)


def test_score_target_declines_when_a_band_has_no_usable_frame():
    aligned = torch.zeros(1, 2, 6, 8, 8)
    condition = conditioning([-240.0, 240.0])
    present = torch.ones(1, 2, 6)
    present[0, :, 4] = 0.0
    truth = np.ones((6, 8, 8), np.float32)
    assert ceiling.score_target(aligned, condition, present, truth, (1.5,)) is None


def test_score_target_scores_a_partial_stack_on_the_bands_held():
    aligned = torch.full((1, 2, 6, 8, 8), 0.2)
    aligned[0, 0, 4] = 0.9                              # absent band, stand-in pixels
    condition = conditioning([-240.0, 240.0])
    present = torch.ones(1, 2, 6)
    present[0, 0, 4] = 0.0
    truth = samples.decode_from_model(np.full((6, 8, 8), 0.2, np.float32))
    scores = ceiling.score_target(aligned, condition, present, truth, (1.5,))
    assert scores["best1"] == float("inf")
    assert scores["uniform"] == float("inf")            # the 0.9 never blended
    assert scores["prior[1.5]"] == float("inf")