"""Per-frame measurements stored once and judged from many times.

What must hold: brightness is robust to a flare kernel and resolves a dimmed side
of the disc; the disc's position is found from its limb and not dragged by bright
regions; a double exposure echoes; and the Earth geometry agrees with a full
ephemeris.
"""

import math

import cv2 as cv
import numpy as np
import pytest

from conftest import solar_disc
from suvi import measures

DIAM = 772.0                      # DIAM_SUN for a 386px radius
RADIUS = 386


def textured(size=1280, radius=RADIUS, centre=None, seed=0):
    """A synthetic disc with unique small-scale structure, as real frames have."""
    rng = np.random.default_rng(seed)
    noise = cv.GaussianBlur(rng.normal(0, 1, (size, size)).astype(np.float32), (0, 0), 6.0)
    disc = solar_disc(size=size, radius=radius, centre=centre)
    return disc * (1.0 + 0.6 * noise / noise.std()).clip(0.2, None)


def measured(image, diam=DIAM):
    return measures.pixel_measures(image, diam)


# ------------------------------------------------------------------ brightness


def test_disc_median_ignores_a_flare_kernel_that_lifts_the_mean():
    quiet = solar_disc(active_region=False)
    flaring = quiet.copy()
    flaring[600:640, 700:740] = 200.0
    assert flaring.mean() > 1.5 * quiet.mean()
    assert measured(flaring)["disc_median"] == pytest.approx(measured(quiet)["disc_median"],
                                                             rel=0.02)


def test_a_dimmed_side_shows_in_its_sectors_only():
    """A grazing eclipse dims one side of the disc."""
    clean = solar_disc(active_region=False)
    dimmed = clean.copy()
    dimmed[:, 640:] *= 0.5                                  # the right half (+x)
    before, after = measured(clean), measured(dimmed)
    ratios = [after[f"sector_{k}"] / before[f"sector_{k}"] for k in range(measures.SECTORS)]
    # Sectors are counted from +x towards +y (image rows): 0, 1 and 6, 7 lie on +x.
    assert max(ratios[k] for k in (0, 7)) < 0.6
    assert min(ratios[k] for k in (3, 4)) > 0.95


def test_sectors_cover_the_disc_band():
    out = measured(solar_disc())
    assert all(out[f"sector_{k}"] is not None for k in range(measures.SECTORS))


# ------------------------------------------------------------------ position


def test_a_centred_disc_sits_at_the_centre():
    assert measured(textured())["ring_offset"] < 0.03


def test_a_displaced_disc_is_found_where_it_is():
    shift = 0.3 * RADIUS
    out = measured(textured(centre=(639.5 + shift, 639.5)))
    assert out["ring_offset"] == pytest.approx(0.3, abs=0.03)
    assert out["ring_dx"] > 0.25 and abs(out["ring_dy"]) < 0.05


def test_a_bright_region_does_not_drag_the_position():
    """What sank geometry_v1: a large active region pulls an intensity centroid."""
    image = textured()
    yy, xx = np.mgrid[:1280, :1280]
    image = image + 20.0 * np.exp(-((xx - 900) ** 2 + (yy - 500) ** 2) / (2 * 25.0 ** 2))
    assert measured(image)["ring_offset"] < 0.05


# ------------------------------------------------------------------ echo


def test_a_double_exposure_echoes_at_its_shift():
    single = textured(seed=1)
    shifted = np.roll(single, (int(0.3 * RADIUS), int(0.2 * RADIUS)), axis=(0, 1))
    double = 0.5 * (single + shifted)
    plain, doubled = measured(single), measured(double)
    assert doubled["echo"] > 0.3 and doubled["echo"] > 3 * plain["echo"]
    assert doubled["echo_lag"] == pytest.approx(math.hypot(0.3, 0.2), abs=0.05)


# ------------------------------------------------------------------ flare index


def test_flare_index_separates_a_kernel_from_a_quiet_disc():
    quiet = solar_disc()
    flaring = quiet.copy()
    flaring[600:612, 700:712] = 500.0
    assert measured(flaring)["flare_index"] > 20 * measured(quiet)["flare_index"]


# ------------------------------------------------------------------ edges


def test_nothing_is_measured_without_an_image_or_a_radius():
    assert measures.pixel_measures(None, DIAM) == {}
    assert measures.pixel_measures(solar_disc(), None) == {}
    assert measures.pixel_measures(solar_disc(), 0.0) == {}
    assert measures.pixel_measures(np.zeros(5), DIAM) == {}


def test_an_empty_frame_measures_without_failing():
    out = measured(np.zeros((1280, 1280), dtype=np.float32))
    assert out["disc_median"] == 0.0
    assert out["echo"] is None and out["flare_index"] is None


# ------------------------------------------------------------------ Earth geometry


GEO_KM = 42164.0


def geostationary(longitude_deg):
    lon = math.radians(longitude_deg)
    return [GEO_KM * 1000 * math.cos(lon), GEO_KM * 1000 * math.sin(lon), 0.0]


def test_earth_geometry_agrees_with_a_full_ephemeris():
    from astropy.coordinates import EarthLocation, get_body, solar_system_ephemeris
    from astropy.time import Time
    import astropy.units as u

    position = geostationary(-75.2)                         # GOES-16
    for when in ("2024-04-13T05:00:21", "2024-08-29T04:56:21", "2023-10-17T04:52:21"):
        moment = Time(when, format="isot", scale="utc")
        satellite = EarthLocation.from_geocentric(*position, unit=u.m).get_gcrs(
            moment).cartesian.xyz.to(u.km).value
        with solar_system_ephemeris.set("builtin"):
            sun = get_body("sun", moment).cartesian.xyz.to(u.km).value
        towards = sun - satellite
        span = np.linalg.norm(towards)
        unit = towards / span
        along = -satellite @ unit
        expected = (np.linalg.norm(satellite + along * unit) - 6378.137
                    - along * (695700.0 / span))
        got = measures.earth_tangent_km(position, moment.unix)
        assert got == pytest.approx(expected, abs=15.0)


def test_the_sun_beside_the_satellite_has_no_tangent_height():
    """At local noon the Sun is on the far side of the satellite from Earth."""
    from datetime import datetime, timezone

    noon = datetime(2024, 3, 20, 17, 0, tzinfo=timezone.utc).timestamp()
    assert measures.earth_tangent_km(geostationary(-75.2), noon) is None


def test_local_midnight_at_equinox_is_behind_the_earth():
    from datetime import datetime, timezone

    midnight = datetime(2024, 3, 20, 5, 0, tzinfo=timezone.utc).timestamp()
    assert measures.earth_tangent_km(geostationary(-75.2), midnight) < 0


def test_header_measures_need_a_position_and_a_valid_time():
    cards = {"OBSGEO-X": geostationary(-75.2)[0], "OBSGEO-Y": geostationary(-75.2)[1],
             "OBSGEO-Z": 0.0, "DATE-OBS": "2024-04-13T05:00:21.000"}
    assert "earth_tangent_km" in measures.header_measures(cards)
    assert measures.header_measures({**cards, "OBSGEO-Z": None}) == {}
    assert measures.header_measures({**cards, "DATE-OBS": "not a date"}) == {}
    assert measures.header_measures(None) == {}