"""The real-defect audit: candidate sources, proposals, corrections and labels.

The audit's output is the ground truth every detector is later judged against, so
the properties pinned here are the ones that would silently corrupt it: a typo in a
correction must fail rather than leave a frame at its guess, a generic far-field
guess can never become a label, and a one-frame defect must not be taken for a
flare when picking the clean-but-hard controls.
"""

import json
import os

import numpy as np
import pytest

import defect_audit as audit
from conftest import solar_disc, write_fits
from suvi import db, detectors, paths

T0 = 1704067200                                  # 2024-01-01 00:00 UTC


# ------------------------------------------------------------------ header spikes


def test_a_one_frame_spread_jump_is_a_spike():
    series = [(T0 + i * 240, 1.0, 1) for i in range(5)]
    series[2] = (T0 + 480, 3.0, 1)
    assert audit.header_spikes(series) == [T0 + 480]


def test_a_sustained_rise_is_not_a_spike():
    """A flare brightens and stays bright; only its first frame could look odd."""
    series = [(T0 + i * 240, s, 1) for i, s in enumerate([1, 1, 3, 3, 3, 3])]
    assert audit.header_spikes(series) == []


def test_spikes_skip_no_long_exposure_frames_and_distant_neighbours():
    series = [(T0, 1.0, 1), (T0 + 240, 3.0, 0), (T0 + 480, 1.0, 1)]
    assert audit.header_spikes(series) == []
    gapped = [(T0, 1.0, 1), (T0 + 5 * 240, 3.0, 1), (T0 + 6 * 240, 1.0, 1)]
    assert audit.header_spikes(gapped) == []


@pytest.mark.parametrize("bad", [None, 0.0, -1.0])
def test_spikes_need_positive_spreads(bad):
    series = [(T0, bad, 1), (T0 + 240, 3.0, 1), (T0 + 480, 1.0, 1)]
    assert audit.header_spikes(series) == []


def test_spikes_on_series_too_short_to_have_neighbours():
    assert audit.header_spikes([]) == []
    assert audit.header_spikes([(T0, 1.0, 1), (T0 + 240, 5.0, 1)]) == []


# -------------------------------------------------------------------- flare peaks


def test_flare_peaks_prefer_a_sustained_rise_over_a_brighter_single_frame():
    times = [T0 + i * 240 for i in range(200)]
    means = [1.0] * 200
    means[50] = 100.0                              # one-frame defect
    for i in range(120, 130):
        means[i] = 10.0                            # a flare
    peaks = audit.flare_peaks(times, means, count=1, separation=3600)
    assert len(peaks) == 1 and T0 + 120 * 240 <= peaks[0] < T0 + 130 * 240


def test_flare_peaks_are_separated_and_counted():
    times = [T0 + i * 240 for i in range(1000)]
    means = np.ones(1000)
    means[100:110] = 10
    means[112:120] = 9                             # the same flare's tail
    means[700:710] = 5
    peaks = audit.flare_peaks(times, list(means), count=5, separation=86400 // 2)
    assert len(peaks) == 2
    assert peaks[1] - peaks[0] >= 86400 // 2
    assert audit.flare_peaks(times[:3], list(means[:3]), count=2) == []


# ---------------------------------------------------------------------- sampling


def test_stratified_sample_caps_each_stratum_without_repeats():
    rng = np.random.default_rng(0)
    rows = [(s, i) for s in "abc" for i in range(10 if s != "c" else 2)]
    chosen = audit.stratified_sample(rows, lambda r: r[0], 4, rng)
    assert len(chosen) == len(set(chosen)) == 10
    assert sum(r[0] == "c" for r in chosen) == 2
    again = audit.stratified_sample(rows, lambda r: r[0], 4, np.random.default_rng(0))
    assert again == chosen


# --------------------------------------------------------------------- proposals


BASE = {171: {"far_field": 0.03, "noise": 0.0004}}


def candidate(**kw):
    out = {"wavelength": 171, "far_field": 0.03, "noise": 0.0004, "n_long": 1,
           "verdict": "good", "reason": None, "sources": ["random"]}
    out.update(kw)
    return out


@pytest.mark.parametrize("fields,expected", [
    ({}, "normal"),
    ({"sources": ["flare"]}, "flare_normal"),
    ({"n_long": 0}, "static"),
    ({"noise": 0.01}, "static"),
    ({"far_field": 0.5}, audit.PENDING),
    ({"verdict": "bad", "reason": "ECLIPSE flag set (2)"}, "dark"),
    ({"verdict": "bad", "reason": "radius_ratio 1.3 outside"}, "other_defect"),
    ({"far_field": None, "noise": None}, "normal"),
])
def test_proposals(fields, expected):
    assert audit.propose(candidate(**fields), BASE) == expected


def test_a_band_without_a_baseline_is_judged_on_header_and_verdict_alone():
    assert audit.propose(candidate(wavelength=94, far_field=50.0), BASE) == "normal"
    assert audit.propose(candidate(wavelength=94, n_long=0), BASE) == "static"


def test_baselines_come_only_from_random_frames_with_a_long_exposure():
    pool = [candidate(far_field=v, noise=v / 100) for v in (1.0, 2.0, 3.0)]
    pool += [candidate(far_field=99.0, sources=["spike"]),
             candidate(far_field=99.0, n_long=0)]
    base = audit.band_baselines(pool)
    assert base[171] == {"far_field": 2.0, "noise": 0.02}
    assert 94 not in base


# ------------------------------------------------------------------- corrections


def test_corrections_parse_with_comments_and_blank_lines():
    text = "# sheet bar_01\n3 block\n\n7 normal   # was a flare loop\n"
    assert audit.parse_corrections(text, {3, 7}) == {3: "block", 7: "normal"}


@pytest.mark.parametrize("text,message", [
    ("3 bar\n9 bar", "no candidate 9"),
    ("3 barr", "unknown category"),
    ("3", "expected 'ID CATEGORY'"),
    ("x bar", "expected 'ID CATEGORY'"),
    (f"3 {audit.PENDING}", "unknown category"),
])
def test_bad_corrections_fail_loudly(text, message):
    with pytest.raises(ValueError, match=message):
        audit.parse_corrections(text, {3})


def test_current_labels_apply_corrections_over_proposals():
    cands = [{"id": 1, "proposed": "normal"}, {"id": 2, "proposed": audit.PENDING}]
    assert audit.current_labels(cands, {2: "bar"}) == {1: "normal", 2: "bar"}


def write_audit(path, candidates):
    with open(path, "w") as handle:
        json.dump({"baselines": {"171": BASE[171]}, "candidates": candidates}, handle)


def test_apply_refuses_while_anything_is_only_proposed(tmp_path):
    source = tmp_path / "c.json"
    write_audit(source, [{"id": 1, "proposed": audit.PENDING, "satellite": 16,
                          "wavelength": 171, "t_start": T0, "sources": ["spike"]}])
    with pytest.raises(SystemExit, match="classify them first"):
        audit.main(["apply", "--candidates", str(source), "--out", str(tmp_path / "l.json")])


def test_apply_writes_sorted_confirmed_labels(tmp_path):
    source, fixes, out = tmp_path / "c.json", tmp_path / "fix.txt", tmp_path / "labels" / "l.json"
    write_audit(source, [
        {"id": 1, "proposed": audit.PENDING, "satellite": 18, "wavelength": 171,
         "t_start": T0, "sources": ["spike"]},
        {"id": 2, "proposed": "normal", "satellite": 16, "wavelength": 171,
         "t_start": T0 + 240, "sources": ["random"]},
    ])
    fixes.write_text("1 bar\n")
    assert audit.main(["apply", "--candidates", str(source), "--corrections", str(fixes),
                       "--out", str(out)]) == 0
    labels = json.loads(out.read_text())["labels"]
    assert [(r["satellite"], r["category"]) for r in labels] == [(16, "normal"), (18, "bar")]


# --------------------------------------------------------- measurement and sheets


def write_frame(root, satellite, wavelength, when, image, **headers):
    name = paths.FrameName(satellite=satellite, wavelength=wavelength, t_start=when,
                           t_end=when + 240, version="1-0-2")
    write_fits(os.path.join(str(root), name.relpath()), image,
               headers={"WAVELNTH": wavelength, **headers})
    return name


def test_measure_reports_far_field_grain_and_verdict(archive):
    image = solar_disc(peak=1.0)
    clean = write_frame(archive, 16, 171, T0, image)
    static = write_frame(archive, 16, 171, T0 + 240,
                         image + np.random.default_rng(1).normal(0, 0.2, image.shape)
                         .astype(np.float32), N_LONG=0)
    a = audit.measure((16, 171, T0, paths.abspath(clean.relpath(), str(archive))))
    b = audit.measure((16, 171, T0 + 240, paths.abspath(static.relpath(), str(archive))))
    assert a["verdict"] == "good" and a["n_long"] == 1 and a["far_field"] < 0.01
    assert b["verdict"] == "bad" and "N_LONG" in b["reason"] and b["noise"] > 100 * a["noise"]


def test_sheets_draw_every_candidate_grouped_by_category(archive, tmp_path):
    image = solar_disc(peak=1.0)
    cands = []
    for i, proposed in enumerate(["normal", audit.PENDING, "normal"]):
        name = write_frame(archive, 16, 171, T0 + i * 240, image)
        cands.append({"id": i + 1, "proposed": proposed, "satellite": 16, "wavelength": 171,
                      "t_start": T0 + i * 240, "path": name.relpath(), "far_field": 0.03,
                      "noise": 0.0004, "verdict": "good", "sources": ["random"]})
    source, fixes = tmp_path / "c.json", tmp_path / "fix.txt"
    write_audit(source, cands)
    fixes.write_text("2 ghost\n")
    out = tmp_path / "sheets"
    assert audit.main(["sheets", "--candidates", str(source), "--corrections", str(fixes),
                       "--out-dir", str(out), "--workers", "1"]) == 0
    assert sorted(os.listdir(out)) == ["ghost_01.jpg", "index.txt", "normal_01.jpg"]
    assert "ids 1..3 (2 frames)" in (out / "index.txt").read_text()


# ------------------------------------------------------------------- the sources


def test_gather_finds_every_source_in_a_scratch_index(tmp_path):
    """Spikes, N_LONG=0, legacy rejections, flare windows and random controls."""
    conn = db.connect(str(tmp_path / "index.sqlite"))
    ids = {}
    for satellite in (16, 18):
        for wavelength in paths.WAVELENGTHS:
            for i in range(400):
                when = T0 + i * 240
                name = paths.FrameName(satellite=satellite, wavelength=wavelength,
                                       t_start=when, t_end=when + 240, version="1-0-2")
                frame_id = db.upsert_frame(conn, name, name.relpath(), size_bytes=1_500_000)
                ids[(satellite, wavelength, when)] = frame_id
                flare = wavelength == 131 and 300 <= i < 310
                db.record_headers(conn, [(frame_id, {
                    "img_sdev": 5.0 if (satellite, wavelength, i) == (18, 171, 100) else 1.0,
                    "img_mean": 9.0 if flare else 1.0,
                    "n_long": 0 if (satellite, wavelength, i) == (16, 94, 200) else 1,
                }, None)])
    run = db.create_detector_run(conn, "legacy_filter", {"source": "test"})
    db.record_detections(conn, run, [(ids[(16, 304, T0 + 50 * 240)], "bad", None, None, 0)])
    conn.commit()

    found, peaks = audit.gather(conn, T0, T0 + 400 * 240, np.random.default_rng(0),
                                random_per_stratum=3, flare_count=1, flare_frames=20)
    by_source = {}
    for key, (_, sources) in found.items():
        for source in sources:
            by_source.setdefault(source, set()).add(key)
    assert by_source["spike"] == {(18, 171, T0 + 100 * 240)}
    assert by_source["n_long0"] == {(16, 94, T0 + 200 * 240)}
    assert by_source["legacy"] == {(16, 304, T0 + 50 * 240)}
    assert T0 + 300 * 240 <= peaks[0] < T0 + 310 * 240
    assert all(abs(t - peaks[0]) <= audit.FLARE_WINDOW for _, _, t in by_source["flare"])
    assert (16, 94, T0 + 200 * 240) not in by_source.get("random", set())