"""Evaluate the pre-registered hypotheses (prereg.md section 3) and build summary tables + Space data.

Usage: uv run python scripts/report.py [results_dir]   (after run_all.py and latency.py)
"""
import itertools
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

from gatestats import stats

ROOT = Path(__file__).resolve().parents[1]
RES = Path(sys.argv[1]) if len(sys.argv) > 1 else ROOT / "results"
PRIMARY = dict(truth="population", metric="recall@5")


def pick(df, **kw):
    for k, v in kw.items():
        df = df[df[k].isin(v) if isinstance(v, (list, tuple)) else np.isclose(df[k], v) if isinstance(v, float) else df[k] == v]
    return df


def verdict(rows, need):
    return {"supported": sum(r["ok"] for r in rows) >= need, "need": need, "per_unit": rows}


def h1(thr, sel=PRIMARY):
    rows = []
    for ds, g in pick(thr, **sel, n=100, d=[-0.02, 0.02]).groupby("dataset"):
        wrong = int(np.where(g.truth_pass, g.R - g.k_pass, g.k_pass).sum())
        lo, hi = stats.wilson_ci(wrong, g.R.sum())
        rows.append(dict(dataset=ds, error=wrong / g.R.sum(), lo=float(lo), hi=float(hi), ok=bool(lo >= 0.15)))
    return verdict(rows, 2)


def h2(pairs, con, sel=PRIMARY):
    rows = []
    for ds in sorted(pairs.dataset.unique()):
        g1 = pick(pairs, **sel, dataset=ds, B="ctl03", gate="G1", param=0.05, n=100).iloc[0]
        c = pick(con, **sel, dataset=ds, B="ctl03", contrast="G1pass-G3pass", param=0.02, n=100).iloc[0]
        rate = g1.k_pass / g1.R
        rows.append(dict(dataset=ds, g1_pass=rate, diff=c["mean"], lo=c.lo, hi=c.hi, ok=bool(rate >= 0.5 and c.lo > 0)))
    return verdict(rows, 2)


def h3(con, sel=PRIMARY):
    rows = []
    for ds in sorted(con.dataset.unique()):
        c = pick(con, **sel, dataset=ds, A="*", param=0.02, n=[50, 100])
        rows.append(dict(dataset=ds, diffs={int(r.n): [r["mean"], r.lo, r.hi] for _, r in c.iterrows()},
                         ok=bool(len(c) == 2 and (c.lo > 0).all())))
    return verdict(rows, 2)


def hnsw_table(pq):
    h = pq[pq.system.str.startswith("hnsw")].copy()
    h["ef"] = h.system.str.extract(r"ef(\d+)").astype(int)
    h["seed"] = h.system.str.extract(r"_s(\d+)").astype(int)
    means = h.groupby(["dataset", "ef", "seed"])["recall@10"].mean().reset_index()
    rows = []
    for (ds, ef), g in means.groupby(["dataset", "ef"]):
        v = g.sort_values("seed")["recall@10"].to_numpy()
        diffs = [abs(x - y) for x, y in itertools.combinations(v, 2)]
        q = h[(h.dataset == ds) & (h.ef == ef)].pivot(index="qid", columns="seed", values="recall@10")
        rows.append(dict(dataset=ds, ef=int(ef), mean=float(v.mean()), sd=float(v.std(ddof=1)),
                         p95_abs_diff=float(np.percentile(diffs, 95)), max_abs_diff=float(max(diffs)),
                         share_queries_changed=float((q.ne(q[0], axis=0)).any(axis=1).mean())))
    return pd.DataFrame(rows)


def h4(tab):
    rows = []
    for ds, g in tab.groupby("dataset"):
        p = g.sort_values("ef").p95_abs_diff.to_numpy()
        rows.append(dict(dataset=ds, p95_by_ef=p.round(5).tolist(), positive_at_16=bool(p[0] > 0),
                         non_increasing=bool(np.all(np.diff(p) <= 1e-12)), ok=bool(np.all(np.diff(p) <= 1e-12))))
    out = verdict(rows, 2)
    out["supported"] = out["supported"] and all(r["positive_at_16"] for r in rows)
    return out


def latency_table(lat):
    rows = []
    for (system, run), g in lat.groupby(["system", "run"]):
        x = g.ms.to_numpy()
        p50, lo50, hi50 = stats.quantile_bootstrap_ci(x, 0.5, seed=run)
        p95, lo95, hi95 = stats.quantile_bootstrap_ci(x, 0.95, seed=run)
        rows.append(dict(system=system, run=run, p50=p50, p95=p95, p99=float(np.quantile(x, 0.99)),
                         rel_w50=(hi50 - lo50) / p50, rel_w95=(hi95 - lo95) / p95))
    return pd.DataFrame(rows)


def h5(lt):
    rows = []
    for system, g in lt.groupby("system"):
        ratio = float((g.rel_w95 / g.rel_w50).median())
        cv50, cv95 = g.p50.std(ddof=1) / g.p50.mean(), g.p95.std(ddof=1) / g.p95.mean()
        rows.append(dict(system=system, width_ratio_median=ratio, cv_p50=float(cv50), cv_p95=float(cv95),
                         ok=bool(ratio >= 2 and cv95 > cv50)))
    return verdict(rows, 4)


def latency_gates(lat, hs=(-0.2, -0.1, -0.05, 0.05, 0.1, 0.2), ns=(50, 100, 200), R=200, B=500, alpha=0.05):
    """Point-estimate P95 gate vs GL (bootstrap upper bound) against the pooled P95 of all runs."""
    rng, rows = np.random.default_rng(20261017), []
    for system, g in lat.groupby("system"):
        true = np.quantile(g.ms, 0.95)
        for n in ns:
            point, gl = {h: [] for h in hs}, {h: [] for h in hs}
            for _, gr in g.groupby("run"):
                S = rng.permuted(np.tile(gr.ms.to_numpy(), (R, 1)), axis=1)[:, :n]
                est = np.quantile(S, 0.95, axis=1)
                idx = rng.integers(0, n, size=(B, n))
                ub = np.array([stats.upper_bound(np.quantile(s[idx], 0.95, axis=1), alpha) for s in S])
                for h in hs:
                    point[h].append((est <= true * (1 + h)).mean())
                    gl[h].append((ub <= true * (1 + h)).mean())
            for h in hs:
                rows.append(dict(system=system, n=n, sla_rel=h, truth_pass=h > 0,
                                 point_pass=float(np.mean(point[h])), gl_pass=float(np.mean(gl[h]))))
    return pd.DataFrame(rows)


def mde_table(pairs):
    g = pick(pairs, **PRIMARY, gate="G3", param=0.02, n=100)[["dataset", "family", "A", "B", "delta_true", "sd_d"]]
    for n in (50, 100, 200, 300, 500, 1000):
        g[f"mde_n{n}"] = stats.mde(g.sd_d, n)
    return g


def md(df, digits=3):
    return df.round(digits).to_markdown()


def article_tables(thr, pairs, pq, hn, env, secondary, lt=None, lg=None):
    """Markdown tables the article quotes; every number in the article comes from here."""
    from scipy.stats import norm

    out = []
    exact = pq[pq.system.isin(["bm25", "bge", "minilm", "e5"])]
    sd = exact.groupby("dataset")["recall@5"].std()
    ds = pd.DataFrame({"N": {k: v["queries"] for k, v in env["datasets"].items()},
                       "docs": {k: v["docs"] for k, v in env["datasets"].items()},
                       "bge_recall@5": pq[pq.system == "bge"].groupby("dataset")["recall@5"].mean(),
                       "sd_recall@5": sd})
    ds["se_n100"] = ds["sd_recall@5"] / 10 * np.sqrt(1 - 100 / ds.N)
    out += ["## A. Veri setleri", md(ds, 4)]

    t = pick(thr, truth="population", metric="recall@5")
    t = t[t.d.abs() > 1e-9].assign(abs_d=lambda x: x.d.abs().round(2))
    tab = t.groupby(["dataset", "n", "abs_d"]).error.mean().unstack("abs_d")
    out += ["## B. G1 hata oranı (Recall@5, popülasyon gerçeği, 4 sistem havuzlanmış)", md(tab)]
    pred = {d: norm.cdf(-d / ds.se_n100) for d in (0.01, 0.02, 0.03, 0.05)}
    comp = pd.DataFrame({f"obs_{d}": tab.xs(100, level="n")[d] for d in pred} | {f"pred_{d}": v for d, v in pred.items()})
    out += ["## B2. n=100: gözlenen ve normal yaklaşımla öngörülen G1 hata oranı", md(comp)]

    q = pick(pairs, truth="population", metric="recall@5", n=100)
    q = q[((q.gate == "G1") & np.isclose(q.param, 0.05)) | ((q.gate != "G1") & np.isclose(q.param, 0.02))].copy()
    q["scenario"] = q.B.str.replace(r"hnsw_ef(16|128)_s[1-9]", r"noop_ef\1", regex=True)
    q.loc[q.family == "natural", "scenario"] = q.A + " -> " + q.B
    g = q.groupby(["dataset", "scenario", "gate"]).agg(pass_rate=("pass_rate", "mean"), delta=("delta_true", "mean"))
    wide = g.pass_rate.unstack("gate").join(g.delta.groupby(level=[0, 1]).mean())
    out += ["## C. PASS oranları, n=100 (G1: eşik = baseline − 0.05; diğerleri m = 0.02)", md(wide)]

    g3 = pick(pairs, truth="population", metric="recall@5", gate="G3", param=0.02).copy()
    g3["scenario"] = g3.B.str.replace(r"hnsw_ef(16|128)_s[1-9]", r"noop_ef\1", regex=True)
    g3 = g3[g3.scenario.isin(["ctl01", "ctl02", "ctl03", "noop_ef16", "noop_ef128"])]
    out += ["## D. G3 PASS oranı golden set boyutuna göre (m = 0.02)",
            md(g3.groupby(["dataset", "scenario", "n"]).pass_rate.mean().unstack("n"))]
    gap = pick(pairs, truth="population", metric="recall@5", gate="G3").assign(
        gap=lambda x: (x.pass_rate - x.analytic_pass).abs())
    out += ["## D2. G3: simülasyon ile normal yaklaşım arasındaki mutlak fark",
            md(gap.groupby(["dataset", "n"]).gap.agg(["median", "max"]).unstack("n"))]

    out += ["## E. HNSW ekleme sırası (Recall@10, 10 seed)", md(hn.set_index(["dataset", "ef"]), 4)]
    if lt is not None:
        out += ["## F. Latency (FiQA, 2 iş parçacığı, koşu ortalaması; ms)",
                md(lt.groupby("system")[["p50", "p95", "p99", "rel_w50", "rel_w95"]].mean(), 3),
                "## F2. Koşular arası değişkenlik (CV)",
                md(lt.groupby("system").agg(cv_p50=("p50", lambda x: x.std(ddof=1) / x.mean()),
                                            cv_p95=("p95", lambda x: x.std(ddof=1) / x.mean())), 4),
                "## F3. Latency gate PASS oranı (SLA = havuzlanmış P95 × (1 + h))",
                md(lg.pivot_table(index=["system", "n"], columns="sla_rel", values=["point_pass", "gl_pass"]))]
    out += ["## G. İkincil analizler (destek kararı)", md(pd.DataFrame(secondary).T, 0)]
    return "\n\n".join(out) + "\n"


def main():
    pq = pd.read_parquet(RES / "per_query.parquet")
    thr = pd.read_parquet(RES / "threshold.parquet")
    pairs = pd.read_parquet(RES / "pairs.parquet")
    con = pd.read_parquet(RES / "contrasts.parquet")
    thr["error"] = np.where(thr.truth_pass, thr.R - thr.k_pass, thr.k_pass) / thr.R
    pairs["pass_rate"] = pairs.k_pass / pairs.R

    hn = hnsw_table(pq)
    hyp = {"H1": h1(thr), "H2": h2(pairs, con), "H3": h3(con), "H4": h4(hn)}
    secondary = {f"{t}/{m}": {k: f(sel=dict(truth=t, metric=m))["supported"] for k, f in
                                [("H1", lambda sel: h1(thr, sel)), ("H2", lambda sel: h2(pairs, con, sel)),
                                 ("H3", lambda sel: h3(con, sel))]}
                 for t, m in [("population", "ndcg@10"), ("split_half", "recall@5")]}
    space = {}
    if (RES / "latency.parquet").exists():
        lat = pd.read_parquet(RES / "latency.parquet")
        lt = latency_table(lat)
        hyp["H5"] = h5(lt)
        lg = latency_gates(lat)
        lt.to_csv(RES / "latency_summary.csv", index=False)
        lg.to_csv(RES / "latency_gates.csv", index=False)
        space["latency"] = lt.groupby("system")[["p50", "p95", "p99", "rel_w50", "rel_w95"]].mean().reset_index().to_dict("records")
        space["latency_gates"] = lg.to_dict("records")
    (RES / "hypotheses.json").write_text(json.dumps(hyp, indent=2, default=float))
    (RES / "secondary.json").write_text(json.dumps(secondary, indent=2))

    mt = mde_table(pairs)
    mt.to_csv(RES / "mde.csv", index=False)
    hn.to_csv(RES / "hnsw.csv", index=False)
    systems = pq.groupby(["dataset", "system"])[["recall@5", "recall@10", "mrr@10", "ndcg@10"]].mean().reset_index()
    systems.to_csv(RES / "system_means.csv", index=False)
    env = json.loads((RES / "env.json").read_text())
    lat_args = (lt, lg) if "H5" in hyp else ()
    (RES / "summary.md").write_text(article_tables(thr, pairs, pq, hn, env, secondary, *lat_args), encoding="utf-8")

    curves = thr.groupby(["dataset", "truth", "metric", "n", "d"]).apply(
        lambda g: np.where(g.truth_pass, g.R - g.k_pass, g.k_pass).sum() / g.R.sum(), include_groups=False)
    space.update({
        "hypotheses": {k: v["supported"] for k, v in hyp.items()},
        "g1_error": curves.rename("error").reset_index().to_dict("records"),
        "gates": pairs[(pairs.truth == "population") & ((pairs.gate == "G1") | np.isclose(pairs.param, 0.02))][
            ["dataset", "truth", "family", "A", "B", "metric", "n", "gate", "param", "delta_true", "sd_d", "pass_rate",
             "analytic_pass"]].replace({np.nan: None}).to_dict("records"),
        "mde": mt.to_dict("records"),
        "hnsw": hn.to_dict("records"),
        "systems": systems.to_dict("records"),
    })
    out = ROOT / "space" / "data"
    out.mkdir(parents=True, exist_ok=True)
    (out / "summary.json").write_text(json.dumps(space, default=float))
    for k, v in hyp.items():
        print(k, "SUPPORTED" if v["supported"] else "NOT SUPPORTED")


if __name__ == "__main__":
    main()