import cv2 as cv
import numpy as np
import pytest

from conftest import solar_disc
from suvi import detectors

BAND = 171


def features(slot_band=BAND, time=1715299200, satellite=16, **header):
    """A FrameFeatures whose header describes a healthy frame, before overrides."""
    base = {
        "empty": 0,
        "degraded": 0,
        "eclipse": 0,
        "num_imgs": 2,
        "wavelnth": slot_band,
        "img_mean": 0.35,
        "img_sdev": 0.57,
        "diam_sun": 771.98,
        "crpix1": 640.5,
        "crpix2": 640.5,
    }
    base.update(header)
    return detectors.FrameFeatures(slot=(satellite, slot_band, time), header=base)


def bright_disc(**kwargs):
    """A disc bright enough to clear the 171A geometry threshold of 1.0."""
    return solar_disc(peak=3.0, **kwargs)


# -------------------------------------------------------------------------- header


def test_header_accepts_a_healthy_frame():
    assert detectors.header_v1(features()).verdict == detectors.GOOD


def test_header_records_scores_even_when_passing():
    verdict = detectors.header_v1(features())
    assert verdict.scores["img_mean"] == pytest.approx(0.35)
    assert "crpix_offset" in verdict.scores


@pytest.mark.parametrize(
    "override,fragment",
    [
        ({"empty": 1}, "EMPTY"),
        ({"eclipse": 2}, "ECLIPSE"),
        ({"num_imgs": 0}, "no source images"),
        ({"img_sdev": 0.0}, "zero variance"),
        ({"wavelnth": 304}, "wavelength"),
        ({"img_mean": 1e-5}, "below"),
        ({"img_mean": 500.0}, "above"),
        ({"diam_sun": 100.0}, "DIAM_SUN"),
        ({"crpix1": 700.0}, "sun centre offset"),
    ],
)
def test_header_rejects(override, fragment):
    verdict = detectors.header_v1(features(**override))
    assert verdict.verdict == detectors.BAD
    assert fragment in verdict.reason


def test_header_does_not_reject_on_degraded_alone():
    """Regression: DEGRADED is set for whole months of the 195A band on good frames.

    Rejecting on it discarded 40/40 good 195A frames on 2024-01-20 while flagging
    0/40 on 2024-07-04.  It is recorded as a score and nothing more.
    """
    verdict = detectors.header_v1(features(degraded=1))
    assert verdict.verdict == detectors.GOOD
    assert verdict.scores["degraded"] == 1.0


def test_header_reports_read_errors_as_bad():
    frame = detectors.FrameFeatures(slot=(16, BAND, 0), header={}, error="truncated")
    verdict = detectors.header_v1(frame)
    assert verdict.verdict == detectors.BAD and "truncated" in verdict.reason


def test_header_abstains_without_a_cached_header():
    frame = detectors.FrameFeatures(slot=(16, BAND, 0), header={})
    assert detectors.header_v1(frame).verdict == detectors.UNKNOWN


def test_header_bounds_are_per_band():
    """A radiance normal for 304A is a dropout for 94A."""
    assert detectors.header_v1(features(slot_band=304, img_mean=3.0, wavelnth=304)).verdict == detectors.GOOD
    assert detectors.header_v1(features(slot_band=94, img_mean=3.0, wavelnth=94)).verdict == detectors.BAD


# ------------------------------------------------------------------------ geometry


def test_geometry_accepts_a_synthetic_disc():
    verdict = detectors.geometry_v1(bright_disc(), BAND)
    assert verdict.verdict == detectors.GOOD, verdict.reason


def test_geometry_rejects_a_blank_frame():
    verdict = detectors.geometry_v1(np.zeros((1280, 1280), np.float32), BAND)
    assert verdict.verdict == detectors.BAD and "too low" in verdict.reason


def test_geometry_rejects_a_saturated_frame():
    verdict = detectors.geometry_v1(np.full((1280, 1280), 99.0, np.float32), BAND)
    assert verdict.verdict == detectors.BAD and "too high" in verdict.reason


def test_geometry_rejects_a_displaced_disc():
    verdict = detectors.geometry_v1(bright_disc(centre=(500, 640)), BAND)
    assert verdict.verdict == detectors.BAD


def test_geometry_reports_missing_and_misshapen_input():
    assert detectors.geometry_v1(None, BAND).verdict == detectors.BAD
    small = detectors.geometry_v1(np.zeros((64, 64), np.float32), BAND)
    assert small.verdict == detectors.BAD and "dimensions" in small.reason


def test_geometry_abstains_on_an_unknown_band():
    verdict = detectors.geometry_v1(bright_disc(), 999)
    assert verdict.verdict == detectors.UNKNOWN


def test_geometry_centre_check_is_one_sided():
    """Documents a real defect in the baseline, so a fix cannot land unnoticed.

    The test is `HALF_DIMS//2 - centre > skew`, so brightness pulled toward low
    indices rejects the frame while the same pull toward high indices passes.  This
    asymmetry is why the filter rejects entire days during high solar activity.
    """
    low = detectors.geometry_v1(bright_disc(centre=(560, 640)), BAND)
    high = detectors.geometry_v1(bright_disc(centre=(720, 640)), BAND)
    assert low.scores["centre_x"] < detectors.HALF_DIMS // 2
    assert high.scores["centre_x"] > detectors.HALF_DIMS // 2
    assert low.verdict == detectors.BAD
    assert high.verdict == detectors.GOOD  # symmetric displacement, opposite verdict


def test_geometry_catches_an_all_nan_frame_only_by_luck():
    """An all-NaN frame is caught by the ratio test, not by the shape tests.

    `NaN > x` is False, so the centroid, radius and goodness-of-fit comparisons all
    evaluate False and report "good" on NaN input.  Only the emptiness of the
    threshold mask saves this case -- a frame that was partly NaN could still slip
    through the shape checks.
    """
    verdict = detectors.geometry_v1(np.full((1280, 1280), np.nan, np.float32), BAND)
    assert verdict.verdict == detectors.BAD
    assert "too low" in verdict.reason

    # Demonstrate the underlying blindness: NaN defeats each shape test directly.
    assert not (detectors.HALF_DIMS // 2 - np.nan > detectors.MAX_CENTER_SKEW)
    assert not (np.nan > detectors.MAX_GOF)


# ------------------------------------------------------------------------ temporal


def series(count=25, means=None, thumbs=None):
    out = []
    for index in range(count):
        mean = 0.35 if means is None else means[index]
        frame = detectors.FrameFeatures(
            slot=(16, BAND, 1715299200 + index * 240),
            header={"img_mean": mean},
        )
        if thumbs is not None:
            frame.thumbnail = thumbs[index]
        out.append(frame)
    return out


def drifting_thumbs(count, rng=None):
    """Thumbnails that change a little each step, like the real Sun."""
    rng = rng or np.random.default_rng(0)
    base = solar_disc(size=detectors.THUMBNAIL_SIZE, radius=38, peak=1.0)
    return [(base + rng.normal(0, 0.01, base.shape)).astype(np.float32) for _ in range(count)]


def test_temporal_accepts_a_steady_series():
    verdicts = detectors.temporal_v1(series(thumbs=drifting_thumbs(25)))
    assert all(v.verdict == detectors.GOOD for v in verdicts)


def test_temporal_catches_a_frozen_frame():
    thumbs = drifting_thumbs(25)
    thumbs[12] = thumbs[11].copy()  # the feed stalled
    verdicts = detectors.temporal_v1(series(thumbs=thumbs))
    assert verdicts[12].verdict == detectors.BAD
    assert "identical to previous" in verdicts[12].reason


def test_temporal_catches_a_brightness_step():
    means = [0.35] * 25
    means[12] = 3.5
    verdicts = detectors.temporal_v1(series(means=means, thumbs=drifting_thumbs(25)))
    assert verdicts[12].verdict == detectors.BAD


def test_temporal_catches_a_structural_jump():
    thumbs = drifting_thumbs(25)
    thumbs[12] = np.roll(thumbs[12], 40, axis=0)  # a frame from somewhere else
    verdicts = detectors.temporal_v1(series(thumbs=thumbs))
    assert verdicts[12].verdict == detectors.BAD


def test_temporal_reports_unreadable_frames():
    frames = series(thumbs=drifting_thumbs(25))
    frames[5].error = "truncated"
    verdicts = detectors.temporal_v1(frames)
    assert verdicts[5].verdict == detectors.BAD


def test_temporal_abstains_without_context():
    verdicts = detectors.temporal_v1(series(count=2))
    assert all(v.verdict == detectors.UNKNOWN for v in verdicts)


def test_temporal_handles_an_empty_series():
    assert detectors.temporal_v1([]) == []


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


def cross_pair(scale_b=1.0, shift_b=0, noise=0.005):
    rng = np.random.default_rng(1)
    base = solar_disc(size=detectors.THUMBNAIL_SIZE, radius=38, peak=1.0)
    a = (base + rng.normal(0, noise, base.shape)).astype(np.float32)
    b = (base + rng.normal(0, noise, base.shape)).astype(np.float32) * scale_b
    if shift_b:
        b = np.roll(b, shift_b, axis=1)
    first = detectors.FrameFeatures(slot=(16, BAND, 0), header={"img_mean": 0.35})
    second = detectors.FrameFeatures(slot=(18, BAND, 0), header={"img_mean": 0.35})
    first.thumbnail, second.thumbnail = a, b
    return first, second


def test_crosssat_agrees_on_matching_views():
    a, b = cross_pair()
    first, second = detectors.crosssat_v1(a, b)
    assert first.verdict == detectors.GOOD and second.verdict == detectors.GOOD


def test_crosssat_tolerates_calibration_differences():
    """A modest scale factor between flight models is normal, not a fault."""
    a, b = cross_pair(scale_b=1.3)
    first, _ = detectors.crosssat_v1(a, b)
    assert first.verdict == detectors.GOOD


def test_crosssat_catches_a_blackout_despite_gain_matching():
    """Regression: least-squares matching rescales a 1e-4 frame into agreement.

    Without an explicit bound on the fitted gain this detector called a total
    blackout a match, because the residual after rescaling is tiny.
    """
    a, b = cross_pair(scale_b=1e-4)
    first, second = detectors.crosssat_v1(a, b)
    assert detectors.BAD in (first.verdict, second.verdict) or first.verdict == detectors.UNKNOWN
    assert "gain" in (first.reason or "")


def test_crosssat_blames_the_frame_its_own_history_disowns():
    a, b = cross_pair(scale_b=1e-4)
    suspect = detectors.Verdict(detectors.BAD, "brightness z=40")
    healthy = detectors.Verdict(detectors.GOOD)
    first, second = detectors.crosssat_v1(a, b, temporal_a=healthy, temporal_b=suspect)
    assert first.verdict == detectors.GOOD
    assert second.verdict == detectors.BAD


def test_crosssat_abstains_when_it_cannot_tell_which_is_wrong():
    a, b = cross_pair(scale_b=1e-4)
    first, second = detectors.crosssat_v1(a, b)
    assert first.verdict == detectors.UNKNOWN and second.verdict == detectors.UNKNOWN


def test_crosssat_abstains_without_a_counterpart():
    a, _ = cross_pair()
    first, second = detectors.crosssat_v1(a, None)
    assert first.verdict == detectors.UNKNOWN and second.verdict == detectors.UNKNOWN


def test_crosssat_abstains_without_thumbnails():
    a = detectors.FrameFeatures(slot=(16, BAND, 0), header={})
    b = detectors.FrameFeatures(slot=(18, BAND, 0), header={})
    first, _ = detectors.crosssat_v1(a, b)
    assert first.verdict == detectors.UNKNOWN


# --------------------------------------------------------------------------- utils


def test_thumbnail_shrinks_and_removes_nans():
    image = solar_disc()
    image[0:10, 0:10] = np.nan
    thumb = detectors.thumbnail(image)
    assert thumb.shape == (detectors.THUMBNAIL_SIZE, detectors.THUMBNAIL_SIZE)
    assert np.isfinite(thumb).all()


def test_align_shift_measures_a_known_translation():
    base = solar_disc(size=128, radius=38)
    shifted = np.roll(base, 5, axis=1)
    assert detectors.align_shift(base, shifted) == pytest.approx(5.0, abs=1.0)
    assert detectors.align_shift(base, base) == pytest.approx(0.0, abs=0.5)


def test_align_shift_returns_none_on_mismatched_shapes():
    assert detectors.align_shift(np.zeros((8, 8)), np.zeros((4, 4))) is None
    assert detectors.align_shift(None, np.zeros((4, 4))) is None


def test_robust_z_handles_a_constant_neighbourhood():
    values = [1.0] * 10
    assert detectors._robust_z(values, 5, 4) == 0.0
    values[5] = 9.0
    assert detectors._robust_z(values, 5, 4) == np.inf


def test_robust_z_needs_enough_neighbours():
    assert detectors._robust_z([1.0, 2.0], 0, 4) is None


# ---------------------------------------------------------------------- disc_v1

DISC_HEADER = {"diam_sun": 772.0}


def quiet_disc(peak=1.1):
    """A disc whose quiet regions sit near the 171A threshold, as real frames do."""
    return solar_disc(peak=peak, active_region=False)


def add_active_region(image, offset, strength=6.0, width=0.30, radius=386):
    """Add a bright compact region at `offset` solar radii along +x from centre."""
    yy, xx = np.mgrid[0 : image.shape[0], 0 : image.shape[1]].astype(np.float32)
    centre = (image.shape[1] - 1) / 2.0
    on_disc = np.hypot(xx - centre, yy - centre) < radius
    spot = np.hypot(xx - (centre + radius * offset), yy - centre)
    return image + strength * np.exp(-((spot / (radius * width)) ** 2)) * on_disc


def test_disc_accepts_a_synthetic_disc():
    assert detectors.disc_v1(bright_disc(), BAND, DISC_HEADER).verdict == detectors.GOOD


def test_disc_measurements_are_unmoved_by_an_active_region():
    """The whole point of the detector, and the fix for the 36.9% rejection rate.

    Every disc_v1 measurement is an average over angle, so where the bright regions
    sit does not move it.  geometry_v1 averages along image rows and columns
    instead, so its centroid swings with the active region -- and because its centre
    test is one-sided, an equal displacement is fatal on one limb and harmless on
    the other.
    """
    base = quiet_disc()
    left = add_active_region(base, -0.45)
    right = add_active_region(base, +0.45)

    measurements = [detectors.disc_profile(img, 772.0 / 4.0) for img in (base, left, right)]
    for name in ("radius_ratio", "limb_width"):
        values = [m[name] for m in measurements]
        assert max(values) - min(values) < 1e-6, f"{name} moved: {values}"
    contrasts = [m["limb_contrast"] for m in measurements]
    assert max(contrasts) - min(contrasts) < 0.01

    # Meanwhile the baseline's centroid swings, in opposite directions.
    centroids = [detectors.geometry_v1(img, BAND).scores["centre_x"] for img in (left, base, right)]
    assert centroids[0] < centroids[1] < centroids[2]
    assert centroids[2] - centroids[0] > 5.0

    for img in (base, left, right):
        assert detectors.disc_v1(img, BAND, DISC_HEADER).verdict == detectors.GOOD


@pytest.mark.parametrize(
    "image,label",
    [
        (np.zeros((1280, 1280), np.float32), "all zero"),
        (np.full((1280, 1280), np.nan, np.float32), "all NaN"),
        (np.full((1280, 1280), 1.0, np.float32), "uniform field"),
    ],
)
def test_disc_rejects_frames_with_no_disc(image, label):
    """A frame with no radial structure is a conclusion, not an abstention."""
    verdict = detectors.disc_v1(image, BAND, DISC_HEADER)
    assert verdict.verdict == detectors.BAD, label
    assert verdict.scores["limb_contrast"] == 0.0


def test_disc_detects_a_wrong_sized_disc():
    small = solar_disc(radius=300, peak=3.0, active_region=False)
    measured = detectors.disc_profile(small, 772.0 / 4.0)
    assert measured["radius_ratio"] < 0.85


def test_disc_detects_a_displaced_disc_through_limb_smearing():
    """A decentred disc smears the azimuthally averaged limb; that is the signal."""
    base = solar_disc(peak=3.0, active_region=False)
    sharp = detectors.disc_profile(base, 772.0 / 4.0)
    for shift in (20, 40):
        moved = detectors.disc_profile(np.roll(base, shift, axis=1), 772.0 / 4.0)
        assert moved["limb_width"] > sharp["limb_width"] * 2, f"shift {shift}"


def test_disc_detects_a_blurred_limb():
    base = solar_disc(peak=3.0, active_region=False)
    blurred = cv.GaussianBlur(base, (81, 81), 25)
    assert (
        detectors.disc_profile(blurred, 772.0 / 4.0)["limb_width"]
        > detectors.disc_profile(base, 772.0 / 4.0)["limb_width"] * 2
    )


def test_disc_is_blind_to_rotation_by_construction():
    """Documents a deliberate limit: rotate/yaw_flip belong to the other detectors.

    Any azimuthally averaged quantity is rotation invariant, and a single frame
    carries no absolute rotation reference beyond the CROTA header.
    """
    image = solar_disc(peak=3.0)
    flipped = np.flip(np.flip(image, 0), 1).copy()
    original = detectors.disc_profile(image, 772.0 / 4.0)
    rotated = detectors.disc_profile(flipped, 772.0 / 4.0)
    for name, value in original.items():
        assert rotated[name] == pytest.approx(value, abs=1e-6), name


def test_disc_abstains_without_the_expected_radius():
    verdict = detectors.disc_v1(bright_disc(), BAND, {})
    assert verdict.verdict == detectors.UNKNOWN and "DIAM_SUN" in verdict.reason


def test_disc_abstains_on_an_uncalibrated_band():
    verdict = detectors.disc_v1(bright_disc(), 999, DISC_HEADER)
    assert verdict.verdict == detectors.UNKNOWN


def test_disc_reports_a_missing_image():
    assert detectors.disc_v1(None, BAND, DISC_HEADER).verdict == detectors.BAD


def test_disc_profile_rejects_nonsense_input():
    for image in (None, np.zeros((4, 4), np.float32), np.zeros((8, 8, 3), np.float32)):
        assert detectors.disc_profile(image, 193.0)["limb_contrast"] is None
    assert detectors.disc_profile(bright_disc(), 0)["limb_contrast"] is None
    assert detectors.disc_profile(bright_disc(), np.nan)["limb_contrast"] is None


def test_disc_honours_supplied_bounds():
    """Bounds are per band and injectable, so calibration is testable in isolation."""
    image = bright_disc()
    permissive = {BAND: {"limb_contrast": (0.0, 1.0), "radius_ratio": None, "limb_width": None}}
    strict = {BAND: {"limb_contrast": (0.999, 1.0), "radius_ratio": None, "limb_width": None}}
    assert detectors.disc_v1(image, BAND, DISC_HEADER, permissive).verdict == detectors.GOOD
    assert detectors.disc_v1(image, BAND, DISC_HEADER, strict).verdict == detectors.BAD


def test_disc_bounds_set_to_none_record_a_score_without_judging():
    image = bright_disc()
    bounds = {BAND: {"limb_contrast": None, "radius_ratio": None, "limb_width": None}}
    verdict = detectors.disc_v1(image, BAND, DISC_HEADER, bounds)
    assert verdict.verdict == detectors.GOOD
    assert "radius_ratio" in verdict.scores  # still measured and reported


def test_disc_is_registered_as_a_frame_detector():
    assert "disc_v1" in detectors.FRAME_DETECTORS
    assert "disc_v1" in detectors.ALL_DETECTORS