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


def own_test_weight(row):
    """Item D (own-test mutation score, metadata only for now): reads runs/traj/<run>/own_test_mutation.json when it exists,
    stores it as extra data and does NOT drop or reweight (Kral + Opus 2026-10-06: do not change the acceptance yet)."""
    run = row["id"].split("_r")[-1]
    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))
                return {"keep": True, "weight": 1.0, "extra": {"own_test_mutation": m.get("score"), "own_test_mutants": m.get("mutants")}}
    return {"keep": True, "weight": 1.0}


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