"""Registered meta-analysis (plan/analysis-plan.md): does depression shift time production in
opposite directions inside and beyond the experienced moment (<= 3 s versus > 3 s)?

Effect size: Hedges' g with small-sample correction J, signed so that positive = depressed
participants UNDERproduce (produce shorter intervals than controls; "acceleration") and
negative = OVERproduce ("dilation").

Inputs:
  data/extraction.csv   group means/SDs extracted from primary papers
  data/secondary.csv    effect sizes taken from Thones & Oberfeld (2015, Table A2) for studies
                        whose primary group means are not available, with the intervals they
                        cover (each row one study x zone)
Usage: python code/meta.py <out.json>
"""
import csv
import json
import math
import sys
from collections import defaultdict

import numpy as np
from scipy import stats

MOMENT_S = 3.0
SEED = 20261008


def pool_groups(rows):
    """Pool depressed subgroups sharing one control group (weighted mean; SD including
    between-subgroup variance)."""
    n = sum(r["n"] for r in rows)
    m = sum(r["n"] * r["m"] for r in rows) / n
    ss = sum((r["n"] - 1) * r["s"] ** 2 + r["n"] * (r["m"] - m) ** 2 for r in rows)
    return n, m, math.sqrt(ss / (n - 1))


def hedges(nd, md, sd, nc, mc, sc, higher_means):
    sp = math.sqrt(((nd - 1) * sd ** 2 + (nc - 1) * sc ** 2) / (nd + nc - 2))
    d = (md - mc) / sp
    if higher_means == "overproduction":  # seconds produced: higher = longer = overproduction
        d = -d
    j = 1 - 3 / (4 * (nd + nc) - 9)
    g = j * d
    v = (nd + nc) / (nd * nc) + g ** 2 / (2 * (nd + nc))
    return g, v


def zone(iv):
    return "moment" if iv <= MOMENT_S else "beyond"


def primary_effects(path):
    """One g per study x interval x occasion; then averaged per study x zone (Thones 2015 rule:
    mean g, variance of the mean computed from the average sample sizes)."""
    groups = defaultdict(list)
    with open(path, newline="") as fh:
        for r in csv.DictReader(fh):
            key = (r["study"], float(r["interval_s"]), r["occasion"], r["higher_value_means"],
                   int(r["n_ctl"]), float(r["mean_ctl"]), float(r["sd_ctl"]))
            groups[key].append({"n": int(r["n_dep"]), "m": float(r["mean_dep"]), "s": float(r["sd_dep"])})
    per = []
    for (study, iv, occ, hm, nc, mc, sc), g in groups.items():
        nd, md, sd = pool_groups(g)
        gg, vv = hedges(nd, md, sd, nc, mc, sc, hm)
        per.append({"study": study, "interval_s": iv, "occasion": occ, "g": gg, "v": vv, "nd": nd, "nc": nc})
    return per


def by_zone(per):
    out = defaultdict(list)
    for e in per:
        out[(e["study"], zone(e["interval_s"]))].append(e)
    rows = []
    for (study, z), es in out.items():
        g = float(np.mean([e["g"] for e in es]))
        nd = float(np.mean([e["nd"] for e in es]))
        nc = float(np.mean([e["nc"] for e in es]))
        v = (nd + nc) / (nd * nc) + g ** 2 / (2 * (nd + nc))
        rows.append({"study": study, "zone": z, "g": g, "v": v, "k_intervals": len(es),
                     "intervals_s": sorted({e["interval_s"] for e in es}), "source": "primary"})
    return rows


def secondary_effects(path):
    rows = []
    with open(path, newline="") as fh:
        for r in csv.DictReader(fh):
            rows.append({"study": r["study"], "zone": r["zone"], "g": float(r["g"]), "v": float(r["var_g"]),
                         "k_intervals": int(r["k_intervals"]), "intervals_s": r["intervals_s"],
                         "source": "Thones & Oberfeld 2015 Table A2"})
    return rows


def reml_tau2(y, v, X):
    def nll(t2):
        w = 1 / (v + t2)
        W = np.diag(w)
        XtWX = X.T @ W @ X
        b = np.linalg.solve(XtWX, X.T @ W @ y)
        r = y - X @ b
        return 0.5 * (np.sum(np.log(v + t2)) + np.log(np.linalg.det(XtWX)) + r @ W @ r)
    from scipy.optimize import minimize_scalar
    res = minimize_scalar(nll, bounds=(0, 10), method="bounded")
    return float(res.x)


def meta_regression(rows):
    """Random-effects meta-regression g ~ zone (indicator beyond), REML tau^2, with
    cluster-robust (study) standard errors, CR1 small-sample factor, t with (studies - 2) df."""
    y = np.array([r["g"] for r in rows])
    v = np.array([r["v"] for r in rows])
    beyond = np.array([1.0 if r["zone"] == "beyond" else 0.0 for r in rows])
    X = np.column_stack([1 - beyond, beyond])  # cell means: moment, beyond
    t2 = reml_tau2(y, v, X)
    w = 1 / (v + t2)
    W = np.diag(w)
    bread = np.linalg.inv(X.T @ W @ X)
    b = bread @ X.T @ W @ y
    resid = y - X @ b
    studies = sorted({r["study"] for r in rows})
    meat = np.zeros((2, 2))
    for s in studies:
        idx = [i for i, r in enumerate(rows) if r["study"] == s]
        u = X[idx].T @ (w[idx] * resid[idx])
        meat += np.outer(u, u)
    m = len(studies)
    V = bread @ meat @ bread * m / (m - 1)
    df = m - 2
    L = np.array([-1.0, 1.0])
    diff = float(L @ b)
    se_diff = float(np.sqrt(L @ V @ L))
    out = {"studies": m, "effects": len(rows), "tau2": round(t2, 4), "df": df}
    for k, name in enumerate(["moment", "beyond"]):
        se = float(np.sqrt(V[k, k]))
        t = float(b[k] / se)
        out[name] = {"g": round(float(b[k]), 4), "se": round(se, 4), "t": round(t, 3),
                     "ci95": [round(float(b[k] - stats.t.ppf(0.975, df) * se), 4), round(float(b[k] + stats.t.ppf(0.975, df) * se), 4)],
                     "p_two_sided": float(2 * stats.t.sf(abs(t), df)),
                     "k_effects": int(sum(1 for r in rows if r["zone"] == name)),
                     "k_studies": len({r["study"] for r in rows if r["zone"] == name})}
    out["moment"]["p_one_sided_negative"] = float(stats.t.cdf(out["moment"]["t"], df))
    out["beyond"]["p_one_sided_positive"] = float(stats.t.sf(out["beyond"]["t"], df))
    t = diff / se_diff
    out["difference_beyond_minus_moment"] = {"g": round(diff, 4), "se": round(se_diff, 4), "t": round(t, 3),
                                             "ci95": [round(diff - stats.t.ppf(0.975, df) * se_diff, 4), round(diff + stats.t.ppf(0.975, df) * se_diff, 4)],
                                             "p_one_sided_positive": float(stats.t.sf(t, df))}
    return out


def slope_log_interval(per_interval):
    """Secondary: random-effects meta-regression of g on log10(interval) with study-clustered SEs."""
    y = np.array([e["g"] for e in per_interval])
    v = np.array([e["v"] for e in per_interval])
    x = np.log10([e["interval_s"] for e in per_interval])
    X = np.column_stack([np.ones_like(x), x])
    t2 = reml_tau2(y, v, X)
    w = 1 / (v + t2)
    bread = np.linalg.inv(X.T @ np.diag(w) @ X)
    b = bread @ X.T @ (w * y)
    r = y - X @ b
    studies = sorted({e["study"] for e in per_interval})
    meat = np.zeros((2, 2))
    for s in studies:
        idx = [i for i, e in enumerate(per_interval) if e["study"] == s]
        u = X[idx].T @ (w[idx] * r[idx])
        meat += np.outer(u, u)
    m = len(studies)
    V = bread @ meat @ bread * m / (m - 1)
    se = float(np.sqrt(V[1, 1]))
    t = float(b[1] / se)
    return {"slope_per_log10_s": round(float(b[1]), 4), "se": round(se, 4), "t": round(t, 3),
            "p_one_sided_positive": float(stats.t.sf(t, m - 2)), "studies": m, "effects": len(per_interval)}


def main():
    out_path = sys.argv[1]
    per = primary_effects("data/extraction.csv")
    rows = by_zone(per) + secondary_effects("data/secondary.csv")
    res = {"moment_bound_s": MOMENT_S, "effects_by_study_zone": rows}
    res["primary"] = meta_regression(rows)
    # sensitivity: leave one study out (difference estimate)
    loo = {}
    for s in sorted({r["study"] for r in rows}):
        sub = [r for r in rows if r["study"] != s]
        if len({r["study"] for r in sub if r["zone"] == "moment"}) >= 1 and len({r["study"] for r in sub}) >= 4:
            mr = meta_regression(sub)
            loo[s] = mr["difference_beyond_minus_moment"]["g"]
    res["leave_one_study_out_difference"] = loo
    # sensitivity: studies with primary extraction only
    prim = [r for r in rows if r["source"] == "primary"]
    if len({r["study"] for r in prim}) >= 4 and any(r["zone"] == "moment" for r in prim):
        res["primary_extraction_only"] = meta_regression(prim)
    # secondary: continuous slope over primary per-interval effects (occasions averaged)
    agg = defaultdict(list)
    for e in per:
        agg[(e["study"], e["interval_s"])].append(e)
    pi = [{"study": s, "interval_s": iv, "g": float(np.mean([e["g"] for e in es])), "v": float(np.mean([e["v"] for e in es]))}
          for (s, iv), es in agg.items()]
    if len({e["study"] for e in pi}) >= 3:
        res["secondary_log_interval_slope_primary"] = slope_log_interval(pi)
    with open(out_path, "w") as fh:
        json.dump(res, fh, indent=2, sort_keys=True, default=list)
        fh.write("\n")
    print(json.dumps({k: v for k, v in res.items() if k != "effects_by_study_zone"}, indent=1))


if __name__ == "__main__":
    main()