Jyotisha/scripts/research/guided_collect_holdout_replay.py

#!/usr/bin/env python3
"""Offline T2 replay: after six probes, inject guided_collect_windows events.

F5 (2026-09-16): the injected event now lands on the **true** candidate's own
boundary date for that window's track and index, so the replay measures what a
user who really lived through that window would report. A second, opposite
pass injects the boundary of the candidate furthest from the true minute, as a
control. Neither pass is a merge gate; both numbers go in the progress note.

Does not change production defaults. Writes a compact JSON summary for the
progress note. Not a merge gate.
"""

from __future__ import annotations

import argparse
import json
import statistics
import sys
import time
import traceback
from calendar import monthrange
from datetime import date
from pathlib import Path
from typing import Any, Sequence
from uuid import NAMESPACE_URL, uuid5

ROOT = Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from scripts.active_rectification_event_engine import (  # noqa: E402
    AYANAMSA,
    NODE_MODE,
    compute_candidate_static_contexts,
)
from scripts.rectification.event_probes import (  # noqa: E402
    GUIDED_ANY_DOMAIN,
    GUIDED_COLLECT_LIMIT,
    _birth_year,
    _context_time,
    _guided_track_starts,
    _guided_year_windows,
    _remaining_contexts,
    _scoreable,
    _static_contexts,
    discriminating_event_probes,
    guided_collect_windows,
)
from scripts.rectification.refinement_packet import window_scan  # noqa: E402
from scripts.rectification.scoring_service import (  # noqa: E402
    build_event_contribution_matrix,
    score_from_matrix,
    scoreable_request,
)
from scripts.rectification_policy import DELIVERY_MAX_WIDTH_MINUTES  # noqa: E402
from scripts.research.cluster_width_lib import (  # noqa: E402
    SEPARATION_LEAD,
    delivery_from_public,
    merge_adjacent_traced,
    public_from_clusters,
    raw_signature_clusters,
    still_valid_public,
)
from scripts.research.cluster_width_probe import replay_public  # noqa: E402
from scripts.research.minute_resolution_sweep import MINUTE_STEP, scoring_request_for  # noqa: E402
from scripts.research.probe_supply_after_six import ASK_COUNT, KIND_BY_DOMAIN  # noqa: E402

HOLDOUT = ROOT / "references" / "real_case_calibration" / "minute_rectification_holdout_v4.json"
REPORT_JSON = ROOT / "docs" / "research" / "guided_collect_holdout_2026_09_16.json"
TODAY = date(2026, 9, 16)
RADII = (10, 30, 60)
RANGE_DELIVERY_TIE_PERCENT = 3


def _hhmm(value: object) -> str | None:
    text = str(value or "")[:5]
    return text if len(text) == 5 and text[2] == ":" else None


def load_cases() -> list[dict[str, Any]]:
    payload = json.loads(HOLDOUT.read_text(encoding="utf-8"))
    return list(payload.get("cases") or [])


#: Windows are asked open now (BUG-749); the replay still needs one concrete
#: domain to write into the ledger, so an open window uses this one.
DEFAULT_OPEN_DOMAIN = "career"


def synthetic_event(
    window: dict[str, Any],
    index: int,
    *,
    when: date,
) -> dict[str, Any]:
    """One month-precision event on `when`, labelled with the window's domain."""
    last = monthrange(when.year, when.month)[1]
    raw_domain = str(window.get("domain") or "")
    domain = DEFAULT_OPEN_DOMAIN if raw_domain in {"", GUIDED_ANY_DOMAIN} else raw_domain
    kind = KIND_BY_DOMAIN.get(domain, "career_change")
    stamp = f"{when.year:04d}-{when.month:02d}"
    return {
        "id": str(uuid5(NAMESPACE_URL, f"guided-collect:{index}:{domain}:{stamp}")),
        "domain": domain,
        "event_kind": kind,
        "date_start": f"{stamp}-01",
        "date_end": f"{stamp}-{last:02d}",
        "precision": "month",
        "summary": f"guided {domain} {stamp}",
    }


def _minutes(value: str) -> int:
    hour, minute = value.split(":")[:2]
    return int(hour) * 60 + int(minute)


def window_boundary_dates(
    request: dict[str, Any],
    built: dict[str, Any],
    *,
    candidate_times: Sequence[str],
    today: date,
) -> dict[tuple[int, int, int], dict[str, date]]:
    """Map each guided window to the per-candidate boundary date behind it.

    Mirrors `guided_collect_windows`'s own loop so the (track, index) pairing
    is the engine's, not a guess: for every window key (year, month_lo,
    month_hi) it returns {candidate_time: that candidate's boundary date}.
    """
    birth_date = str(request.get("birth_date") or "").strip()
    birth_year = _birth_year(birth_date)
    if birth_year is None:
        return {}
    remaining = _remaining_contexts(built, candidate_times)
    if len(remaining) < 2:
        remaining = _static_contexts(built)
    work = [item for item in remaining if _scoreable(item)]
    if len(work) < 2:
        return {}
    lo, hi = birth_year + 5, min(today.year, birth_year + 80)
    vim_cache: dict[tuple[Any, ...], list[date]] = {}
    narayana_cache: dict[tuple[Any, ...], list[date] | None] = {}
    starts_by_time: dict[str, dict[str, list[date]]] = {}
    for context in work:
        stamp = _context_time(context)
        if not stamp:
            continue
        starts_by_time[stamp] = _guided_track_starts(
            context,
            birth_date=birth_date,
            lo=lo,
            hi=hi,
            include_pratyantar=True,
            varga_narayana=True,
            vim_cache=vim_cache,
            narayana_cache=narayana_cache,
        )
    out: dict[tuple[int, int, int], dict[str, date]] = {}
    for track in ("vim", "nara", "nara:d9", "nara:d10"):
        lengths = [len(starts_by_time.get(str(_context_time(item)) or "", {}).get(track) or []) for item in work]
        if not lengths or min(lengths) == 0:
            continue
        for index in range(min(lengths)):
            per_time: dict[str, date] = {}
            for context in work:
                stamp = _context_time(context)
                dates = (starts_by_time.get(stamp or "") or {}).get(track) or []
                if stamp and index < len(dates):
                    per_time[stamp] = dates[index]
            gathered = list(per_time.values())
            if len(gathered) < 2 or min(gathered) == max(gathered):
                continue
            for year, month_lo, month_hi in _guided_year_windows(gathered):
                if year < lo or year > hi:
                    continue
                out.setdefault((year, month_lo, month_hi), per_time)
    return out


def posterior_state(
    *,
    rows: Sequence[dict[str, Any]],
    contexts: Sequence[dict[str, Any]],
    probes: Sequence[dict[str, Any]],
    true_time: str,
) -> dict[str, Any]:
    raw = raw_signature_clusters(contexts)
    by_time = {stamp: row for row in rows if (stamp := _hhmm(row.get("time")))}
    merged, _trace = merge_adjacent_traced(raw, by_time)
    public = public_from_clusters(merged, rows)
    prior = {stamp: float(row.get("score") or 0) for row in public if (stamp := _hhmm(row.get("time")))}
    replay = replay_public(probes=probes, public=public, prior=prior, true_time=true_time)
    posterior = []
    eliminated = set(replay["eliminated"])
    scores = dict(replay["scores"])
    for row in public:
        stamp = _hhmm(row.get("time"))
        if not stamp:
            continue
        posterior.append({**row, "score": scores.get(stamp, row.get("score") or 0)})
    valid = still_valid_public(posterior, scores, eliminated, lead=SEPARATION_LEAD)
    delivery = delivery_from_public(valid)
    return {
        "scores": scores,
        "eliminated": eliminated,
        "public": posterior,
        "valid": valid,
        "delivery": delivery,
        "merged": merged,
    }


def precision_gate(valid: Sequence[dict[str, Any]], scores: dict[str, float]) -> dict[str, Any]:
    delivery = delivery_from_public(valid)
    width = delivery.get("width")
    ranked = sorted(
        valid,
        key=lambda row: (
            -float(scores.get(_hhmm(row.get("time")) or "", row.get("score") or 0)),
            str(row.get("time") or ""),
        ),
    )
    if not ranked:
        return {
            "met": False,
            "width": width,
            "tied_for_first": False,
            "gap": None,
            "percents": [],
        }
    if len(ranked) == 1:
        return {
            "met": width is not None and width <= DELIVERY_MAX_WIDTH_MINUTES,
            "width": width,
            "tied_for_first": False,
            "gap": None,
            "percents": [100],
        }
    s0 = float(scores.get(_hhmm(ranked[0].get("time")) or "", ranked[0].get("score") or 0))
    s1 = float(scores.get(_hhmm(ranked[1].get("time")) or "", ranked[1].get("score") or 0))
    tied = s0 == s1
    total = sum(
        max(float(scores.get(_hhmm(row.get("time")) or "", row.get("score") or 0)), 0.0)
        for row in ranked
    )
    percents = []
    for row in ranked[:3]:
        score = max(float(scores.get(_hhmm(row.get("time")) or "", row.get("score") or 0)), 0.0)
        percents.append(round(score / total * 100) if total > 0 else 0)
    gap = percents[0] - percents[1]
    met = (
        width is not None
        and width <= DELIVERY_MAX_WIDTH_MINUTES
        and gap > RANGE_DELIVERY_TIE_PERCENT
        and not tied
    )
    return {
        "met": met,
        "width": width,
        "tied_for_first": tied,
        "gap": gap,
        "percents": percents,
    }


def remaining_times(state: dict[str, Any]) -> list[str]:
    times: list[str] = []
    seen: set[str] = set()
    for row in state["valid"]:
        stamp = _hhmm(row.get("time"))
        if stamp and stamp not in seen and stamp not in state["eliminated"]:
            seen.add(stamp)
            times.append(stamp)
        for member in row.get("cluster_times") or []:
            clock = _hhmm(member)
            if clock and clock not in seen and clock not in state["eliminated"]:
                seen.add(clock)
                times.append(clock)
    return times


def evaluate_case(
    case: dict[str, Any],
    radius: int,
    direction: str = "truth",
) -> dict[str, Any]:
    true_time = str(case["birth"]["time"])[:5]
    request = scoring_request_for(case, radius)
    request["ayanamsa"] = AYANAMSA
    request["node_mode"] = NODE_MODE
    request["minute_step"] = MINUTE_STEP
    static_contexts = compute_candidate_static_contexts(request)
    built = build_event_contribution_matrix(request, static_contexts=static_contexts)
    rows = score_from_matrix(request, built)
    times = [stamp for row in rows if (stamp := _hhmm(row.get("time")))]
    probes = discriminating_event_probes(
        {**request, "refresh_probes": False, "asked_probe_keys": []},
        built,
        scan=window_scan(built),
        candidate_times=times,
        representative_time=true_time,
        today=TODAY,
    )
    state = posterior_state(rows=rows, contexts=static_contexts, probes=probes[:ASK_COUNT], true_time=true_time)
    gate = precision_gate(state["valid"], state["scores"])
    if gate["met"]:
        return {
            "case_id": case.get("case_id"),
            "radius": radius,
            "direction": direction,
            "events_needed": 0,
            "windows": 0,
            "met": True,
            "after_six": gate,
        }
    remaining = remaining_times(state)
    windows = guided_collect_windows(
        request,
        built,
        candidate_times=remaining or times,
        today=TODAY,
    )
    boundaries = window_boundary_dates(
        request,
        built,
        candidate_times=remaining or times,
        today=TODAY,
    )
    pool = remaining or times
    if direction == "truth":
        source_time = true_time
    else:
        source_time = max(
            pool,
            key=lambda stamp: abs(_minutes(stamp) - _minutes(true_time)),
        ) if pool else true_time
    extras: list[dict[str, Any]] = []
    last_gate = gate
    injected_from: list[str] = []
    for index, window in enumerate(windows):
        key = (int(window["year"]), int(window["month_lo"]), int(window["month_hi"]))
        per_time = boundaries.get(key) or {}
        when = per_time.get(source_time)
        if when is None and per_time:
            when = per_time[min(per_time, key=lambda stamp: abs(_minutes(stamp) - _minutes(source_time)))]
        if when is None:
            when = date(int(window["year"]), int(window["month_lo"]), 1)
        injected_from.append(f"{source_time}:{when.isoformat()}")
        extras.append(synthetic_event(window, index, when=when))
        injected = {**request, "events": list(request["events"]) + extras}
        rebuilt = build_event_contribution_matrix(
            scoreable_request(injected),
            static_contexts=static_contexts,
        )
        new_rows = score_from_matrix(scoreable_request(injected), rebuilt)
        state = posterior_state(
            rows=new_rows,
            contexts=static_contexts,
            probes=probes[:ASK_COUNT],
            true_time=true_time,
        )
        last_gate = precision_gate(state["valid"], state["scores"])
        if last_gate["met"]:
            return {
                "case_id": case.get("case_id"),
                "radius": radius,
                "direction": direction,
                "events_needed": len(extras),
                "windows": len(windows),
                "met": True,
                "after_six": gate,
                "final": last_gate,
                "injected": injected_from,
            }
    return {
        "case_id": case.get("case_id"),
        "radius": radius,
        "direction": direction,
        "events_needed": None,
        "windows": len(windows),
        "met": False,
        "after_six": gate,
        "final": last_gate,
        "injected": injected_from,
    }


def summarize(rows: Sequence[dict[str, Any]], radius: int, direction: str) -> dict[str, Any]:
    subset = [
        row for row in rows
        if row.get("radius") == radius and row.get("direction") == direction and not row.get("error")
    ]
    needed = [int(row["events_needed"]) for row in subset if row.get("events_needed") is not None]
    # Cases the six probes had not already settled: the only ones the guided
    # windows can be credited for.
    from_guided = [
        int(row["events_needed"]) for row in subset
        if row.get("events_needed") not in (None, 0)
    ]
    met = sum(1 for row in subset if row.get("met"))
    return {
        "radius": radius,
        "direction": direction,
        "n": len(subset),
        "met": met,
        "met_after_six_only": sum(1 for row in subset if row.get("events_needed") == 0),
        "met_via_guided": len(from_guided),
        "median_events_to_gate": statistics.median(needed) if needed else None,
        "median_events_via_guided": statistics.median(from_guided) if from_guided else None,
        "median_windows": statistics.median([int(row.get("windows") or 0) for row in subset]) if subset else None,
        "unmet": sum(1 for row in subset if not row.get("met")),
        "errors": sum(
            1 for row in rows
            if row.get("radius") == radius and row.get("direction") == direction and row.get("error")
        ),
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--radii", default="10,30,60")
    parser.add_argument("--directions", default="truth,opposite")
    parser.add_argument("--json-out", default=str(REPORT_JSON))
    args = parser.parse_args()
    radii = tuple(int(item) for item in str(args.radii).split(",") if item.strip())
    directions = tuple(item.strip() for item in str(args.directions).split(",") if item.strip())
    cases = load_cases()
    if args.limit:
        cases = cases[: args.limit]
    started = time.perf_counter()
    rows: list[dict[str, Any]] = []
    for case in cases:
        for radius in radii:
            for direction in directions:
                label = f"{case.get('case_id')} ±{radius} {direction}"
                try:
                    result = evaluate_case(case, radius, direction)
                    rows.append(result)
                    print(
                        f"{label} needed={result.get('events_needed')} "
                        f"met={result.get('met')} windows={result.get('windows')}",
                        flush=True,
                    )
                except Exception as exc:  # noqa: BLE001
                    rows.append({
                        "case_id": case.get("case_id"),
                        "radius": radius,
                        "direction": direction,
                        "error": f"{type(exc).__name__}: {exc}",
                        "trace": traceback.format_exc(limit=8),
                        "met": False,
                        "events_needed": None,
                    })
                    print(f"{label} ERROR {type(exc).__name__}: {exc}", flush=True)
    summaries = [summarize(rows, radius, direction) for radius in radii for direction in directions]
    payload = {
        "generated_at": TODAY.isoformat(),
        "holdout": str(HOLDOUT.relative_to(ROOT)).replace("\\", "/"),
        "ask_count": ASK_COUNT,
        "guided_collect_limit": GUIDED_COLLECT_LIMIT,
        "injection": (
            "truth = the true candidate's own boundary date for that window's "
            "track and index; opposite = the boundary of the remaining "
            "candidate furthest from the true minute"
        ),
        "delivery_max_width_minutes": DELIVERY_MAX_WIDTH_MINUTES,
        "tie_percent": RANGE_DELIVERY_TIE_PERCENT,
        "elapsed_s": round(time.perf_counter() - started, 1),
        "summaries": summaries,
        "rows": [
            {key: value for key, value in row.items() if key != "trace"}
            for row in rows
        ],
        "errors": [row for row in rows if row.get("error")],
    }
    out = Path(args.json_out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(summaries, ensure_ascii=False, indent=2), flush=True)
    print(f"wrote {out} in {payload['elapsed_s']}s", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())