#!/usr/bin/env python3
"""Varga-resolution research runner (TASK-rectification-varga-resolution-research-20260930).

Stages (``--stages``, comma list, default all):

* ``m0``     — Table 1 (signs per window) and Table 2 (chart type inside the
               six-probe delivery interval) of the brief, reproduced from
               `varga_resolution_lib`; written to ``--baseline-out``.
* ``m1``     — segment shares after the production six probes, raw / percent /
               uniform mass; full-fit threshold table and leave-one-case-out
               confidence labels per varga × radius; truth-segment retention.
* ``m2``     — production probe order vs segment-information-gain order on the
               same pool, per target varga; also with a stop rule.
* ``m3``     — question-domain strategies (career / relationship / general):
               joint segments, "no rectification needed" share, accuracy,
               questions with stop rule.
* ``robust`` — 1 / 2 flipped answers (5 seeds), ±7-day shift of day-precision
               events, LMT-era cases split out.

Everything is offline; no production file is touched. ``PYTHONHASHSEED=0`` and
two runs must be byte-identical (``--json-out`` is written with sorted keys).

Usage::

    PYTHONHASHSEED=0 python3 scripts/research/varga_resolution_probe.py \
        --stages m0,m1,m2,m3,robust --radii 10,30,60 \
        --cache-dir /tmp/varga-cache \
        --baseline-out docs/research/varga_resolution_baseline_2026_09_30.json \
        --json-out docs/research/varga_resolution_research_2026_09_30.json
"""

from __future__ import annotations

import argparse
import statistics
import sys
import time
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any, Sequence

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

from scripts.research import varga_resolution_lib as vr  # noqa: E402

BASELINE_JSON = ROOT / "docs" / "research" / "varga_resolution_baseline_2026_09_30.json"
REPORT_JSON = ROOT / "docs" / "research" / "varga_resolution_research_2026_09_30.json"
TABLE_KEYS: dict[str, tuple[str, ...]] = {
    "D1": ("D1",), "D9": ("D9",), "D10": ("D10",), "D12": ("D12",), "D1xD9xD10": ("D1", "D9", "D10"),
}
STOP_SHARES: tuple[float, ...] = (0.7, 0.8)


def _rate(values: Sequence[bool]) -> float | None:
    return round(sum(1 for v in values if v) / len(values), 4) if values else None


def _median(values: Sequence[float | int | None]) -> float | None:
    clean = [float(v) for v in values if v is not None]
    return statistics.median(clean) if clean else None


# ---------------------------------------------------------------------------
# M0
# ---------------------------------------------------------------------------

def interval_offsets(state: dict[str, Any], true_time: str, radius: int) -> list[int]:
    """Every minute (step 1) from delivery start to end, as offsets from the true minute."""
    start, end = state["delivery"].get("start"), state["delivery"].get("end")
    if start is None or end is None:
        return []
    # Forward span from start to end (through midnight when the window straddles
    # it, e.g. a 00:01–23:59 range around a 23:15 birth), every offset folded
    # into [-720, 720] and clipped to the search window. Interior minutes that
    # were eliminated stay inside: this is the range the product shows.
    span = (vr.clock(end) - vr.clock(start)) % 1440
    first = vr.offset_of(start, true_time)
    offsets = []
    for step in range(span + 1):
        offset = ((first + step + 720) % 1440) - 720
        if -radius <= offset <= radius:
            offsets.append(offset)
    return sorted(set(offsets))


def table2_row(signs: dict[int, dict[str, int]], offsets: Sequence[int]) -> dict[str, Any]:
    row: dict[str, Any] = {}
    for label, vargas in TABLE_KEYS.items():
        keys = [tuple(signs[o][v] for v in vargas) for o in offsets if o in signs]
        if not keys:
            row[label] = {"kinds": 0, "single": False, "le2": False, "mode_is_truth": False}
            continue
        counter = Counter(keys)  # insertion order = ascending offset; ties go to the earliest sign
        ranked = counter.most_common()
        mode_key = ranked[0][0]
        tied = len(ranked) > 1 and ranked[1][1] == ranked[0][1]
        truth_key = tuple(signs[0][v] for v in vargas)
        row[label] = {
            "kinds": len(counter),
            "single": len(counter) == 1,
            "le2": len(counter) <= 2,
            "mode_is_truth": mode_key == truth_key,
            "mode_tied": tied,
        }
    return row


def stage_m0(replays: dict[tuple[str, int], vr.CaseReplay], signs60: dict[str, dict[int, dict[str, int]]], radii: Sequence[int]) -> dict[str, Any]:
    table1 = vr.count_table(signs60)
    per_case: list[dict[str, Any]] = []
    for (case_id, radius), rep in sorted(replays.items()):
        state = vr.replay_scores(rep.public, rep.probes, rep.true_time)
        offsets = interval_offsets(state, rep.true_time, radius)
        per_case.append({
            "case_id": case_id, "radius": radius,
            "start": state["delivery"].get("start"), "end": state["delivery"].get("end"),
            "width": state["delivery"].get("width"),
            "truth_in_range": vr.truth_in_delivery(state, rep.true_time),
            "table2": table2_row(rep.signs, offsets),
        })
    table2: dict[str, Any] = {}
    for radius in radii:
        rows = [r for r in per_case if r["radius"] == radius]
        summary: dict[str, Any] = {
            "n": len(rows),
            "truth_in_range": sum(1 for r in rows if r["truth_in_range"]),
            "width_median": _median([r["width"] for r in rows]),
        }
        for label in TABLE_KEYS:
            summary[label] = {
                "single": sum(1 for r in rows if r["table2"][label]["single"]),
                "le2": sum(1 for r in rows if r["table2"][label]["le2"]),
                "mode_is_truth": sum(1 for r in rows if r["table2"][label]["mode_is_truth"]),
                "mode_tied": sum(1 for r in rows if r["table2"][label]["mode_tied"]),
                "mean_kinds": round(sum(r["table2"][label]["kinds"] for r in rows) / len(rows), 2) if rows else None,
            }
        table2[str(radius)] = summary
    return {"table1_signs_per_window": table1, "table2_delivery_interval": table2, "per_case": per_case}


# ---------------------------------------------------------------------------
# M1
# ---------------------------------------------------------------------------

def share_points(replays: dict[tuple[str, int], vr.CaseReplay], radius: int, vargas: Sequence[str], mode: str,
                 *, cases_filter: set[str] | None = None) -> list[dict[str, Any]]:
    out = []
    for (case_id, r), rep in sorted(replays.items()):
        if r != radius or (cases_filter is not None and case_id not in cases_filter):
            continue
        state = vr.replay_scores(rep.public, rep.probes, rep.true_time)
        share = vr.segment_shares(vr.minute_weights(state, mode), rep.segments(vargas), rep.true_time)
        out.append({"case_id": case_id, **share, "segments_in_window": len(rep.segments(vargas))})
    return out


def summarize_shares(points: Sequence[dict[str, Any]]) -> dict[str, Any]:
    pairs = [(p["top_share"], p["top_is_truth"]) for p in points]
    return {
        "n": len(points),
        "truth_retained": sum(1 for p in points if p["truth_retained"]),
        "top_is_truth": sum(1 for p in points if p["top_is_truth"]),
        "alive_le2": sum(1 for p in points if p["alive_segments"] <= 2),
        "alive_single": sum(1 for p in points if p["alive_segments"] == 1),
        "top_share_median": _median([p["top_share"] for p in points]),
        "full_fit": vr.threshold_table(pairs),
        "loo": [vr.loo_threshold(pairs, target=t) for t in vr.TARGET_ACCURACIES],
    }


def stage_m1(replays: dict[tuple[str, int], vr.CaseReplay], radii: Sequence[int], lmt_ids: set[str]) -> dict[str, Any]:
    out: dict[str, Any] = {"by_radius": {}, "lmt_era_case_ids": sorted(lmt_ids)}
    for radius in radii:
        block: dict[str, Any] = {}
        for label, vargas in TABLE_KEYS.items():
            block[label] = {}
            for mode in vr.MODES:
                points = share_points(replays, radius, vargas, mode)
                summary = summarize_shares(points)
                non_lmt = [p for p in points if p["case_id"] not in lmt_ids]
                summary["excluding_lmt_era"] = {
                    "n": len(non_lmt),
                    "truth_retained": sum(1 for p in non_lmt if p["truth_retained"]),
                    "top_is_truth": sum(1 for p in non_lmt if p["top_is_truth"]),
                }
                summary["per_case"] = [
                    {k: p[k] for k in ("case_id", "top_share", "top_is_truth", "truth_retained", "alive_segments", "segments_in_window")}
                    for p in points
                ]
                block[label][mode] = summary
        out["by_radius"][str(radius)] = block
    return out


# ---------------------------------------------------------------------------
# M2
# ---------------------------------------------------------------------------

def order_outcome(rep: vr.CaseReplay, state: dict[str, Any], vargas: Sequence[str], mode: str) -> dict[str, Any]:
    share = vr.segment_shares(vr.minute_weights(state, mode), rep.segments(vargas), rep.true_time)
    return {
        "questions": len(state.get("asked", [])) if "asked" in state else state["answered"],
        "answered": state["answered"],
        "truth_retained": share["truth_retained"],
        "top_is_truth": share["top_is_truth"],
        "top_share": share["top_share"],
        "alive_segments": share["alive_segments"],
        "width": state["delivery"].get("width"),
        "truth_in_range": vr.truth_in_delivery(state, rep.true_time),
    }


def stage_m2(replays: dict[tuple[str, int], vr.CaseReplay], radii: Sequence[int], mode: str = "raw") -> dict[str, Any]:
    targets = {"D9": ("D9",), "D10": ("D10",), "D1xD9xD10": ("D1", "D9", "D10")}
    out: dict[str, Any] = {"mode": mode, "by_radius": {}}
    for radius in radii:
        block: dict[str, Any] = {}
        for label, vargas in targets.items():
            rows = []
            for (case_id, r), rep in sorted(replays.items()):
                if r != radius:
                    continue
                segments = rep.segments(vargas)
                prod = vr.replay_scores(rep.public, rep.probes, rep.true_time)
                prod["asked"] = list(rep.probes)[:vr.ASK_COUNT]
                seg = vr.select_probes_by_segment(rep.public, rep.probes, segments, rep.true_time, mode=mode)
                row: dict[str, Any] = {
                    "case_id": case_id,
                    "pool": len(rep.probes),
                    "production": order_outcome(rep, prod, vargas, mode),
                    "segment_order": order_outcome(rep, seg, vargas, mode),
                    "same_first_six": [p.get("semantic_key") for p in prod["asked"]] == [p.get("semantic_key") for p in seg["asked"]],
                }
                for stop in STOP_SHARES:
                    prod_stop = vr.production_order_with_stop(rep.public, rep.probes, segments, rep.true_time, mode=mode, stop_share=stop)
                    seg_stop = vr.select_probes_by_segment(rep.public, rep.probes, segments, rep.true_time, mode=mode, stop_share=stop)
                    row[f"stop_{stop}"] = {
                        "production": order_outcome(rep, prod_stop, vargas, mode),
                        "segment_order": order_outcome(rep, seg_stop, vargas, mode),
                    }
                rows.append(row)

            def agg(getter) -> dict[str, Any]:
                items = [getter(r) for r in rows]
                return {
                    "n": len(items),
                    "mean_questions": round(sum(i["questions"] for i in items) / len(items), 2) if items else None,
                    "truth_retained": sum(1 for i in items if i["truth_retained"]),
                    "top_is_truth": sum(1 for i in items if i["top_is_truth"]),
                    "alive_le2": sum(1 for i in items if i["alive_segments"] <= 2),
                    "top_share_median": _median([i["top_share"] for i in items]),
                    "truth_in_range": sum(1 for i in items if i["truth_in_range"]),
                    "width_median": _median([i["width"] for i in items]),
                }

            summary: dict[str, Any] = {
                "six": {"production": agg(lambda r: r["production"]), "segment_order": agg(lambda r: r["segment_order"])},
                "same_first_six_cases": sum(1 for r in rows if r["same_first_six"]),
                "pool_median": _median([r["pool"] for r in rows]),
            }
            for stop in STOP_SHARES:
                summary[f"stop_{stop}"] = {
                    "production": agg(lambda r, s=stop: r[f"stop_{s}"]["production"]),
                    "segment_order": agg(lambda r, s=stop: r[f"stop_{s}"]["segment_order"]),
                }
            summary["per_case"] = rows
            block[label] = summary
        out["by_radius"][str(radius)] = block
    return out


# ---------------------------------------------------------------------------
# M3
# ---------------------------------------------------------------------------

def stage_m3(replays: dict[tuple[str, int], vr.CaseReplay], signs60: dict[str, dict[int, dict[str, int]]],
             radii: Sequence[int], mode: str = "raw") -> dict[str, Any]:
    out: dict[str, Any] = {"mode": mode, "strategies": {k: list(v) for k, v in vr.STRATEGIES.items()}, "by_radius": {}}
    for radius in list(radii) + [r for r in vr.COUNT_RADII if r not in radii]:
        block: dict[str, Any] = {}
        for name, vargas in vr.STRATEGIES.items():
            # "no rectification needed": the target vargas keep one joint sign across the whole window
            no_need = sum(1 for signs in signs60.values() if vr.distinct_signs(signs, radius, vargas) == 1)
            entry: dict[str, Any] = {"no_rectification_needed": no_need, "n_cases": len(signs60)}
            if radius in radii:
                rows = []
                for (case_id, r), rep in sorted(replays.items()):
                    if r != radius:
                        continue
                    segments = rep.segments(vargas)
                    prod = vr.replay_scores(rep.public, rep.probes, rep.true_time)
                    prod["asked"] = list(rep.probes)[:vr.ASK_COUNT]
                    seg = vr.select_probes_by_segment(rep.public, rep.probes, segments, rep.true_time, mode=mode)
                    row = {
                        "case_id": case_id,
                        "segments_in_window": len(segments),
                        "production": order_outcome(rep, prod, vargas, mode),
                        "segment_order": order_outcome(rep, seg, vargas, mode),
                    }
                    for stop in STOP_SHARES:
                        seg_stop = vr.select_probes_by_segment(rep.public, rep.probes, segments, rep.true_time, mode=mode, stop_share=stop)
                        row[f"stop_{stop}"] = order_outcome(rep, seg_stop, vargas, mode)
                    rows.append(row)

                def agg(items: Sequence[dict[str, Any]]) -> dict[str, Any]:
                    return {
                        "n": len(items),
                        "mean_questions": round(sum(i["questions"] for i in items) / len(items), 2) if items else None,
                        "truth_retained": sum(1 for i in items if i["truth_retained"]),
                        "top_is_truth": sum(1 for i in items if i["top_is_truth"]),
                        "alive_le2": sum(1 for i in items if i["alive_segments"] <= 2),
                        "alive_single": sum(1 for i in items if i["alive_segments"] == 1),
                        "top_share_median": _median([i["top_share"] for i in items]),
                    }

                entry.update({
                    "segments_in_window_mean": round(sum(r["segments_in_window"] for r in rows) / len(rows), 2) if rows else None,
                    "production_six": agg([r["production"] for r in rows]),
                    "segment_order_six": agg([r["segment_order"] for r in rows]),
                    **{f"segment_order_stop_{s}": agg([r[f"stop_{s}"] for r in rows]) for s in STOP_SHARES},
                    "per_case": rows,
                })
            block[name] = entry
        out["by_radius"][str(radius)] = block
    return out


# ---------------------------------------------------------------------------
# robustness
# ---------------------------------------------------------------------------

def stage_robust(replays: dict[tuple[str, int], vr.CaseReplay], radii: Sequence[int], lmt_ids: set[str], mode: str = "raw") -> dict[str, Any]:
    targets = {"D9": ("D9",), "D10": ("D10",)}
    out: dict[str, Any] = {"mode": mode, "by_radius": {}}
    for radius in radii:
        block: dict[str, Any] = {}
        reps = [(cid, rep) for (cid, r), rep in sorted(replays.items()) if r == radius]
        for label, vargas in targets.items():
            entry: dict[str, Any] = {}
            base = []
            for cid, rep in reps:
                state = vr.replay_scores(rep.public, rep.probes, rep.true_time)
                base.append((cid, vr.segment_shares(vr.minute_weights(state, mode), rep.segments(vargas), rep.true_time)))
            entry["baseline"] = {
                "n": len(base),
                "truth_retained": sum(1 for _c, s in base if s["truth_retained"]),
                "top_is_truth": sum(1 for _c, s in base if s["top_is_truth"]),
                "lmt_era": {
                    "n": sum(1 for c, _s in base if c in lmt_ids),
                    "truth_retained": sum(1 for c, s in base if c in lmt_ids and s["truth_retained"]),
                    "top_is_truth": sum(1 for c, s in base if c in lmt_ids and s["top_is_truth"]),
                },
            }
            for flips in (1, 2):
                retained: list[bool] = []
                correct: list[bool] = []
                for seed in range(vr.FLIP_SEEDS):
                    for cid, rep in reps:
                        answers = vr.flipped_answers(rep.probes, rep.true_time, flips=flips, seed=f"{cid}:{radius}:{flips}:{seed}")
                        state = vr.replay_scores(rep.public, rep.probes, rep.true_time, answers=answers)
                        share = vr.segment_shares(vr.minute_weights(state, mode), rep.segments(vargas), rep.true_time)
                        retained.append(share["truth_retained"])
                        correct.append(share["top_is_truth"])
                entry[f"flip_{flips}"] = {
                    "replays": len(retained),
                    "truth_retained_rate": _rate(retained),
                    "top_is_truth_rate": _rate(correct),
                }
            shifted_retained: list[bool] = []
            shifted_correct: list[bool] = []
            for cid, rep in reps:
                rows, probes, public = rep.shifted(vr.DAY_SHIFT, seed=f"{cid}:{radius}:shift")
                state = vr.replay_scores(public, probes, rep.true_time)
                share = vr.segment_shares(vr.minute_weights(state, mode), rep.segments(vargas), rep.true_time)
                shifted_retained.append(share["truth_retained"])
                shifted_correct.append(share["top_is_truth"])
            entry["shift_7_days"] = {
                "n": len(shifted_retained),
                "truth_retained": sum(shifted_retained),
                "top_is_truth": sum(shifted_correct),
            }
            block[label] = entry
        out["by_radius"][str(radius)] = block
    return out


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------

def build_replays(cases: Sequence[dict[str, Any]], radii: Sequence[int], cache_dir: Path | None, log) -> dict[tuple[str, int], vr.CaseReplay]:
    replays: dict[tuple[str, int], vr.CaseReplay] = {}
    for case in cases:
        for radius in radii:
            started = time.perf_counter()
            replays[(str(case["case_id"]), radius)] = vr.CaseReplay(case, radius, cache_dir=cache_dir)
            log(f"{case['case_id']} ±{radius} ready in {time.perf_counter() - started:.1f}s")
    return replays


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", default=str(vr.HOLDOUT_V5))
    parser.add_argument("--stages", default="m0,m1,m2,m3,robust")
    parser.add_argument("--radii", default=",".join(str(r) for r in vr.RADII))
    parser.add_argument("--vargas", default=",".join(vr.VARGAS))
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--cache-dir", default="")
    parser.add_argument("--baseline-out", default=str(BASELINE_JSON))
    parser.add_argument("--json-out", default=str(REPORT_JSON))
    parser.add_argument("--quiet", action="store_true")
    args = parser.parse_args()
    stages = {s.strip() for s in args.stages.split(",") if s.strip()}
    radii = tuple(int(r) for r in args.radii.split(",") if r.strip())
    cache_dir = Path(args.cache_dir) if args.cache_dir else None
    log = (lambda *_a, **_k: None) if args.quiet else (lambda msg: print(msg, flush=True))

    started = time.perf_counter()
    cases = vr.load_cases(Path(args.dataset))
    if args.limit:
        cases = cases[: args.limit]
    lmt_ids = {str(c["case_id"]) for c in cases if vr.is_lmt_era(c)}
    signs60 = {str(c["case_id"]): vr.signs_per_minute(c, max(vr.COUNT_RADII), cache_dir=cache_dir) for c in cases}
    log(f"signs ready for {len(signs60)} cases")
    replays = build_replays(cases, radii, cache_dir, log)

    meta = {
        "dataset": Path(args.dataset).name,
        "case_count": len(cases),
        "radii": list(radii),
        "vargas": list(vr.VARGAS),
        "ayanamsa": vr.AYANAMSA,
        "node_mode": vr.NODE_MODE,
        "minute_step_candidates": vr.MINUTE_STEP,
        "minute_step_signs": 1,
        "ask_count": vr.ASK_COUNT,
        "separation_lead": vr.SEPARATION_LEAD,
        "lmt_era_before_year": vr.LMT_ERA_BEFORE_YEAR,
        "open_set_not_blind": True,
    }
    if "m0" in stages:
        baseline = {**meta, **stage_m0(replays, signs60, radii)}
        Path(args.baseline_out).write_text(vr.stable_json(baseline), encoding="utf-8")
        log(f"m0 written to {args.baseline_out}")
    report: dict[str, Any] = {**meta, "stages": sorted(stages - {"m0"})}
    if "m1" in stages:
        report["m1"] = stage_m1(replays, radii, lmt_ids)
        log("m1 done")
    if "m2" in stages:
        report["m2"] = stage_m2(replays, radii)
        log("m2 done")
    if "m3" in stages:
        report["m3"] = stage_m3(replays, signs60, radii)
        log("m3 done")
    if "robust" in stages:
        report["robust"] = stage_robust(replays, radii, lmt_ids)
        log("robust done")
    if stages - {"m0"}:
        # No wall-clock value goes into the JSON so two runs stay byte-identical.
        Path(args.json_out).write_text(vr.stable_json(report), encoding="utf-8")
        log(f"report written to {args.json_out}")
    log(f"elapsed {time.perf_counter() - started:.0f}s")
    return 0


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