Jyotisha/scripts/research/guided_collect_holdout_replay.py

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

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
    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 [])


def synthetic_event(window: dict[str, Any], index: int) -> dict[str, Any]:
    year = int(window["year"])
    month = int(window["month_lo"])
    last = monthrange(year, month)[1]
    domain = str(window.get("domain") or "career")
    kind = KIND_BY_DOMAIN.get(domain, "career_change")
    stamp = f"{year:04d}-{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 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,
) -> 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,
            "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,
    )
    extras: list[dict[str, Any]] = []
    last_gate = gate
    for index, window in enumerate(windows):
        extras.append(synthetic_event(window, index))
        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,
                "events_needed": len(extras),
                "windows": len(windows),
                "met": True,
                "after_six": gate,
                "final": last_gate,
            }
    return {
        "case_id": case.get("case_id"),
        "radius": radius,
        "events_needed": None,
        "windows": len(windows),
        "met": False,
        "after_six": gate,
        "final": last_gate,
    }


def summarize(rows: Sequence[dict[str, Any]], radius: int) -> dict[str, Any]:
    subset = [row for row in rows if row.get("radius") == radius and not row.get("error")]
    needed = [int(row["events_needed"]) for row in subset if row.get("events_needed") is not None]
    met = sum(1 for row in subset if row.get("met"))
    return {
        "radius": radius,
        "n": len(subset),
        "met": met,
        "median_events_to_gate": statistics.median(needed) if needed 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("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("--json-out", default=str(REPORT_JSON))
    args = parser.parse_args()
    radii = tuple(int(item) for item in str(args.radii).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:
            label = f"{case.get('case_id')} ±{radius}"
            try:
                result = evaluate_case(case, radius)
                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,
                    "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) for radius in radii]
    payload = {
        "generated_at": TODAY.isoformat(),
        "holdout": str(HOLDOUT.relative_to(ROOT)).replace("\\", "/"),
        "ask_count": ASK_COUNT,
        "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())