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_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


def test_every_registered_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, filler in fillers.FILLERS.items():
        result = filler(ctx)
        assert result is not None, name
        assert result.shape == before.shape, name
        assert np.isfinite(result).all(), name