#!/usr/bin/env python
import argparse
import json
import itertools
from pathlib import Path

import numpy as np
from scipy.stats import wilcoxon

HERE = Path(__file__).resolve().parent
_ap = argparse.ArgumentParser()
_ap.add_argument(
    "--pattern",
    default="trajectory_cv_tilt_k*_seed*.json",
    help="glob over data/p53 selecting the K/seed runs to pool; the default "
    "matches the original site_drug sweep",
)
_ap.add_argument(
    "--out", type=Path, default=HERE / "data/p53/trajectory_cv_tilt_robustness.json"
)
_args = _ap.parse_args()
files = sorted((HERE / "data/p53").glob(_args.pattern))
RNG_SEED = 20260729


def mean_bootstrap_ci(x, n_boot=200_000):
    """Trajectory-fold bootstrap interval for the arithmetic mean."""
    x = np.asarray(x, float)
    rng = np.random.default_rng(RNG_SEED)
    idx = rng.integers(0, len(x), size=(n_boot, len(x)))
    return np.percentile(x[idx].mean(axis=1), [2.5, 97.5]).tolist()


def exact_signflip_p(x, null=1.0, alternative="two-sided"):
    """Exact paired randomization test aligned to a mean-difference estimand."""
    delta = np.asarray(x, float) - null
    observed = delta.mean()
    signs = np.asarray(list(itertools.product((-1.0, 1.0), repeat=len(delta))))
    stats = (signs * delta).mean(axis=1)
    if alternative == "greater":
        return float(np.mean(stats >= observed - 1e-15))
    return float(np.mean(np.abs(stats) >= abs(observed) - 1e-15))


print("TRAJECTORY-HELD-OUT FK ROBUSTNESS")
print(f"{'file':<42}{'mean':>8}{'median':>9}{'>1':>7}{'p':>9}")
all_means = []
for path in files:
    d = json.loads(path.read_text())
    x = d["summary"]["fk_doob"]
    mean = x["mean_ratio_to_equilibrium"]
    all_means.append(mean)
    print(
        f"{path.name:<42}{mean:>8.3f}"
        f"{x['median_ratio_to_equilibrium']:>9.3f}"
        f"{x['folds_above_1']:>3}/{x['n_folds']:<3}"
        f"{x['wilcoxon_p_vs_1']:>9.3g}"
    )

if all_means:
    print(
        f"\nAcross completed discretizations: mean={np.mean(all_means):.3f}, "
        f"range=[{np.min(all_means):.3f}, {np.max(all_means):.3f}]"
    )
else:
    print("\nNo completed robustness outputs.")

if files:
    _base = ("equilibrium", "uniform_states", "tpt_flux", "doob_untilted", "fk_doob")
    _per_file = [set(json.loads(p.read_text())["folds"][0]["selectors"]) for p in files]
    _common = set.intersection(*_per_file) if _per_file else set()
    _alg1 = sorted(
        (s for s in _common if s.startswith("alg1[a=")),
        key=lambda s: float(s.split("=")[1].rstrip("]")),
    )
    selectors = tuple(s for s in _base if s in _common) + tuple(_alg1)
    dropped = sorted((set().union(*_per_file) - _common) if _per_file else [])
    if dropped:
        print(f"[note] selectors absent from some runs, excluded: {dropped}")
    raw = {}
    for selector in selectors:
        raw[selector] = np.asarray(
            [
                [
                    fold["selectors"][selector]["ratio_to_equilibrium"]
                    for fold in json.loads(path.read_text())["folds"]
                ]
                for path in files
            ]
        )

    collapsed = {k: v.mean(axis=0) for k, v in raw.items()}
    print("\nINFERENCE AFTER COLLAPSING K/SEED REPEATS WITHIN HELD-OUT FOLD")
    print(f"{'selector':<18}{'mean':>8}{'median':>9}{'>1':>7}{'p vs 1':>10}")
    aggregate = {}
    for selector in selectors:
        x = collapsed[selector]
        p = (
            1.0
            if selector == "equilibrium"
            else float(wilcoxon(x - 1, alternative="two-sided").pvalue)
        )
        print(
            f"{selector:<18}{x.mean():>8.3f}{np.median(x):>9.3f}"
            f"{int((x > 1).sum()):>3}/{len(x):<3}{p:>10.3g}"
        )
        aggregate[selector] = {
            "mean": float(x.mean()),
            "mean_fold_bootstrap_95ci": (
                [1.0, 1.0] if selector == "equilibrium" else mean_bootstrap_ci(x)
            ),
            "median": float(np.median(x)),
            "n_positive": int((x > 1).sum()),
            "n_folds": int(len(x)),
            "wilcoxon_p_vs_1": p,
            "exact_mean_signflip_p_two_sided": (
                1.0
                if selector == "equilibrium"
                else exact_signflip_p(x, null=1.0, alternative="two-sided")
            ),
            "exact_mean_signflip_p_one_sided_greater": (
                1.0
                if selector == "equilibrium"
                else exact_signflip_p(x, null=1.0, alternative="greater")
            ),
            "fold_values": x.tolist(),
        }

    if _alg1:
        print(
            "\nALGORITHM 1 TILT-STRENGTH SWEEP  (w=exp(-alpha_rho*rho), no committor shaping)"
        )
        print(
            f"{'alpha_rho':>10}{'mean':>8}{'median':>9}{'>1':>7}{'95% CI':>18}{'p vs 1':>10}"
        )
        for s in _alg1:
            a = aggregate[s]
            ci = a["mean_fold_bootstrap_95ci"]
            star = " *" if (ci[0] > 1 or ci[1] < 1) else ""
            print(
                f"{float(s.split('=')[1].rstrip(']')):>10.2f}{a['mean']:>8.3f}"
                f"{a['median']:>9.3f}{a['n_positive']:>3}/{a['n_folds']:<3}"
                f"{f'[{ci[0]:.3f}, {ci[1]:.3f}]':>18}{a['wilcoxon_p_vs_1']:>10.3g}{star}"
            )
        print("  * = 95% CI excludes 1.   alpha_rho=0 is the no-density-tilt control.")

    print("\nPAIRED FK ADVANTAGE")
    fk = collapsed["fk_doob"]
    paired = {}
    for comparator in ("tpt_flux", "uniform_states", "doob_untilted"):
        delta = fk - collapsed[comparator]
        p = float(wilcoxon(delta, alternative="two-sided").pvalue)
        print(
            f"FK - {comparator:<14} mean={delta.mean():+.3f}; "
            f"wins={int((delta > 0).sum())}/{len(delta)}; p={p:.3g}"
        )
        paired[comparator] = {
            "mean_difference": float(delta.mean()),
            "mean_difference_fold_bootstrap_95ci": mean_bootstrap_ci(delta),
            "fk_wins": int((delta > 0).sum()),
            "n_folds": int(len(delta)),
            "wilcoxon_p": p,
            "exact_mean_signflip_p_two_sided": exact_signflip_p(
                delta, null=0.0, alternative="two-sided"
            ),
            "exact_mean_signflip_p_one_sided_greater": exact_signflip_p(
                delta, null=0.0, alternative="greater"
            ),
        }

    out = {
        "n_configurations": len(files),
        "configuration_mean_range": [
            float(np.min(all_means)),
            float(np.max(all_means)),
        ],
        "aggregate": aggregate,
        "paired_fk_advantage": paired,
        "note": "K/seed configurations are averaged within fold before inference.",
    }
    out_path = _args.out
    out["source_pattern"] = _args.pattern
    out["source_files"] = [p.name for p in files]
    out["observable"] = sorted(
        {json.loads(p.read_text())["config"]["observable"] for p in files}
    )
    out_path.write_text(json.dumps(out, indent=2))
    print(f"\n-> {out_path}")