"""The deterministic alignment stage: ephemeris, rotation warp, photometric transfer.

This is the physics the model no longer has to learn, so its correctness bounds the
whole system: a wrong warp or a wrong gain poisons every candidate the network is
allowed to blend.
"""

import datetime as dt
import math

import numpy as np
import pytest

from conftest import solar_disc
from suvi import align, samples

torch = pytest.importorskip("torch")


def unix(year, month, day):
    return dt.datetime(year, month, day, tzinfo=dt.timezone.utc).timestamp()


# ---------------------------------------------------------------------- ephemeris


def test_ephemeris_b0_reaches_its_annual_extremes():
    """B0 swings +/-7.25 deg, peaking in early March and September."""
    assert math.degrees(align.solar_ephemeris(unix(2024, 3, 7))[0]) == pytest.approx(
        -7.25, abs=0.15
    )
    assert math.degrees(align.solar_ephemeris(unix(2024, 9, 8))[0]) == pytest.approx(
        7.25, abs=0.15
    )


def test_ephemeris_b0_crosses_zero_in_june_and_december():
    for when in (unix(2024, 6, 6), unix(2024, 12, 7)):
        assert abs(math.degrees(align.solar_ephemeris(when)[0])) < 0.5


def test_ephemeris_radius_tracks_the_orbit():
    """Apparent radius runs ~944 arcsec at aphelion (July) to ~976 at perihelion."""
    def arcsec(when):
        return align.solar_ephemeris(when)[1] * align.NATIVE_SIZE * align.PLATE_SCALE

    assert arcsec(unix(2024, 1, 3)) == pytest.approx(975.9, abs=1.0)
    assert arcsec(unix(2024, 7, 5)) == pytest.approx(943.9, abs=1.0)


def test_ephemeris_matches_a_real_archive_header():
    """Pinned against dr_suvi-l2-ci094_g16_s20240510T000000Z (verified by hand):
    DIAM_SUN = 760.2932 px, SOLAR_B0 = -3.205694 deg."""
    when = dt.datetime(2024, 5, 10, 0, 2, tzinfo=dt.timezone.utc).timestamp()
    b0, radius_fraction = align.solar_ephemeris(when)
    assert 2 * radius_fraction * 1280 == pytest.approx(760.2932, rel=5e-4)
    assert math.degrees(b0) == pytest.approx(-3.205694, abs=0.05)


# ------------------------------------------------------------------ rotation warp


def numpy_reference(size, dt_seconds, b0, radius_fraction):
    """The already-validated numpy map from suvi.fillers, on the same geometry."""
    from suvi import fillers

    header = {
        "diam_sun": 2 * radius_fraction * size,
        "crpix1": (size + 1) / 2,
        "crpix2": (size + 1) / 2,
        "solar_b0": math.degrees(b0),
    }
    return fillers._rotation_map((size, size), header, dt_seconds)


def test_rotation_grid_matches_the_numpy_map():
    """One implementation trains, the other filled the published baselines; a drift
    between them would score the learned filler against different physics."""
    size, lag, b0, radius = 128, 20 * 3600.0, -0.056, 760.29 / 2 / 1280
    map_x, map_y, visible = numpy_reference(size, lag, b0, radius)

    grid = align.rotation_grid(size, torch.tensor([lag]), torch.tensor([b0]),
                               torch.tensor([radius]))[0].numpy()
    got_x = (grid[..., 0] + 1) * size / 2 - 0.5
    got_y = (grid[..., 1] + 1) * size / 2 - 0.5
    assert np.abs(got_x - map_x)[visible].max() < 1e-3
    assert np.abs(got_y - map_y)[visible].max() < 1e-3


def test_rotation_grid_is_identity_where_nothing_co_rotates():
    """Off-disc and behind-the-limb pixels keep their own values: the corona above
    the limb does not rotate with the photosphere."""
    size, lag, b0, radius = 64, 20 * 3600.0, 0.02, 0.3
    _, _, visible = numpy_reference(size, lag, b0, radius)
    grid = align.rotation_grid(size, torch.tensor([lag]), torch.tensor([b0]),
                               torch.tensor([radius]))[0].numpy()
    grid_y, grid_x = np.mgrid[0:size, 0:size]
    got_x = (grid[..., 0] + 1) * size / 2 - 0.5
    got_y = (grid[..., 1] + 1) * size / 2 - 0.5
    assert np.abs(got_x - grid_x)[~visible].max() == 0.0
    assert np.abs(got_y - grid_y)[~visible].max() == 0.0


def test_rotate_at_zero_dt_is_the_identity():
    frames = torch.randn(2, 6, 96, 96)
    out = align.rotate(frames, torch.zeros(2), torch.tensor([0.05, -0.05]),
                       torch.full((2,), 0.3))
    assert float((out - frames).abs().max()) < 1e-3


def test_rotate_moves_on_disc_content():
    disc = torch.from_numpy(
        np.stack([solar_disc(size=96, radius=30, peak=2.0)] * 6)[None]
    ).float()
    # An off-centre bright blob, so rotation has something visible to move.
    disc[..., 40:48, 30:38] += 3.0
    out = align.rotate(disc, torch.tensor([12 * 3600.0]), torch.tensor([0.0]),
                       torch.tensor([30 / 96]))
    assert float((out - disc).abs().max()) > 0.5


def test_rotate_round_trips():
    """Forward then back must land where it started, on the visible interior."""
    disc = torch.from_numpy(
        np.stack([solar_disc(size=96, radius=30, peak=2.0)] * 6)[None]
    ).float()
    lag = torch.tensor([6 * 3600.0])
    b0, radius = torch.tensor([0.03]), torch.tensor([30 / 96])
    there = align.rotate(disc, lag, b0, radius)
    back = align.rotate(there, -lag, b0, radius)
    centre = (slice(None), slice(None), slice(38, 58), slice(38, 58))
    assert float((back[centre] - disc[centre]).abs().max()) < 0.05


# -------------------------------------------------------------------- photometry


def test_apply_photometry_matches_radiance_arithmetic():
    """The transfer is defined on radiance; applying it in coded space would model a
    different fault entirely."""
    radiance = np.stack([solar_disc(size=32, radius=10, peak=1.0 + b) for b in range(6)])
    coded = torch.from_numpy(samples.encode_for_model(radiance))[None]
    gain = torch.tensor([[0.81, 0.885, 1.49, 1.0, 1.2, 0.86]])
    offset = torch.tensor([[0.0, 0.01, -0.02, 0.0, 0.005, 0.0]])

    moved = align.apply_photometry(coded, gain, offset)
    got = samples.decode_from_model(moved[0].numpy())
    expected = radiance * gain[0, :, None, None].numpy() + offset[0, :, None, None].numpy()
    np.testing.assert_allclose(got, expected, rtol=3e-3, atol=1e-5)


def test_apply_photometry_stays_in_the_coded_range():
    coded = torch.full((1, 6, 8, 8), 1.0)
    moved = align.apply_photometry(coded, torch.full((1, 6), 100.0),
                                   torch.full((1, 6), 1e6))
    assert torch.isfinite(moved).all()
    assert float(moved.abs().max()) <= 1.0


def test_fit_photometry_recovers_a_known_transfer():
    rng = np.random.default_rng(0)
    local = rng.uniform(0.0, 2.0, (6, 32, 32)).astype(np.float32)
    gains = np.array([0.81, 0.885, 1.49, 1.0, 1.2, 0.86], dtype=np.float32)
    offsets = np.array([0.0, 0.01, -0.02, 0.0, 0.005, 0.03], dtype=np.float32)
    counterpart = (local - offsets[:, None, None]) / gains[:, None, None]

    got_gains, got_offsets = align.fit_photometry(counterpart, local)
    np.testing.assert_allclose(got_gains, gains, rtol=1e-4)
    np.testing.assert_allclose(got_offsets, offsets, atol=1e-4)


# ------------------------------------------------------------------- whole stacks


def coded_disc(peak=1.0, size=64):
    radiance = np.stack([solar_disc(size=size, radius=size // 3, peak=peak + b * 0.2)
                         for b in range(6)])
    return radiance, torch.from_numpy(samples.encode_for_model(radiance))


def test_align_stack_puts_the_counterpart_on_the_target_scale():
    gains = np.array([0.81, 0.885, 1.49, 1.0, 1.2, 0.86], dtype=np.float32)
    local_radiance, local = coded_disc()
    counter_radiance = local_radiance / gains[:, None, None]
    counter = torch.from_numpy(samples.encode_for_model(counter_radiance))

    stack = torch.stack([counter])[None]
    aligned = align.align_stack(
        stack, dts=torch.zeros(1, 1), valid=torch.ones(1, 1),
        gains=torch.from_numpy(gains)[None, None],
        offsets=torch.zeros(1, 1, 6),
        b0=torch.tensor([0.0]), radius_fraction=torch.tensor([0.3]),
    )
    assert float((aligned[0, 0] - local).abs().max()) < 5e-3


def test_align_stack_leaves_missing_frames_at_zero():
    """An offset applied to a frame of zeros would manufacture an image from nothing."""
    stack = torch.zeros(1, 1, 6, 32, 32)
    aligned = align.align_stack(
        stack, dts=torch.zeros(1, 1), valid=torch.zeros(1, 1),
        gains=torch.full((1, 1, 6), 3.0), offsets=torch.full((1, 1, 6), 0.5),
        b0=torch.tensor([0.0]), radius_fraction=torch.tensor([0.3]),
    )
    assert float(aligned.abs().max()) == 0.0


def test_align_stack_leaves_a_simultaneous_local_frame_alone():
    _, local = coded_disc()
    stack = torch.stack([local])[None]
    aligned = align.align_stack(
        stack, dts=torch.zeros(1, 1), valid=torch.ones(1, 1),
        gains=torch.ones(1, 1, 6), offsets=torch.zeros(1, 1, 6),
        b0=torch.tensor([0.05]), radius_fraction=torch.tensor([0.3]),
    )
    assert float((aligned[0, 0] - local).abs().max()) < 1e-3


def test_align_stack_warps_toward_the_target_not_away():
    """The direction test the first training run lacked.

    A 'before' frame (dt < 0) must be rotated *forward* onto the target instant.
    The sign error this pins -- warping by dt instead of -dt -- doubled the
    misalignment at long gaps while staying sub-pixel at short ones, so only a
    test that compares against an independently-warped truth can catch it.
    """
    from suvi import fillers

    size, lag, radius = 96, 15 * 3600.0, 30 / 96
    base = np.stack([solar_disc(size=size, radius=30, peak=2.0)] * 6)
    base[:, 40:48, 30:38] += 3.0                        # feature rotation will move
    header = {"diam_sun": 2 * radius * size, "crpix1": (size + 1) / 2,
              "crpix2": (size + 1) / 2, "solar_b0": 0.0}
    # Independent reference: the validated numpy warp advances `base` by +lag.
    truth = np.stack([fillers._warp(band, header, lag, True)[0] for band in base])

    coded = torch.from_numpy(samples.encode_for_model(base))[None, None]
    aligned = align.align_stack(
        coded, dts=torch.full((1, 1), -lag), valid=torch.ones(1, 1),
        gains=torch.ones(1, 1, 6), offsets=torch.zeros(1, 1, 6),
        b0=torch.tensor([0.0]), radius_fraction=torch.tensor([radius]),
    )
    got = samples.decode_from_model(aligned[0, 0].numpy())
    interior = (slice(None), slice(30, 66), slice(30, 66))
    aligned_error = float(np.abs(got[interior] - truth[interior]).mean())
    unwarped_error = float(np.abs(base[interior] - truth[interior]).mean())
    assert aligned_error < unwarped_error * 0.35, (
        f"aligned {aligned_error:.4f} vs unwarped {unwarped_error:.4f}: "
        "the rotation warp is not moving frames onto the target instant"
    )


def test_align_stack_warps_stale_frames_toward_the_target_instant():
    _, local = coded_disc()
    marked = local.clone()
    marked[..., 20:26, 14:20] += 0.4                     # off-centre, on-disc
    stack = torch.stack([marked])[None]
    aligned = align.align_stack(
        stack, dts=torch.full((1, 1), 15 * 3600.0), valid=torch.ones(1, 1),
        gains=torch.ones(1, 1, 6), offsets=torch.zeros(1, 1, 6),
        b0=torch.tensor([0.0]), radius_fraction=torch.tensor([0.3]),
    )
    assert float((aligned[0, 0] - marked).abs().max()) > 0.05


def test_shard_and_ephemeris_share_one_asinh_convention():
    """align.apply_photometry decodes with dataset's constants; if those drift apart
    the photometric transfer silently degrades."""
    values = np.array([[-3.7, 0.0, 1e-3, 0.19, 50.0]], dtype=np.float32)
    coded = torch.from_numpy(samples.encode_for_model(values))
    identity = align.apply_photometry(coded[None, :, None], torch.ones(1, 1),
                                      torch.zeros(1, 1))
    restored = samples.decode_from_model(identity[0, :, 0].numpy())
    np.testing.assert_allclose(restored, values, rtol=2e-3, atol=1e-7)