"""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