"""The rotation-profile fitting script: the measurement that set suvi.align's rate.

The archive-facing plumbing (frame discovery, vetting) reuses tested modules; what
needs pinning here is the fitting math -- the profile parameterisation, the grid
refinement, the holdout bookkeeping, and that the fit recovers a known profile from
a synthetic pair.
"""

import numpy as np
import pytest

import fit_rotation
from conftest import solar_disc
from suvi import fillers


def test_profile_rate_recovers_snodgrass_at_unit_scale():
    """s = 1 must reproduce the literature profile exactly, s = 0 a rigid body."""
    latitudes = np.radians([0.0, 30.0, 60.0])
    snodgrass = 14.713 + fit_rotation.B0_TERM * np.sin(latitudes) ** 2 \
        + fit_rotation.C0_TERM * np.sin(latitudes) ** 4
    np.testing.assert_allclose(
        fit_rotation.profile_rate(14.713, 1.0, latitudes, synodic=False), snodgrass)
    np.testing.assert_allclose(
        fit_rotation.profile_rate(14.713, 0.0, latitudes, synodic=False), 14.713)


def test_profile_rate_synodic_subtracts_earth_orbit():
    sidereal = fit_rotation.profile_rate(14.7, 0.5, 0.3, synodic=False)
    synodic = fit_rotation.profile_rate(14.7, 0.5, 0.3, synodic=True)
    assert sidereal - synodic == pytest.approx(fillers.EARTH_ORBIT_DEG_PER_DAY)


def test_refine_minimum_finds_an_off_grid_parabola_vertex():
    grid = np.arange(10.0, 15.0, 0.5)
    errors = (grid - 12.6) ** 2
    assert fit_rotation.refine_minimum(grid, errors) == pytest.approx(12.6, abs=1e-6)


def test_refine_minimum_clamps_at_the_grid_edge():
    grid = np.arange(10.0, 15.0, 0.5)
    assert fit_rotation.refine_minimum(grid, grid.copy()) == 10.0
    assert fit_rotation.refine_minimum(grid, -grid) == grid[-1]


def test_date_list_spans_the_era_and_hits_the_b0_extremes():
    dates = fit_rotation.build_date_list("2019-12-15", "2025-03-15")
    assert dates[0] == "2019-12-15" and dates[-1] == "2025-03-15"
    months = {d[5:7] for d in dates}
    assert {"03", "09"} <= months  # both B0 extremes are sampled


def test_holdout_split_is_disjoint_and_complete():
    dates = [str(i) for i in range(10)]
    fit, held = fit_rotation.split_holdout(dates, every=3)
    assert set(fit) | set(held) == set(dates)
    assert set(fit) & set(held) == set()
    assert held == ["2", "5", "8"]


def synthetic_pair(a, s, gap_seconds=24 * 3600.0):
    """A frame and its future self, evolved with a known rotation profile."""
    header = {"crpix1": 64.5, "crpix2": 64.5, "diam_sun": 76.0, "solar_b0": 3.0}
    rng = np.random.default_rng(7)
    source = solar_disc(size=128, radius=38, peak=2.0, active_region=False)
    for _ in range(12):  # texture, so the fit has features at every latitude
        row, col = rng.integers(30, 98, 2)
        source[row - 2:row + 3, col - 2:col + 3] += rng.uniform(0.5, 2.0)
    original = fillers.rotation_rate
    fillers.rotation_rate = lambda lat, synodic=True: fit_rotation.profile_rate(
        a, s, lat, synodic)
    try:
        target, _ = fillers._warp(source, header, gap_seconds, True)
    finally:
        fillers.rotation_rate = original
    target = target + rng.normal(0.0, 0.02, target.shape)  # sensor noise: real
    return source, target.astype(np.float32), header, gap_seconds  # pairs never score 0


def test_pair_surface_minimum_lands_on_the_generating_profile():
    """Fitting a pair synthesised with a known (A, s) must recover that (A, s)."""
    a_true, s_true = 14.75, 0.5
    source, target, header, gap = synthetic_pair(a_true, s_true)
    surface = fit_rotation.pair_surface(source, target, header, gap)
    i, j = np.unravel_index(np.argmin(surface), surface.shape)
    assert fit_rotation.A_GRID[i] == pytest.approx(a_true, abs=0.26)
    assert fit_rotation.S_GRID[j] == pytest.approx(s_true, abs=0.26)
    assert surface.min() == pytest.approx(1.0)  # normalised by its own best


def test_holdout_score_interpolates_between_grid_points():
    surface = np.zeros((len(fit_rotation.A_GRID), len(fit_rotation.S_GRID)))
    surface[:, :] = 2.0
    surface[3, :] = 1.0
    surfaces = {"date": {86400: surface}}
    on_grid = fit_rotation.holdout_score(surfaces, fit_rotation.A_GRID[3], 0.5)
    halfway = fit_rotation.holdout_score(
        surfaces, (fit_rotation.A_GRID[3] + fit_rotation.A_GRID[4]) / 2, 0.5)
    assert on_grid == pytest.approx(1.0)
    assert halfway == pytest.approx(1.5)


def test_holdout_score_of_nothing_is_nan():
    assert np.isnan(fit_rotation.holdout_score({}, 14.7, 0.5))