from __future__ import annotations

import json
from pathlib import Path

import pytest

from scripts.research import varga_resolution_m1 as m1
from scripts.research.varga_resolution_lib import (
    RADII,
    THRESHOLDS,
    VARGA_PREFIXES,
    choose_loo_threshold,
    segment_metrics,
    segment_rows,
    threshold_scan,
    valid_minute_scores,
)


ROOT = Path(__file__).resolve().parents[1]


def test_segments_keep_non_contiguous_repeated_signs_separate() -> None:
    rows = [
        {"time": "12:00", "D9": 1},
        {"time": "12:01", "D9": 2},
        {"time": "12:02", "D9": 1},
    ]
    segments = segment_rows(rows, "D9")
    assert [item["value"] for item in segments] == [1, 2, 1]
    assert [item["segment_id"] for item in segments] == [0, 1, 2]


def test_segments_treat_midnight_as_contiguous() -> None:
    rows = [
        {"time": "23:59", "D10": 4},
        {"time": "00:00", "D10": 4},
        {"time": "00:01", "D10": 5},
    ]
    segments = segment_rows(rows, "D10")
    assert len(segments) == 2
    assert segments[0]["times"] == ["23:59", "00:00"]
    assert segments[1]["times"] == ["00:01"]


def test_segments_do_not_bridge_a_missing_minute() -> None:
    rows = [
        {"time": "23:59", "D1": 4},
        {"time": "00:01", "D1": 4},
    ]
    segments = segment_rows(rows, "D1")
    assert len(segments) == 2


def test_m0_metadata_declares_grid_and_set_envelope_distinction() -> None:
    payload = json.loads(
        (ROOT / "artifacts" / "varga-resolution" / "m0-smoke.json").read_text(encoding="utf-8")
    )
    metadata = payload["metadata"]
    assert metadata["scoring_candidate_step_minutes"] == 2
    assert metadata["segment_scan_step_minutes"] == 1
    assert metadata["refresh_probes"] is False
    assert "real_valid_candidate_set" in metadata
    assert "interval_envelope" in metadata


# These are deliberately synthetic arithmetic fixtures, not engine-response
# goldens or evidence of a deterministic full-engine repeat.
def _metric(share, correct=True, retained=True, count=1, tie=False):
    return {
        "top_share": share,
        "top_segment_correct": correct,
        "truth_retained": retained,
        "valid_segment_count": count,
        "top_segment_tie": tie,
    }


def _item(case_id, row, lmt=False):
    return {
        "case_id": case_id,
        "lmt_before_1900": lmt,
        "by_varga": {"D1": {"raw": row}},
    }


@pytest.mark.parametrize(
    ("training", "thresholds", "minimum", "expected"),
    [
        ([], THRESHOLDS, 5, None),
        ([_metric(None), _metric(0.49)], THRESHOLDS, 5, None),
        # A small fold requires all its rows, not an impossible five rows.
        ([_metric(0.6), _metric(0.9)], THRESHOLDS, 5, 0.5),
        # The high-accuracy subset is too small at minimum=5, but not at 4.
        ([_metric(0.9)] * 4 + [_metric(0.5, False)], THRESHOLDS, 5, 0.5),
        ([_metric(0.9)] * 4 + [_metric(0.5, False)], THRESHOLDS, 4, 0.6),
        # Accuracy wins even when retention and coverage are worse.
        ([_metric(0.9, True, False), _metric(0.5, False, True)], (0.5, 0.9), 1, 0.9),
        # Then retention, then coverage, then the lower threshold break ties.
        ([_metric(0.9), _metric(0.5, True, False)], (0.5, 0.9), 1, 0.9),
        ([_metric(0.9), _metric(0.5)], (0.9, 0.5), 1, 0.5),
        ([_metric(0.9)], (0.9, 0.7, 0.5), 1, 0.5),
        ([_metric(0.6)], (0.6,), 5, 0.6),
    ],
    ids=["empty", "no-coverage", "small-fold", "minimum-five", "minimum-four",
         "accuracy-first", "retention-tiebreak", "coverage-tiebreak",
         "lower-threshold-tiebreak", "inclusive-boundary"],
)
def test_choose_loo_threshold(training, thresholds, minimum, expected) -> None:
    assert choose_loo_threshold(training, thresholds, minimum) == expected


def test_loo_selected_excludes_each_held_out_identity(monkeypatch) -> None:
    rows = [_metric(0.5), _metric(0.8, False), _metric(0.9, True, False)]
    items = [_item(f"synthetic-{index}", row, index == 0) for index, row in enumerate(rows)]
    calls = []

    def select(training):
        held_out = len(calls)
        expected = [row for index, row in enumerate(rows) if index != held_out]
        assert len(training) == len(expected)
        assert all(actual is wanted for actual, wanted in zip(training, expected))
        assert all(row is not rows[held_out] for row in training)
        calls.append(training)
        return (0.5, 0.9, 0.9)[held_out]

    monkeypatch.setattr(m1, "choose_loo_threshold", select)
    actual = m1.loo_selected(items, "D1", "raw")
    assert len(calls) == 3
    assert actual == {
        "eligible": 2, "denominator": 3, "coverage": 0.66666667,
        "accuracy": 1.0, "truth_retained": 0.5,
        "selected_threshold_counts": {"0.5": 1, "0.6": 0, "0.7": 0, "0.8": 0, "0.9": 2},
        "validation_rows": [
            {"case_id": "synthetic-0", "threshold": 0.5, "eligible": True,
             "top_segment_correct": True, "truth_retained": True, "lmt_before_1900": True},
            {"case_id": "synthetic-1", "threshold": 0.9, "eligible": False,
             "top_segment_correct": None, "truth_retained": None, "lmt_before_1900": False},
            {"case_id": "synthetic-2", "threshold": 0.9, "eligible": True,
             "top_segment_correct": True, "truth_retained": False, "lmt_before_1900": False},
        ],
    }


def test_loo_selected_singleton_has_no_training_threshold() -> None:
    result = m1.loo_selected([_item("synthetic-only", _metric(1.0))], "D1", "raw")
    assert result == {
        "eligible": 0, "denominator": 1, "coverage": 0.0,
        "accuracy": None, "truth_retained": None,
        "selected_threshold_counts": {str(value): 0 for value in THRESHOLDS},
        "validation_rows": [{
            "case_id": "synthetic-only", "threshold": None, "eligible": False,
            "top_segment_correct": None, "truth_retained": None, "lmt_before_1900": False,
        }],
    }


@pytest.fixture
def synthetic_cluster_state():
    # Two-minute clusters project their representative's score to BOTH minutes.
    # The high-scoring third cluster is posterior but not in the valid set.
    posterior = [
        {"time": "12:00", "cluster_times": ["12:00", "12:01"], "score": 999},
        {"time": "12:02", "score": 999},
        {"time": "12:03", "cluster_times": ["12:03", "12:04"], "score": 999},
        {"time": "12:05", "score": 999},
    ]
    state = {
        "posterior": posterior,
        "valid": [posterior[0], posterior[1], posterior[3]],
        "scores": {"12:00": 2.0, "12:02": 3.0, "12:03": 100.0, "12:05": -1.0},
    }
    rows = [{"time": f"12:0{index}", "D1": sign} for index, sign in enumerate([1, 1, 2, 3, 3, 4])]
    return state, rows


def test_minute_projection_includes_posterior_before_valid_filtering(synthetic_cluster_state) -> None:
    state, rows = synthetic_cluster_state
    assert valid_minute_scores(state, rows) == {
        "12:00": 2.0, "12:01": 2.0, "12:02": 3.0,
        "12:03": 100.0, "12:04": 100.0, "12:05": -1.0,
    }
    # Explicit scores override representatives, but missing keys use row.score.
    projected = valid_minute_scores(state, rows, {"12:00": 7.0})
    assert projected["12:00"] == projected["12:01"] == 7.0
    assert projected["12:02"] == 999.0


@pytest.mark.parametrize(
    ("mode", "qualities", "share"),
    [("raw", [4.0, 3.0, 0.0, -1.0], 0.66666667),
     ("percent", [3.80952381, 2.85714286, 0.0, 0.0], 0.57142857),
     ("uniform", [2.0, 1.0, 0.0, 1.0], 0.5)],
)
def test_segment_metrics_cluster_projection_and_valid_filtering(
    synthetic_cluster_state, mode, qualities, share
) -> None:
    state, rows = synthetic_cluster_state
    # Percent normalizes nonnegative representative scores (total=105), then
    # projects clusters; only segment_metrics filters the invalid 100-point row.
    assert segment_metrics(state, rows, "D1", "12:01", mode) == {
        "prefix": "D1", "mode": mode, "segment_count_window": 4,
        "valid_segment_count": 3, "truth_segment_id": 0, "truth_retained": True,
        "top_segment_correct": True, "top_segment_tie": False,
        "top_segment_ids": [0], "top_share": share, "segment_qualities": qualities,
    }
    excluded = segment_metrics(state, rows, "D1", "12:04", mode)
    assert excluded["truth_segment_id"] == 2
    assert excluded["truth_retained"] is False
    assert excluded["top_segment_correct"] is False


@pytest.mark.parametrize("mode", m1.MODES)
def test_segment_metrics_counts_truth_in_any_tied_top_segment(mode) -> None:
    rows = [{"time": "12:00", "D1": 1}, {"time": "12:01", "D1": 2}]
    state = {"posterior": rows, "valid": rows, "scores": {"12:00": 2.0, "12:01": 2.0}}
    result = segment_metrics(state, rows, "D1", "12:01", mode)
    assert result["top_segment_ids"] == [0, 1]
    assert result["top_segment_tie"] is True
    assert result["top_segment_correct"] is True
    assert result["truth_retained"] is True
    assert result["top_share"] == 0.5
    absent = segment_metrics(state, rows, "D1", "12:02", mode)
    assert absent["truth_segment_id"] is None
    assert absent["truth_retained"] is False
    assert absent["top_segment_correct"] is False


@pytest.mark.parametrize("mode", m1.MODES)
def test_segment_metrics_empty_window_and_empty_valid_set(mode) -> None:
    empty = {"posterior": [], "valid": [], "scores": {}}
    result = segment_metrics(empty, [], "D1", "12:00", mode)
    assert result == {
        "prefix": "D1", "mode": mode, "segment_count_window": 0,
        "valid_segment_count": 0, "truth_segment_id": None, "truth_retained": False,
        "top_segment_correct": False, "top_segment_tie": False,
        "top_segment_ids": [], "top_share": None, "segment_qualities": [],
    }
    rows = [{"time": "12:00", "D1": 1}]
    result = segment_metrics({**empty, "posterior": rows, "scores": {"12:00": 2}}, rows, "D1", "12:00", mode)
    assert result == {
        "prefix": "D1", "mode": mode, "segment_count_window": 1,
        "valid_segment_count": 0, "truth_segment_id": 0, "truth_retained": False,
        "top_segment_correct": False, "top_segment_tie": False,
        "top_segment_ids": [], "top_share": None, "segment_qualities": [0.0],
    }


@pytest.mark.parametrize("mode", m1.MODES)
@pytest.mark.parametrize("score", [0.0, -2.0])
def test_segment_metrics_nonpositive_scores_do_not_create_a_scored_leader(mode, score) -> None:
    rows = [{"time": "12:00", "D1": 1}]
    state = {"posterior": rows, "valid": rows, "scores": {"12:00": score}}
    result = segment_metrics(state, rows, "D1", "12:00", mode)
    assert result["valid_segment_count"] == 1
    assert result["truth_retained"] is True
    assert result["top_segment_tie"] is False
    assert result["segment_qualities"] == ([1.0] if mode == "uniform" else [0.0] if mode == "percent" else [score])
    assert result["top_share"] == (1.0 if mode == "uniform" else None)
    assert result["top_segment_ids"] == ([0] if mode == "uniform" else [])
    assert result["top_segment_correct"] is (mode == "uniform")


def test_threshold_scan_empty_no_coverage_and_inclusive_boundary() -> None:
    assert threshold_scan([], [0.5]) == {
        "0.5": {"n": 0, "denominator": 0, "coverage": None, "accuracy": None, "truth_retained": None}
    }
    rows = [_metric(0.5), _metric(0.9, False, False), _metric(0.49999999), _metric(None)]
    assert threshold_scan(rows, [0.5, 0.9, 1.0]) == {
        "0.5": {"n": 2, "denominator": 4, "coverage": 0.5, "accuracy": 0.5, "truth_retained": 0.5},
        "0.9": {"n": 1, "denominator": 4, "coverage": 0.25, "accuracy": 0.0, "truth_retained": 0.0},
        "1.0": {"n": 0, "denominator": 4, "coverage": 0.0, "accuracy": None, "truth_retained": None},
    }


@pytest.mark.parametrize(
    ("case", "expected"),
    [({"birth": {"date": "1899-12-31"}}, True),
     ({"birth": {"date": "1900-01-01"}}, False),
     ({"birth": {"date": "2000-01-01"}}, False),
     ({"birth": {}}, False), ({}, False)],
)
def test_is_lmt_synthetic_year_boundary(case, expected) -> None:
    assert m1.is_lmt(case) is expected


def test_stratified_rows_preserves_denominators_and_empty_strata() -> None:
    items = [
        _item("synthetic-old", _metric(0.9, True, True, 1, True), True),
        _item("synthetic-new", _metric(0.5, False, True, 3)),
        _item("synthetic-empty", _metric(None, False, False, 0)),
    ]
    result = m1.stratified_rows(items, "D1", "raw")
    for label, expected in {
        "all": (3, 0.46666667, 1, 2, 1, 2),
        "lmt_before_1900": (1, 0.9, 1, 1, 1, 1),
        "post_1900": (2, 0.25, 0, 1, 0, 1),
    }.items():
        row = result[label]
        assert tuple(row[key] for key in (
            "denominator", "top_share_mean", "top_segment_correct", "truth_retained",
            "top_segment_ties", "valid_segment_count_le_2",
        )) == expected
        denominator = expected[0]
        for count in ("top_segment_correct", "truth_retained", "valid_segment_count_le_2"):
            assert row[f"{count}_rate"] == round(row[count] / denominator, 8)
        assert row["top_segment_tie_rate"] == round(row["top_segment_ties"] / denominator, 8)
        assert row["truth_excluded"] == denominator - row["truth_retained"]
    empty = m1.stratified_rows(items[1:], "D1", "raw")["lmt_before_1900"]
    assert empty["denominator"] == 0
    assert empty["top_share_mean"] is None
    assert empty["truth_retained_rate"] is None
    assert empty["top_segment_correct_rate"] is None
    assert empty["top_segment_tie_rate"] is None
    assert empty["valid_segment_count_le_2_rate"] is None
    assert all(row["coverage"] is None for row in empty["thresholds_full_fit"].values())


@pytest.fixture(scope="module")
def m1_artifact():
    # Read the completed engine artifact; never regenerate it in this suite.
    return json.loads(
        (ROOT / "artifacts" / "varga-resolution" / "varga_resolution_m1.json").read_text(encoding="utf-8")
    )


def test_m1_artifact_has_all_unique_case_radius_mode_combinations(m1_artifact) -> None:
    payload = m1_artifact
    assert payload["schema"] == m1.SCHEMA
    assert payload["errors"] == []
    metadata = payload["metadata"]
    assert metadata["case_count_requested"] == metadata["case_count_completed"] == 77
    assert metadata["radii"] == list(RADII) == [10, 15, 30, 60]
    assert metadata["vargas"] == list(VARGA_PREFIXES) == ["D1", "D9", "D10"]
    assert metadata["modes"] == list(m1.MODES) == ["raw", "percent", "uniform"]
    assert metadata["thresholds"] == list(THRESHOLDS)
    cases = m1.load_cases()
    case_ids = {str(case["case_id"]) for case in cases}
    assert len(cases) == len(case_ids) == 77
    items = payload["items"]
    pairs = [(item["case_id"], item["radius"]) for item in items]
    assert len(pairs) == len(set(pairs)) == 308
    assert set(pairs) == {(case_id, radius) for case_id in case_ids for radius in RADII}
    lmt_by_case = {str(case["case_id"]): m1.is_lmt(case) for case in cases}
    for item in items:
        assert item["lmt_before_1900"] is lmt_by_case[item["case_id"]]
        assert set(item["by_varga"]) == set(VARGA_PREFIXES)
        for prefix, modes in item["by_varga"].items():
            assert set(modes) == set(m1.MODES)
            for mode, row in modes.items():
                assert (row["prefix"], row["mode"]) == (prefix, mode)
    aggregates = payload["aggregates"]
    assert len(aggregates) == len(RADII)
    assert {row["radius"] for row in aggregates} == set(RADII)
    for aggregate in aggregates:
        assert aggregate["case_count"] == 77
        assert aggregate["lmt_case_count"] == sum(lmt_by_case.values())
        assert set(aggregate["by_varga"]) == set(VARGA_PREFIXES)
        assert all(set(modes) == set(m1.MODES) for modes in aggregate["by_varga"].values())


@pytest.mark.parametrize("radius", RADII)
@pytest.mark.parametrize("prefix", VARGA_PREFIXES)
@pytest.mark.parametrize("mode", m1.MODES)
def test_m1_stored_aggregates_and_loo_rows_match_offline_recomputation(m1_artifact, radius, prefix, mode) -> None:
    # Recompute only small summary arithmetic, not native_case/engine scoring.
    items = [item for item in m1_artifact["items"] if item["radius"] == radius]
    aggregate = next(row for row in m1_artifact["aggregates"] if row["radius"] == radius)
    stored = aggregate["by_varga"][prefix][mode]
    assert stored["full_fit"] == m1.stratified_rows(items, prefix, mode)
    assert stored["loo_fixed_thresholds"] == m1.loo_fixed(items, prefix, mode)
    loo = m1.loo_selected(items, prefix, mode)
    validation = loo["validation_rows"]
    assert stored["loo_selected_threshold"] == {
        **m1.summarize_loo_rows(validation),
        "selected_threshold_counts": loo["selected_threshold_counts"],
        "validation_rows": validation,
        "lmt_before_1900": m1.summarize_loo_rows([row for row in validation if row["lmt_before_1900"]]),
        "post_1900": m1.summarize_loo_rows([row for row in validation if not row["lmt_before_1900"]]),
    }


def test_m1_report_raw_table_matches_stored_artifact(m1_artifact) -> None:
    report = (ROOT / "docs" / "research" / "rectification_varga_resolution_2026_09_30.md").read_text(encoding="utf-8")
    section = report.split("## 4. M1", 1)[1].split("## 5.", 1)[0]
    actual = [
        [cell.strip() for cell in line.strip().strip("|").split("|")]
        for line in section.splitlines() if line.startswith("| ±")
    ]
    expected = []
    for aggregate in m1_artifact["aggregates"]:
        for prefix in VARGA_PREFIXES:
            raw = aggregate["by_varga"][prefix]["raw"]
            full = raw["full_fit"]["all"]
            loo = raw["loo_selected_threshold"]
            denominator = full["denominator"]
            expected.append([
                f"±{aggregate['radius']}", prefix,
                f"{full['top_segment_correct']}/{denominator}",
                f"{full['truth_retained']}/{denominator}",
                f"{full['valid_segment_count_le_2']}/{denominator}",
                f"{loo['eligible']} / {loo['accuracy']:.3f} / {loo['truth_retained']:.3f}",
            ])
    assert actual == expected