"""Hooks for train/build_stage2.py (--hook hooks_example:own_test_weight). A hook gets one sample (a dict) and returns
None / False (drop), a number (repeat weight), or {"keep": bool, "weight": float, "extra": {...}}."""
import json
import os

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def identity(row):
    return 1.0


# Weights from the own-test mutation score (Kral + Opus 2026-10-06: down-weight, do not drop). Weight w means: w copies per epoch, a fraction is a
# copy in that share of the epochs (0.5 = every second epoch on average). The acceptance filter itself is unchanged.
OWN_TEST_WEIGHTS = {
    "reliable, score >= 0.75": 1.0,         # the own tests pass on the correct reference and kill at least 3 of 4 mutants
    "reliable, score 0.5 to 0.75": 0.75,
    "reliable, score < 0.5": 0.5,
    "unreliable (own tests fail on the correct reference)": 0.5,   # they may encode model specific behavior
    "no own tests": 0.5,                    # accepted at 85 points at most; the habit of writing tests is part of the behavior we want
    "no signal (PROG, no mutants, not scored)": 1.0,
}


def own_test_class(m):
    if m is None:
        return "no signal (PROG, no mutants, not scored)"
    st = m.get("status")
    if st == "no_own_tests":
        return "no own tests"
    if st != "scored":
        return "no signal (PROG, no mutants, not scored)"
    if not m.get("tests_pass_on_reference"):
        return "unreliable (own tests fail on the correct reference)"
    s = m.get("score")
    if s is None:
        return "no signal (PROG, no mutants, not scored)"
    return "reliable, score >= 0.75" if s >= 0.75 else "reliable, score 0.5 to 0.75" if s >= 0.5 else "reliable, score < 0.5"


def own_test_weight(row):
    """Item D: reads runs/traj/<run>/own_test_mutation.json, stores the score as extra data and sets the weight by OWN_TEST_WEIGHTS."""
    run = row["id"].split("_r")[-1]
    m = None
    for d in os.listdir(os.path.join(ROOT, "runs", "traj")):
        if d.startswith(run + "_"):
            p = os.path.join(ROOT, "runs", "traj", d, "own_test_mutation.json")
            if os.path.exists(p):
                m = json.load(open(p))
            break
    cls = own_test_class(m)
    return {"keep": True, "weight": OWN_TEST_WEIGHTS[cls],
            "extra": {"own_test_class": cls, "own_test_mutation": (m or {}).get("score"), "own_test_reliable": bool((m or {}).get("tests_pass_on_reference"))}}


def repair_up(row):
    """Example of a weight: a trajectory with a repair counts twice."""
    return 2.0 if row.get("repair") else 1.0