from __future__ import annotations

from scripts.research.cluster_width_lib import (
    DEFAULT_MAX_CLUSTERS,
    delivery_from_public,
    merge_adjacent_traced,
    metrics_bundle,
    minute_mass_from_public,
    public_from_clusters,
    range_width,
    smallest_mass_interval,
    still_valid_public,
    time_in_delivery,
    truth_cluster_independent,
)
from scripts.research.cluster_width_probe import previous_sweep_width_is_full_window, verdict


def _cluster(key: str, times: list[str]) -> dict:
    return {
        "signature_key": key,
        "times": times,
        "contexts": [],
        "representative_time": times[len(times) // 2],
        "representative": None,
    }


def _rows(scores: dict[str, float]) -> list[dict]:
    return [{"time": time, "score": score} for time, score in scores.items()]


def test_union_of_clusters_equals_the_window_even_after_merge() -> None:
    times = [f"04:{minute:02d}" for minute in range(0, 21)]
    clusters = [_cluster(str(index), [time]) for index, time in enumerate(times)]
    by_time = {time: {"time": time, "score": 20 if index < 19 else 1} for index, time in enumerate(times)}
    merged, trace = merge_adjacent_traced(clusters, by_time, max_clusters=10)
    covered = {item for cluster in merged for item in cluster["times"]}
    assert covered == set(times)
    assert range_width([item for cluster in merged for item in cluster["times"]]) == 21
    assert len(merged) == 10
    assert len(trace) == 11


def test_w2_refuses_high_delta_merge_and_may_exceed_the_cap() -> None:
    times = [f"05:{minute:02d}" for minute in range(0, 8)]
    clusters = [_cluster(str(index), [time]) for index, time in enumerate(times)]
    by_time = {
        time: {"time": time, "score": 30 if index % 2 == 0 else 1}
        for index, time in enumerate(times)
    }
    merged, trace = merge_adjacent_traced(
        clusters, by_time, max_clusters=3, max_peak_delta=2.0,
    )
    assert len(merged) == 8
    assert trace == []


def test_w2_still_merges_close_scores() -> None:
    times = [f"06:{minute:02d}" for minute in range(0, 6)]
    clusters = [_cluster(str(index), [time]) for index, time in enumerate(times)]
    by_time = {time: {"time": time, "score": 10 + index * 0.1} for index, time in enumerate(times)}
    merged, trace = merge_adjacent_traced(
        clusters, by_time, max_clusters=3, max_peak_delta=2.0,
    )
    assert len(merged) == 3
    assert len(trace) == 3
    assert all(item["delta"] <= 2.0 for item in trace)


def test_no_cap_keeps_raw_clusters() -> None:
    times = [f"07:{minute:02d}" for minute in range(0, 5)]
    clusters = [_cluster(str(index), [time]) for index, time in enumerate(times)]
    by_time = {time: {"time": time, "score": 8} for time in times}
    merged, trace = merge_adjacent_traced(clusters, by_time, max_clusters=None)
    assert len(merged) == 5
    assert trace == []


def test_truth_independence_is_lost_when_the_true_cluster_is_merged() -> None:
    raw = [
        _cluster("a", ["08:00", "08:01"]),
        _cluster("b", ["08:02"]),
        _cluster("c", ["08:03"]),
    ]
    by_time = {
        "08:00": {"score": 1},
        "08:01": {"score": 1},
        "08:02": {"score": 1},
        "08:03": {"score": 20},
    }
    merged, _trace = merge_adjacent_traced(raw, by_time, max_clusters=2)
    assert truth_cluster_independent("08:03", raw, merged) is True
    assert truth_cluster_independent("08:00", raw, merged) is False


def test_mass_interval_shrinks_from_the_weak_tail() -> None:
    mass = {
        "04:50": 10.0,
        "04:51": 10.0,
        "04:52": 10.0,
        "04:53": 1.0,
        "04:54": 1.0,
        "04:55": 1.0,
    }
    interval = smallest_mass_interval(mass, 0.80)
    assert interval["start"] == "04:50"
    assert interval["end"] == "04:52"
    assert interval["width"] == 3
    assert time_in_delivery("04:51", interval["times"]) is True
    assert time_in_delivery("04:55", interval["times"]) is False


def test_mass_interval_does_not_drop_a_peak_in_the_middle() -> None:
    mass = {"10:00": 1.0, "10:05": 20.0, "10:10": 1.0}
    interval = smallest_mass_interval(mass, 0.80)
    assert interval["start"] == "10:05"
    assert interval["end"] == "10:05"
    assert time_in_delivery("10:05", interval["times"]) is True


def test_lead_filter_drops_clusters_eight_or_more_behind() -> None:
    public = [
        {"time": "11:00", "score": 20, "cluster_times": ["11:00", "11:01"]},
        {"time": "11:04", "score": 13, "cluster_times": ["11:04"]},
        {"time": "11:08", "score": 11, "cluster_times": ["11:08"]},
    ]
    scores = {"11:00": 20.0, "11:04": 13.0, "11:08": 11.0}
    valid = still_valid_public(public, scores, set(), lead=8)
    times = delivery_from_public(valid)["times"]
    assert "11:00" in times
    assert "11:04" in times
    assert "11:08" not in times


def test_eliminated_cluster_leaves_the_union() -> None:
    public = [
        {"time": "12:00", "score": 16, "cluster_times": ["12:00", "12:01", "12:02"]},
        {"time": "12:10", "score": 16, "cluster_times": ["12:10", "12:11"]},
    ]
    scores = {"12:00": 16.0, "12:10": 16.0}
    valid = still_valid_public(public, scores, {"12:10"}, lead=8)
    delivered = delivery_from_public(valid)
    assert delivered["start"] == "12:00"
    assert delivered["end"] == "12:02"
    assert delivered["width"] == 3


def test_public_from_clusters_keeps_the_peak_minute() -> None:
    clusters = [_cluster("x", ["13:00", "13:01", "13:02"])]
    rows = _rows({"13:00": 4, "13:01": 9, "13:02": 5})
    public = public_from_clusters(clusters, rows)
    assert len(public) == 1
    assert public[0]["time"] == "13:01"
    assert public[0]["cluster_times"] == ["13:00", "13:01", "13:02"]


def test_minute_mass_uses_the_cluster_peak() -> None:
    public = [{"time": "14:00", "score": 7, "cluster_times": ["14:00", "14:01"]}]
    mass = minute_mass_from_public(public, {"14:00": 7})
    assert mass["14:00"] == 7
    assert mass["14:01"] == 7


def test_metrics_mark_a_squeezed_true_minute() -> None:
    public = [{"time": "15:00", "score": 5, "cluster_times": ["15:00"]}]
    metrics = metrics_bundle(
        public=public,
        true_time="15:10",
        window_times=["15:00", "15:10"],
        delivery_times=["15:00"],
        delivery_width=1,
        independent=True,
    )
    assert metrics["coverage"] is False
    assert metrics["truth_squeezed"] is True
    assert metrics["top1_hit"] is False


def test_previous_sweep_width_is_the_full_window() -> None:
    assert previous_sweep_width_is_full_window({
        "10": {"width_median": 21, "n": 20},
        "30": {"width_median": 61, "n": 20},
        "60": {"width_median": 121, "n": 20},
    }) is True
    assert previous_sweep_width_is_full_window({
        "10": {"width_median": 12, "n": 20},
    }) is False


def test_verdict_blocks_a_coverage_drop() -> None:
    baseline = {
        "n": 10, "top1": 0.3, "coverage": 1.0, "width_median": 21, "tie": 0.2,
        "squeezed": 0, "independent": 0.5, "entropy0": 2.0, "entropy6": 1.5,
    }
    narrower = {**baseline, "width_median": 9, "coverage": 0.8, "squeezed": 2}
    better = {**baseline, "width_median": 9, "independent": 0.8}
    same = {**baseline}
    assert verdict(baseline, narrower) == "no_benefit"
    assert verdict(baseline, better) == "benefit"
    assert verdict(baseline, same) == "no_benefit"
    assert DEFAULT_MAX_CLUSTERS == 64