"""tools/train_run.py gate: recall-aware mask score and the pass rule."""
import json

from tools import train_run


def _report(f1, iou, tp, fn, classes=None):
    return {"overall": {"f1": f1, "mean_mask_iou": iou, "tp": tp, "fn": fn, "fp": 0, "precision": 1.0, "recall": tp / (tp + fn)},
            "per_class": classes or {}}


def test_mask_score_weights_by_recall():
    assert abs(train_run.mask_score({"mean_mask_iou": 0.9, "tp": 60, "fn": 40}) - 0.54) < 1e-9
    assert train_run.mask_score({"mean_mask_iou": 0.9, "tp": 0, "fn": 0}) == 0.0


def test_gate_prefers_finding_more_with_slightly_looser_masks(tmp_path):
    base = tmp_path / "base.json"
    base.write_text(json.dumps(_report(0.69, 0.903, 60, 40, {"loop_all": {"f1": 0.71}, "pin": {"f1": 0.84}})))
    new = _report(0.85, 0.883, 85, 15, {"loop_all": {"f1": 0.90}, "pin": {"f1": 0.92}})
    g = train_run.gate(new, base)
    assert g["pass"] is True and g["mask_iou_delta"] > 0 and g["matched_only_iou_delta"] < 0


def test_gate_refuses_a_guarded_class_regression(tmp_path):
    base = tmp_path / "base.json"
    base.write_text(json.dumps(_report(0.70, 0.90, 60, 40, {"loop_all": {"f1": 0.80}, "pin": {"f1": 0.84}})))
    new = _report(0.80, 0.90, 80, 20, {"loop_all": {"f1": 0.75}, "pin": {"f1": 0.90}})
    assert train_run.gate(new, base)["pass"] is False


def test_no_baseline_means_no_gate():
    assert train_run.gate(_report(0.8, 0.9, 1, 0), None) is None


def _frames(per_frame_tp, fp=0, objects=2, iou=0.9):
    """One class, ``objects`` gold objects per frame, ``tp`` of them found."""
    return {f"f{i}": {0: (float(tp), float(fp), float(objects - tp), iou * tp)} for i, tp in enumerate(per_frame_tp)}


def test_paired_bootstrap_separates_a_real_gain_from_a_tie():
    from tools import gate_stats
    base = _frames([1] * 60)
    better = _frames([2] * 40 + [1] * 20)            # finds more on two thirds of the frames
    same = _frames([1] * 60)
    r = gate_stats.paired_bootstrap(better, base, draws=400)
    assert r["frames"] == 60 and r["f1"]["win"] > 0.99 and r["f1"]["lo"] > 0 and r["mask_score"]["win"] > 0.99
    t = gate_stats.paired_bootstrap(same, base, draws=400)
    assert t["f1"]["delta"] == 0 and t["f1"]["win"] == 0.0     # never strictly better


def test_paired_bootstrap_pairs_only_on_common_frames():
    from tools import gate_stats
    cand = _frames([2] * 10); base = _frames([1] * 6)
    assert gate_stats.paired_bootstrap(cand, base, draws=50)["frames"] == 6
    assert gate_stats.paired_bootstrap(cand, {}, draws=50) == {"frames": 0, "draws": 0}


def test_gate_with_bootstrap_needs_a_clear_win(tmp_path):
    base = tmp_path / "base.json"
    base.write_text(json.dumps(_report(0.80, 0.90, 80, 20, {"loop_all": {"f1": 0.90}, "pin": {"f1": 0.90}})))
    new = _report(0.806, 0.90, 81, 19, {"loop_all": {"f1": 0.90}, "pin": {"f1": 0.90}})
    noisy = {"frames": 150, "draws": 2000, "f1": {"delta": 0.006, "lo": -0.02, "hi": 0.03, "win": 0.66},
             "mask_score": {"delta": 0.004, "lo": -0.02, "hi": 0.03, "win": 0.60}}
    g = train_run.gate(new, base, noisy)
    assert g["pass_point_rule"] is True and g["pass"] is False and g["rule"] == "paired bootstrap"
    clear = {"frames": 150, "draws": 2000, "f1": {"delta": 0.05, "lo": 0.02, "hi": 0.08, "win": 1.0},
             "mask_score": {"delta": 0.04, "lo": 0.01, "hi": 0.07, "win": 0.99}}
    assert train_run.gate(new, base, clear)["pass"] is True
    lopsided = {"frames": 150, "draws": 2000, "f1": {"delta": 0.03, "lo": 0.0, "hi": 0.06, "win": 0.97},
                "mask_score": {"delta": -0.04, "lo": -0.07, "hi": -0.01, "win": 0.01}}
    assert train_run.gate(new, base, lopsided)["pass"] is False      # wins F1 but clearly loses the masks


def test_gate_without_predictions_keeps_the_point_rule(tmp_path):
    base = tmp_path / "base.json"
    base.write_text(json.dumps(_report(0.69, 0.903, 60, 40)))
    g = train_run.gate(_report(0.85, 0.883, 85, 15), base)
    assert g["rule"] == "point" and g["pass"] is True and g["pass_point_rule"] is True


def test_same_environment_trusts_legacy_reports_only_on_the_spark():
    here = {"host": "gpu102", "ultralytics": "8.4.37", "torch": "2.5.1"}
    assert train_run.same_environment({"ultralytics": "8.4.37", "torch": "2.5.1", "host": "mx-spark-01"}, here) is True
    assert train_run.same_environment({"ultralytics": "8.4.160", "torch": "2.5.1"}, here) is False
    assert train_run.same_environment(None, here) is False                                   # legacy report, read on the cluster
    assert train_run.same_environment(None, {"host": "mx-spark-01", "ultralytics": "x", "torch": "y"}) is True


def test_baseline_from_another_environment_is_measured_again_or_refused(tmp_path, monkeypatch):
    reg = tmp_path / "models" / "registry" / "v-old"; reg.mkdir(parents=True)
    stored = reg / "gold_report.json"
    stored.write_text(json.dumps(dict(_report(0.84, 0.865, 80, 20), environment={"ultralytics": "8.4.160", "torch": "2.14.0"})))
    monkeypatch.setattr(train_run, "ROOT", tmp_path)
    monkeypatch.setattr(train_run, "environment", lambda: {"host": "gpu102", "ultralytics": "8.4.37", "torch": "2.5.1"})
    # no weights here: no baseline to gate against, and a reason
    path, note = train_run.baseline_here(stored, ["loop_all"], 0.25)
    assert path is None and "weights are not here" in note
    # an lfs pointer is not weights either
    (reg / "weights.pt").write_text("version https://git-lfs.github.com/spec/v1")
    assert train_run.baseline_here(stored, ["loop_all"], 0.25)[0] is None
    # real weights: measured again here, the local report is the baseline
    (reg / "weights.pt").write_bytes(b"0" * 1_100_000)
    monkeypatch.setattr(train_run, "measure", lambda w, work, names, conf=0.25: _report(0.86, 0.886, 85, 15))
    path, note = train_run.baseline_here(stored, ["loop_all"], 0.25)
    assert path.name == "gold_report.local.json" and json.loads(path.read_text())["overall"]["f1"] == 0.86
    assert "measured again" in note
    # same environment: the stored report is used as it is
    monkeypatch.setattr(train_run, "environment", lambda: {"host": "x", "ultralytics": "8.4.160", "torch": "2.14.0"})
    assert train_run.baseline_here(stored, ["loop_all"], 0.25) == (stored, None)


def test_bootstrap_by_mount_widens_the_interval_of_repeated_frames():
    """Ten frames of one mount all won by the candidate: frame-wise draws call it a sure win,
    mount-wise draws see one lucky mount among eleven."""
    from tools import gate_stats
    base = _frames([1] * 20)
    cand = _frames([2] * 10 + [1] * 10)
    groups = {f"f{i}": "mount-a" for i in range(10)}
    by_frame = gate_stats.paired_bootstrap(cand, base, draws=800)
    by_mount = gate_stats.paired_bootstrap(cand, base, draws=800, groups=groups)
    assert by_frame["mounts"] == 20 and by_mount["mounts"] == 11
    assert by_frame["f1"]["delta"] == by_mount["f1"]["delta"]          # the point estimate does not move
    assert by_frame["f1"]["win"] > 0.99 and by_mount["f1"]["win"] < 0.75
    assert by_mount["f1"]["lo"] <= 0 < by_frame["f1"]["lo"]


def test_gold_groups_lists_shared_mounts_and_dropped_frames(tmp_path):
    from tools import gate_stats, gold_groups
    gold = tmp_path / "data" / "gold"
    (gold / "labels").mkdir(parents=True); (gold / "images").mkdir()
    manifest = {"images": {"tpi_1": {"tpid": 1, "mount": "PXII:s6416"}, "tpi_2": {"tpid": 2, "mount": "PXII:s6416"},
                           "tpi_3": {"tpid": 3, "mount": "PXIII:s7"}, "tpi_4": {"tpid": 4}},
                "dropped": {"tpi_9": "same image as tpi_1"}}
    (gold / "gold_manifest.json").write_text(json.dumps(manifest))
    for stem in ("tpi_1", "tpi_9"):
        (gold / "labels" / f"{stem}.txt").write_text("")
    assert gold_groups.prune(gold) == 1 and not (gold / "labels" / "tpi_9.txt").exists()
    g = gold_groups.build(gold, with_light=False)
    assert g["mount"] == {"tpi_1": "PXII:s6416", "tpi_2": "PXII:s6416"}      # single-frame mounts are not listed
    assert g["excluded"] == {"tpi_9": "same image as tpi_1"}
    (gold / "groups.json").write_text(json.dumps(g))
    assert gate_stats.load_groups(gold)["mount"]["tpi_2"] == "PXII:s6416"
    assert gate_stats.load_groups(tmp_path)["mount"] == {}                   # no file: every frame its own mount


def test_light_state_thresholds():
    from tools import gold_groups
    assert gold_groups.light_state(0.6, 200) == "saturated"
    assert gold_groups.light_state(0.0, 12) == "dark"
    assert gold_groups.light_state(0.1, 120) == "ok"