"""Angle-timing research runner (BUG-1141).

A1: for each technique × orb, compare where the truth minute ranks inside its
window when events keep their real dates versus shuffled dates. Pre-registered
tests: TR / SP / SA × orbs 1° / 2° / 3° (9 tests), primary radius ±30, with
±10 / ±60 reported. Score per candidate = number of events with a contact within
the orb; truth rank = share of other window minutes scoring higher (ties half);
0.5 = chance. Real mean rank is compared with the shuffled-date distribution.

A2 (only for techniques significant in A1): see run_a2.

Usage:
  PYTHONHASHSEED=0 python3 -m scripts.research.angle_timing_research --stages a1 \
      --shuffles 200 --json-out docs/research/angle_timing_a1_2026_10_01.json
"""
from __future__ import annotations

import argparse
import platform
import sys
import time
from pathlib import Path
from random import Random
from typing import Any

import numpy as np

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

from scripts.research import angle_timing_lib as at  # noqa: E402

RADII = (10, 30, 60)
PRIMARY_RADIUS = 30
N_TESTS = len(at.TECHNIQUES) * len(at.ORBS)
BONFERRONI_PERCENTILE = 100.0 * (1.0 - 0.05 / N_TESTS)
UNCORRECTED_PERCENTILE = 95.0


def ranks_for(chart: at.CaseChart, raw_events: list[dict[str, Any]] | None) -> dict[str, dict[str, dict[str, float]]]:
    events = chart.events(raw_events)
    out: dict[str, dict[str, dict[str, float]]] = {}
    for tech in at.TECHNIQUES:
        seps = chart.separations(tech, events)
        out[tech] = {}
        for orb in at.ORBS:
            scores = at.hit_counts(seps, orb) if seps.shape[0] else np.zeros(len(chart.offsets))
            out[tech][str(orb)] = {str(r): at.truth_rank(scores, r) for r in RADII}
    return out


def run_a1(cases: list[dict[str, Any]], shuffles: int, log) -> dict[str, Any]:
    charts = [at.CaseChart(case) for case in cases]
    real = {chart.case_id: ranks_for(chart, None) for chart in charts}
    usable = {chart.case_id: {t: int(chart.separations(t, chart.events()).shape[0]) for t in at.TECHNIQUES} for chart in charts}
    shuffled_means: dict[str, dict[str, dict[str, list[float]]]] = {
        t: {str(o): {str(r): [] for r in RADII} for o in at.ORBS} for t in at.TECHNIQUES}
    started = time.perf_counter()
    for rep in range(shuffles):
        per_case = []
        for chart in charts:
            rng = Random(f"angle-timing|{chart.case_id}|{rep}")
            per_case.append(ranks_for(chart, at.shuffled_events(chart.case, rng)))
        for t in at.TECHNIQUES:
            for o in at.ORBS:
                for r in RADII:
                    shuffled_means[t][str(o)][str(r)].append(round(float(np.mean([pc[t][str(o)][str(r)] for pc in per_case])), 6))
        if (rep + 1) % 20 == 0:
            log(f"shuffle {rep + 1}/{shuffles} ({time.perf_counter() - started:.0f}s)")
    table: dict[str, Any] = {}
    for t in at.TECHNIQUES:
        table[t] = {}
        for o in at.ORBS:
            table[t][str(o)] = {}
            for r in RADII:
                real_mean = round(float(np.mean([real[cid][t][str(o)][str(r)] for cid in real])), 6)
                dist = np.array(shuffled_means[t][str(o)][str(r)])
                # percentile: share of shuffles the real dates beat (lower rank is better)
                pct = round(100.0 * float(((dist > real_mean).sum() + 0.5 * (dist == real_mean).sum()) / len(dist)), 3)
                table[t][str(o)][str(r)] = {
                    "real_mean_rank": real_mean,
                    "shuffle_mean": round(float(dist.mean()), 6),
                    "shuffle_sd": round(float(dist.std(ddof=1)), 6) if len(dist) > 1 else 0.0,
                    "shuffle_p05": round(float(np.percentile(dist, 5)), 6),
                    "percentile_beaten": pct,
                    "significant_uncorrected": bool(pct >= UNCORRECTED_PERCENTILE),
                    "significant_bonferroni": bool(pct >= BONFERRONI_PERCENTILE),
                }
    primary = {t: {str(o): table[t][str(o)][str(PRIMARY_RADIUS)] for o in at.ORBS} for t in at.TECHNIQUES}
    significant = sorted(f"{t}@{o}" for t in at.TECHNIQUES for o in at.ORBS
                         if primary[t][str(o)]["significant_bonferroni"])
    return {
        "tests": N_TESTS,
        "primary_radius": PRIMARY_RADIUS,
        "bonferroni_percentile": round(BONFERRONI_PERCENTILE, 4),
        "uncorrected_percentile": UNCORRECTED_PERCENTILE,
        "shuffles": shuffles,
        "table": table,
        "significant_primary_bonferroni": significant,
        "usable_events_per_case": usable,
        "real_per_case": real,
        "shuffled_means": shuffled_means,
    }


def _precision_filter(raw_events: list[dict[str, Any]], precision: str) -> list[dict[str, Any]]:
    return [e for e in raw_events if str(e.get("precision")) == precision]


def run_jitter_null(cases: list[dict[str, Any]], shuffles: int, log) -> dict[str, Any]:
    """Second null (age-structure preserving): real ranks vs events moved ±1-3 years, all radii."""
    charts = [at.CaseChart(case) for case in cases]
    real = [ranks_for(chart, None) for chart in charts]
    dist: dict[str, dict[str, dict[str, list[float]]]] = {t: {str(o): {str(r): [] for r in RADII} for o in at.ORBS} for t in at.TECHNIQUES}
    for rep in range(shuffles):
        per = [ranks_for(c, at.jittered_events(c.case, Random(f"angle-timing-jitter|{c.case_id}|{rep}"))) for c in charts]
        for t in at.TECHNIQUES:
            for o in at.ORBS:
                for r in RADII:
                    dist[t][str(o)][str(r)].append(round(float(np.mean([pc[t][str(o)][str(r)] for pc in per])), 6))
        if (rep + 1) % 10 == 0:
            log(f"jitter {rep + 1}/{shuffles}")
    table: dict[str, Any] = {}
    for t in at.TECHNIQUES:
        table[t] = {}
        for o in at.ORBS:
            table[t][str(o)] = {}
            for r in RADII:
                real_mean = round(float(np.mean([x[t][str(o)][str(r)] for x in real])), 6)
                d = np.array(dist[t][str(o)][str(r)])
                pct = round(100.0 * float(((d > real_mean).sum() + 0.5 * (d == real_mean).sum()) / len(d)), 3)
                table[t][str(o)][str(r)] = {"real_mean_rank": real_mean, "jitter_mean": round(float(d.mean()), 6),
                                            "percentile_beaten": pct}
    return {"shuffles": shuffles, "null": "each event moved 1-3 years earlier/later, precision kept", "table": table}


def run_subsets(cases: list[dict[str, Any]], shuffles: int, log) -> dict[str, Any]:
    """Robustness: real vs shuffled by event precision and by LMT era (primary radius only)."""
    charts = [at.CaseChart(case) for case in cases]
    lmt = {chart.case_id for chart in charts if chart.local.year < 1900}
    subsets: dict[str, tuple[list[at.CaseChart], str | None]] = {
        "precision=day": (charts, "day"), "precision=month": (charts, "month"), "precision=year": (charts, "year"),
        "era=lmt_before_1900": ([c for c in charts if c.case_id in lmt], None),
        "era=1900_and_later": ([c for c in charts if c.case_id not in lmt], None),
    }
    out: dict[str, Any] = {}
    for name, (members, precision) in subsets.items():
        def ranks(chart: at.CaseChart, raw: list[dict[str, Any]]):
            events = _precision_filter(raw, precision) if precision else raw
            return ranks_for(chart, events)
        real = [ranks(c, list(c.case["events"])) for c in members]
        dist: dict[str, dict[str, list[float]]] = {t: {str(o): [] for o in at.ORBS} for t in at.TECHNIQUES}
        for rep in range(shuffles):
            per = [ranks(c, at.shuffled_events(c.case, Random(f"angle-timing-sub|{name}|{c.case_id}|{rep}"))) for c in members]
            for t in at.TECHNIQUES:
                for o in at.ORBS:
                    dist[t][str(o)].append(round(float(np.mean([pc[t][str(o)][str(PRIMARY_RADIUS)] for pc in per])), 6))
        out[name] = {"cases": len(members), "table": {}}
        for t in at.TECHNIQUES:
            out[name]["table"][t] = {}
            for o in at.ORBS:
                real_mean = round(float(np.mean([r[t][str(o)][str(PRIMARY_RADIUS)] for r in real])), 6)
                d = np.array(dist[t][str(o)])
                pct = round(100.0 * float(((d > real_mean).sum() + 0.5 * (d == real_mean).sum()) / len(d)), 3)
                out[name]["table"][t][str(o)] = {"real_mean_rank": real_mean, "shuffle_mean": round(float(d.mean()), 6),
                                                 "percentile_beaten": pct}
        log(f"subset {name} done")
    return {"radius": PRIMARY_RADIUS, "shuffles": shuffles, "subsets": out,
            "note": "descriptive robustness; significance is judged only on the pre-registered A1 tests"}


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--stages", default="a1")
    parser.add_argument("--shuffles", type=int, default=200)
    parser.add_argument("--subset-shuffles", type=int, default=50)
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--json-out", required=True)
    parser.add_argument("--quiet", action="store_true")
    args = parser.parse_args()
    log = (lambda _m: None) if args.quiet else (lambda m: print(m, file=sys.stderr, flush=True))
    cases = at.load_cases()
    if args.limit:
        cases = cases[: args.limit]
    payload: dict[str, Any] = {
        "dataset": "v5",
        "holdout_sha256": at.sha256_of(at.HOLDOUT_V5),
        "case_count": len(cases),
        "zodiac": "tropical (aspects are zodiac-independent)",
        "node": "mean",
        "techniques": {
            "TR": "transit Saturn / Jupiter / mean Rahu-Ketu (+Mars for day events) to natal ASC / MC, 0/90/180",
            "SP": "secondary-progressed ASC / MC (quotidian date method, same location) to natal Sun..Saturn, 0/90/180",
            "SA": "solar-arc ASC / MC to natal Sun..Saturn, 0/90/180",
        },
        "orbs": list(at.ORBS),
        "radii": list(RADII),
        "event_cutoff": at.EVENT_CUTOFF.isoformat(),
        "python_version": platform.python_version(),
    }
    stages = {s.strip() for s in args.stages.split(",") if s.strip()}
    if "a1" in stages:
        payload["a1"] = run_a1(cases, args.shuffles, log)
    if "jitter" in stages:
        payload["jitter_null"] = run_jitter_null(cases, args.subset_shuffles, log)
    if "subsets" in stages:
        payload["subsets"] = run_subsets(cases, args.subset_shuffles, log)
    out = Path(args.json_out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(at.stable_json(payload), encoding="utf-8")
    log(f"wrote {out}")
    return 0


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