import numpy as np
import pytest

from conftest import solar_disc
from suvi import fillers

HEADER = {
    "crpix1": 640.5,
    "crpix2": 640.5,
    "cdelt1": 2.5,
    "diam_sun": 771.98,
    "dsun_obs": 148781338180.972,
    "solar_b0": -7.170855,
}


def context(**kwargs):
    base = dict(dt_before=240.0, dt_after=240.0, header=HEADER)
    base.update(kwargs)
    return fillers.FillContext(**base)


# -------------------------------------------------------------------------- context


def test_alpha_locates_the_frame_within_its_gap():
    assert context(dt_before=240, dt_after=240).alpha == pytest.approx(0.5)
    assert context(dt_before=240, dt_after=720).alpha == pytest.approx(0.25)
    assert context(dt_before=0, dt_after=0).alpha == 0.0


def test_gap_frames_counts_slots():
    assert context(dt_before=240, dt_after=240).gap_frames == 2
    assert context(dt_before=240, dt_after=2160).gap_frames == 10


# --------------------------------------------------------------------- baselines


def test_hold_last_repeats_the_preceding_frame():
    before, after = np.ones((8, 8), np.float32), np.zeros((8, 8), np.float32)
    np.testing.assert_array_equal(fillers.hold_last(context(before=before, after=after)), before)


def test_hold_last_falls_back_to_the_following_frame():
    after = np.full((8, 8), 3.0, np.float32)
    np.testing.assert_array_equal(fillers.hold_last(context(after=after)), after)


def test_hold_last_gives_up_with_nothing_to_hold():
    assert fillers.hold_last(context()) is None


def test_linear_blend_is_exact_on_a_linear_ramp():
    """A quantity changing linearly in time must be recovered exactly."""
    before = np.full((8, 8), 10.0, np.float32)
    after = np.full((8, 8), 20.0, np.float32)
    filled = fillers.linear_blend(context(before=before, after=after))
    np.testing.assert_allclose(filled, 15.0, rtol=1e-6)

    skewed = fillers.linear_blend(
        context(before=before, after=after, dt_before=240, dt_after=720)
    )
    np.testing.assert_allclose(skewed, 12.5, rtol=1e-6)


def test_linear_blend_degrades_to_hold_last_at_a_boundary():
    before = np.full((8, 8), 7.0, np.float32)
    np.testing.assert_array_equal(fillers.linear_blend(context(before=before)), before)


def test_fillers_tolerate_non_finite_input():
    before = np.full((8, 8), np.nan, np.float32)
    after = np.ones((8, 8), np.float32)
    filled = fillers.linear_blend(context(before=before, after=after))
    assert np.isfinite(filled).all()


# ------------------------------------------------------------------ optical flow


def test_optical_flow_tracks_a_translation():
    """A feature moving at constant speed should land mid-way, not appear twice."""
    base = solar_disc(size=256, radius=70, peak=3.0)
    before = np.roll(base, -8, axis=1)
    after = np.roll(base, 8, axis=1)
    filled = fillers.optical_flow(context(before=before, after=after))
    blended = fillers.linear_blend(context(before=before, after=after))
    # The truth is the untranslated frame; flow should beat a plain cross-fade.
    assert np.abs(filled - base).mean() < np.abs(blended - base).mean()


def test_optical_flow_falls_back_without_two_brackets():
    before = np.ones((32, 32), np.float32)
    np.testing.assert_array_equal(fillers.optical_flow(context(before=before)), before)


def test_optical_flow_falls_back_on_mismatched_shapes():
    before = np.ones((32, 32), np.float32)
    after = np.ones((16, 16), np.float32)
    result = fillers.optical_flow(context(before=before, after=after))
    np.testing.assert_array_equal(result, before)


# ---------------------------------------------------------------------- crosssat


def test_crosssat_returns_nothing_without_a_counterpart():
    assert fillers.crosssat(context(before=np.ones((8, 8), np.float32))) is None


def test_crosssat_corrects_the_calibration_difference():
    """The other satellite's radiance scale must be matched, not copied blindly."""
    truth = solar_disc(size=128, radius=38, peak=2.0)
    counterpart = truth * 0.5 + 0.3  # a different flight model's response
    filled = fillers.crosssat(context(before=truth, counterpart=counterpart), align=False)
    np.testing.assert_allclose(filled, truth, atol=1e-3)


def test_crosssat_aligns_a_parallax_shift():
    truth = solar_disc(size=128, radius=38, peak=2.0)
    counterpart = np.roll(truth, 4, axis=1)
    filled = fillers.crosssat(context(before=truth, counterpart=counterpart), align=True)
    unaligned = fillers.crosssat(context(before=truth, counterpart=counterpart), align=False)
    assert np.abs(filled - truth).mean() < np.abs(unaligned - truth).mean()


def test_crosssat_calibrates_from_the_simultaneous_pair_not_the_bracket():
    """The defect this fixes: a stale bracket poisons the instrument gain.

    The counterpart is a perfect observation of the target instant on the other
    instrument's scale.  The bracketing frame is twenty hours old, so the Sun in it
    is genuinely brighter -- fitting the gain against it folds that evolution into
    what should be an instrument constant.  A simultaneous pair, however distant in
    time, does not have that problem.
    """
    truth = solar_disc(size=128, radius=38, peak=2.0)
    gain, offset = 0.5, 0.3                      # the other instrument's response
    counterpart = truth * gain + offset
    stale = truth * 1.8                          # the Sun, twenty hours earlier

    pair = (stale * gain + offset, stale)        # both satellites at that same instant
    calibrated = fillers.crosssat(
        context(before=stale, dt_before=72000.0, counterpart=counterpart,
                calibration=pair), align=False
    )
    bracketed = fillers.crosssat(
        context(before=stale, dt_before=72000.0, counterpart=counterpart), align=False
    )
    np.testing.assert_allclose(calibrated, truth, rtol=1e-3, atol=1e-3)
    assert np.abs(bracketed - truth).mean() > 10 * np.abs(calibrated - truth).mean()


def test_crosssat_calibration_is_independent_of_how_distant_the_pair_is():
    """A simultaneous pair is equally valid at any separation from the target."""
    truth = solar_disc(size=128, radius=38, peak=2.0)
    counterpart = truth * 0.5 + 0.3
    errors = []
    for brightness in (1.0, 1.8, 4.0):           # ever more distant, ever more evolved
        evolved = truth * brightness
        pair = (evolved * 0.5 + 0.3, evolved)
        filled = fillers.crosssat(
            context(before=evolved, counterpart=counterpart, calibration=pair),
            align=False,
        )
        errors.append(float(np.abs(filled - truth).mean()))
    assert max(errors) < 1e-3


def test_crosssat_falls_back_to_the_bracket_without_a_pair():
    """No slot anywhere had both satellites good; the bracket is all there is."""
    truth = solar_disc(size=128, radius=38, peak=2.0)
    counterpart = truth * 0.5 + 0.3
    filled = fillers.crosssat(context(before=truth, counterpart=counterpart), align=False)
    np.testing.assert_allclose(filled, truth, atol=1e-3)


def test_crosssat_ignores_a_malformed_calibration_pair():
    truth = solar_disc(size=128, radius=38, peak=2.0)
    counterpart = truth * 0.5 + 0.3
    for pair in ((None, truth), (truth, None), (np.ones((8, 8), np.float32), truth)):
        filled = fillers.crosssat(
            context(before=truth, counterpart=counterpart, calibration=pair), align=False
        )
        np.testing.assert_allclose(filled, truth, atol=1e-3)


def test_gain_fit_returns_identity_for_a_constant_source():
    assert fillers.gain_fit(np.ones((8, 8), np.float32), np.zeros((8, 8), np.float32)) == (
        1.0,
        0.0,
    )


def test_gain_match_recovers_an_affine_transform():
    source = solar_disc(size=64, radius=20, peak=1.0)
    reference = source * 3.0 - 0.5
    np.testing.assert_allclose(fillers.gain_match(source, reference), reference, atol=1e-4)


def test_gain_match_survives_a_constant_source():
    flat = np.ones((8, 8), np.float32)
    result = fillers.gain_match(flat, np.arange(64, dtype=np.float32).reshape(8, 8))
    assert np.isfinite(result).all()


# --------------------------------------------------------------- solar rotation


def test_rotation_rate_is_fastest_at_the_equator():
    equator = fillers.rotation_rate(0.0)
    mid = fillers.rotation_rate(np.radians(45))
    pole = fillers.rotation_rate(np.radians(80))
    assert equator > mid > pole


def test_synodic_rate_is_slower_than_sidereal():
    """An Earth-orbiting observer sees the Sun turn more slowly than the stars do."""
    assert fillers.rotation_rate(0.0, synodic=True) < fillers.rotation_rate(0.0, synodic=False)
    difference = fillers.rotation_rate(0.0, synodic=False) - fillers.rotation_rate(0.0, True)
    assert difference == pytest.approx(fillers.EARTH_ORBIT_DEG_PER_DAY)


def test_rotation_map_is_the_identity_at_zero_lag():
    shape = (128, 128)
    header = dict(HEADER, crpix1=64.5, crpix2=64.5, diam_sun=76.0)
    map_x, map_y, visible = fillers._rotation_map(shape, header, 0.0)
    grid_x, grid_y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
    np.testing.assert_allclose(map_x[visible], grid_x[visible], atol=0.01)
    np.testing.assert_allclose(map_y[visible], grid_y[visible], atol=0.01)


def test_rotation_map_marks_the_off_disc_region_invisible():
    header = dict(HEADER, crpix1=64.5, crpix2=64.5, diam_sun=76.0)
    _, _, visible = fillers._rotation_map((128, 128), header, 0.0)
    assert not visible[0, 0]  # a corner is outside the disc
    assert visible[64, 64]  # disc centre


def test_rotation_warp_round_trips():
    """Warping forward by dt then back by dt must return the original on-disc."""
    header = dict(HEADER, crpix1=64.5, crpix2=64.5, diam_sun=76.0)
    image = solar_disc(size=128, radius=38, peak=2.0)
    forward, visible = fillers._warp(image, header, 3600.0, True)
    back, visible_back = fillers._warp(forward, header, -3600.0, True)
    core = np.zeros_like(visible)
    core[50:78, 50:78] = True  # well inside the disc, away from limb foreshortening
    np.testing.assert_allclose(back[core], image[core], atol=0.05)


def test_solar_rotation_reduces_to_a_blend_off_disc():
    header = dict(HEADER, crpix1=64.5, crpix2=64.5, diam_sun=76.0)
    before = np.full((128, 128), 1.0, np.float32)
    after = np.full((128, 128), 3.0, np.float32)
    filled = fillers.solar_rotation(context(before=before, after=after, header=header))
    assert filled[0, 0] == pytest.approx(2.0, abs=1e-4)  # corner: pure cross-fade


def test_solar_rotation_handles_a_single_bracket():
    header = dict(HEADER, crpix1=64.5, crpix2=64.5, diam_sun=76.0)
    image = solar_disc(size=128, radius=38, peak=2.0)
    filled = fillers.solar_rotation(context(before=image, header=header))
    assert filled is not None and filled.shape == image.shape


def test_solar_rotation_gives_up_with_no_brackets():
    assert fillers.solar_rotation(context()) is None


# ---------------------------------------------------------------------- registry


#: Fillers that reconstruct one band from the two-frame context.  `learned` is not one
#: of them: it fuses a whole stack of six-band frames and reports "not applicable"
#: without one, which is why it is excluded here rather than exempted inside the loop.
CONTEXT_FILLERS = ("hold_last", "linear_blend", "optical_flow", "crosssat",
                   "solar_rotation")


def test_every_registered_filler_is_covered_here():
    """A new filler must be classified deliberately, not silently skipped."""
    assert set(fillers.FILLERS) == set(CONTEXT_FILLERS) | {"learned"}


def test_every_context_filler_is_callable_and_shape_preserving():
    before = solar_disc(size=64, radius=20, peak=2.0)
    after = solar_disc(size=64, radius=20, peak=2.2)
    header = dict(HEADER, crpix1=32.5, crpix2=32.5, diam_sun=40.0)
    ctx = context(before=before, after=after, counterpart=after, header=header)
    for name in CONTEXT_FILLERS:
        result = fillers.FILLERS[name](ctx)
        assert result is not None, name
        assert result.shape == before.shape, name
        assert np.isfinite(result).all(), name


def test_learned_declines_a_two_frame_context():
    """Reporting 'not applicable' beats returning something plausible but untrained."""
    before = solar_disc(size=64, radius=20, peak=2.0)
    assert fillers.learned(context(before=before, counterpart=before)) is None