#!/usr/bin/env python3
"""
tDCS reaching study — GLM hypothesis tests.

QUESTION
    Two KNOWN conditions anchor the scale and differ from each other:
        H2 (anchor check): Electrode-Box-B2 performs BETTER than Electrode-Box-A2.
    Two UNKNOWN conditions then get CLASSIFIED against those anchors:
        For "Electrode-Box-A" and for "Right-Electrode", is each one
        A2-like or B2-like? (We do not assume either belongs to B2.)
    Each unknown is compared to BOTH anchors; a group it cannot be
    distinguished from is its likely class.

Each comparison is examined through THREE complementary GLMs:

    (A) LEVEL / count      Poisson GEE on `success` (successful reaches per
                           session), adjusting for training day. Answers
                           "who makes more successful reaches overall?"

    (B) LEVEL / rate       Binomial GLM on success / attempts
                           (`success` out of `total`), cluster-robust.
                           Controls for differing numbers of attempts —
                           "who is more accurate per attempt?"

    (C) LEARNING RATE      Poisson GEE with a group x day interaction.
                           Answers "do the groups improve at different
                           RATES?" (slope of the learning curve).

All models adjust for time with day (and day^2 for the plateau) and account
for repeated measures on each subject via GEE (subject = cluster,
exchangeable correlation, robust SE) or cluster-robust standard errors.
Group is coded with Electrode-Box-B2 as the REFERENCE, so H1 and H2 read
directly off the group terms:
    Right-Electrode term  -> H1 (expect ~0 / non-significant)
    Electrode-Box-A2 term -> H2 (expect < 0: B2 above A2)
exp(coef) is an incidence-rate ratio (count model) or odds ratio (rate
model) relative to Box-B2.

CAVEATS
    * Tiny groups: Box-A2 n=3, Box-B2 n=3, Naive n=4, and each UNKNOWN
      (Electrode-Box-A, Right-Electrode) has only n=1 SUBJECT. Classifying a
      1-subject condition is weak: "matches anchor X" means "not statistically
      distinguishable from X", NOT proof of equivalence, and cluster-robust
      inference with a single cluster in a group is fragile (can show
      artificially small SEs, especially in the learning-rate interaction).

DATA PROVENANCE
    analysis/tdcs_reach_data.csv was reconstructed from the experiment
    database (experiment "tDCS", outcome = analysis_summary.counts.Success,
    attempts = analysis_summary.total, x-axis = the "# Days Reach" field) and
    verified cell-by-cell against the exported matrix (71/71 unambiguous
    cells on days 0-5 matched exactly).

USAGE
    python3 analysis/tdcs_glm.py
    Requires: pandas, numpy, scipy, statsmodels, matplotlib.
    Writes: analysis/tdcs_learning_curves.png
"""
import os
import argparse
import numpy as np
import pandas as pd
import patsy
import statsmodels.api as sm
import statsmodels.formula.api as smf
from scipy import stats
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

REF = "Electrode-Box-B2"
# Two KNOWN anchor conditions that differ (Box-B2 > Box-A2, see H2).
ANCHOR_LOW = "Electrode-Box-A2"
ANCHOR_HIGH = "Electrode-Box-B2"
# Two UNKNOWN conditions to classify: is each one A2-like or B2-like?
CANON_UNKNOWNS = ["Electrode-Box-A", "Right-Electrode"]
ALL_GROUPS = ["Naive", ANCHOR_LOW, ANCHOR_HIGH] + CANON_UNKNOWNS
# The active group config is set at runtime (after any --merge/--assign remap).
UNKNOWNS = list(CANON_UNKNOWNS)
MAIN = list(ALL_GROUPS)
OTHERS = [g for g in MAIN if g != REF]
# --merge preset: assume each unknown resolves into an anchor.
MERGE_MAP = {"Electrode-Box-A": ANCHOR_LOW, "Right-Electrode": ANCHOR_HIGH}
HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.join(HERE, "tdcs_reach_data.csv")
PLOT = os.path.join(HERE, "tdcs_learning_curves.png")
RHS = f"C(group, Treatment('{REF}')) + day_c + day_c2"


def gterm(name):
    """statsmodels/patsy column name for a group level (vs the reference)."""
    return f"C(group, Treatment('{REF}'))[T.{name}]"


def xterm(name):
    """Interaction column name (group level x day slope)."""
    return f"C(group, Treatment('{REF}'))[T.{name}]:day_c"


def load(max_day=None, mapping=None):
    df = pd.read_csv(DATA)
    if mapping:                                      # remap unknown -> target group
        df["group"] = df["group"].map(lambda g: mapping.get(g, g))
    df = df[df["group"].isin(ALL_GROUPS)].copy()    # keep canonical groups
    df = df[df["day"] >= 0].copy()                   # training days only (day 0..)
    if max_day is not None:
        df = df[df["day"] <= max_day].copy()         # restrict the analysis window
    df["day_c"] = df["day"] - df["day"].mean()       # center day
    df["day_c2"] = df["day_c"] ** 2
    present = [g for g in ALL_GROUPS if g != REF and g in set(df["group"])]
    df["group"] = pd.Categorical(df["group"], categories=[REF] + present)
    return df


def contrast(res, tname):
    """coef, robust SE, p, effect=exp(coef) and its 95% CI for one term."""
    ci = res.conf_int()
    lo, hi = ci.loc[tname]
    return dict(coef=res.params[tname], se=res.bse[tname], p=res.pvalues[tname],
                eff=np.exp(res.params[tname]), lo=np.exp(lo), hi=np.exp(hi))


def one_sided_p_below(c):
    """One-sided p for H: coef < 0 (i.e. Box-B2 above the compared group)."""
    return stats.norm.cdf(c["coef"] / c["se"])


def diff_contrast(res, a, b):
    """Compare group `a` vs group `b` (both coded relative to REF) via a linear
    contrast. Returns coef, robust SE, p, effect=exp(coef)=ratio a/b, and CI.
    Works for any pair, including b == REF (then it is just the `a` term)."""
    names = list(res.params.index)
    v = np.zeros(len(names))
    if a != REF:
        v[names.index(gterm(a))] += 1.0
    if b != REF:
        v[names.index(gterm(b))] -= 1.0
    tt = res.t_test(v)
    coef = float(np.ravel(tt.effect)[0])
    se = float(np.ravel(tt.sd)[0])
    p = float(np.ravel(tt.pvalue)[0])
    lo, hi = np.ravel(tt.conf_int())[:2]
    return dict(coef=coef, se=se, p=p, eff=np.exp(coef),
                lo=np.exp(lo), hi=np.exp(hi))


# ----------------------------------------------------------------------------- models
def fit_count_level(df):
    """(A) Poisson GEE on success counts; group main effects vs Box-B2."""
    return smf.gee(f"success ~ {RHS}", groups="subject", data=df,
                   family=sm.families.Poisson(),
                   cov_struct=sm.cov_struct.Exchangeable()).fit()


def fit_rate_level(df):
    """(B) Binomial GLM on success/attempts with cluster-robust SE.
    Zero-attempt sessions (total==0) carry no rate information and are dropped."""
    d = df[df["total"] > 0].copy()
    dropped = len(df) - len(d)
    if dropped:
        print(f"[rate model] dropped {dropped} zero-attempt session(s) (undefined rate)")
    X = patsy.dmatrix(RHS, d, return_type="dataframe")
    endog = np.column_stack([d["success"].values,
                             (d["total"] - d["success"]).values])
    return sm.GLM(endog, X, family=sm.families.Binomial()).fit(
        cov_type="cluster", cov_kwds={"groups": d["subject"].values})


def fit_learning_rate(df):
    """(C) Poisson GEE with group x day interaction (slope differences)."""
    return smf.gee(
        f"success ~ C(group, Treatment('{REF}')) * day_c + day_c2",
        groups="subject", data=df, family=sm.families.Poisson(),
        cov_struct=sm.cov_struct.Exchangeable()).fit()


def fit_mixed_count(df):
    """Subject random-intercept Poisson mixed model (MAP / Laplace).

    Adds a per-subject random intercept so each subject has its own baseline
    level; the group fixed effects are then estimated after allowing for that
    individual variation. A subject-specific complement to the population-
    average GEE. statsmodels has no frequentist Poisson GLMM, so this is the
    posterior-mode (MAP) fit with a Laplace covariance for the fixed effects
    (fit_vb failed to converge on these large counts; MAP is stable)."""
    import warnings
    from statsmodels.genmod.bayes_mixed_glm import PoissonBayesMixedGLM
    vc = {"subject": "0 + C(subject)"}          # random intercept per subject
    model = PoissonBayesMixedGLM.from_formula(f"success ~ {RHS}", vc, df)
    with warnings.catch_warnings():
        # MAP stops at |gradient|~5e-5 (effectively converged); silence the
        # over-strict "did not converge" notice.
        warnings.filterwarnings("ignore", message="Laplace fitting did not converge")
        return model.fit_map()


def mixed_contrast(res, tname):
    """coef, SD, approx p and IRR + 95% interval from a BayesMixedGLM fit."""
    names = list(res.model.exog_names)
    i = names.index(tname)
    mean, sd = float(res.fe_mean[i]), float(res.fe_sd[i])
    p = 2.0 * stats.norm.sf(abs(mean / sd))
    return dict(coef=mean, se=sd, p=p, eff=np.exp(mean),
                lo=np.exp(mean - 1.96 * sd), hi=np.exp(mean + 1.96 * sd))


# ----------------------------------------------------------------------------- reports
def describe(df):
    print("=" * 78)
    print("DESCRIPTIVES")
    print("=" * 78)
    grp = df.groupby("group", observed=True)
    tbl = pd.DataFrame({
        "n_subj": grp["subject"].nunique(),
        "n_sessions": grp.size(),
        "mean_success": grp["success"].mean().round(1),
        "mean_rate": (grp["success"].sum() / grp["total"].sum()).round(3),
        "max_day": grp["day"].max(),
    })
    print(tbl.to_string())


def report_level(title, res, effect_label):
    print("\n" + "=" * 78)
    print(title)
    print("=" * 78)
    print(f"Effect vs {REF} ({effect_label} relative to Box-B2):")
    h = f"{'group':20s} {'effect':>7s} {'95% CI':>16s} {'coef':>8s} {'SE':>7s} {'p':>9s}"
    print(h)
    print("-" * len(h))
    out = {}
    for g in OTHERS:
        c = contrast(res, gterm(g))
        out[g] = c
        ci = f"[{c['lo']:.2f}, {c['hi']:.2f}]"
        print(f"{g:20s} {c['eff']:7.3f} {ci:>16s} {c['coef']:8.3f} {c['se']:7.3f} {c['p']:9.4f}")
    return out


def report_learning(res):
    print("\n" + "=" * 78)
    print("(C) LEARNING RATE — Poisson GEE, group x day interaction")
    print("=" * 78)
    base = res.params["day_c"]
    print(f"Box-B2 learning slope: {np.exp(base):.3f}x successes per training day "
          f"(baseline).")
    print("Slope DIFFERENCE vs Box-B2 (exp(coef) = per-day multiplier on the rate ratio):")
    h = f"{'group':20s} {'slopeΔ/day':>10s} {'coef':>8s} {'SE':>7s} {'p':>9s}"
    print(h)
    print("-" * len(h))
    for g in OTHERS:
        c = contrast(res, xterm(g))
        print(f"{g:20s} {c['eff']:10.3f} {c['coef']:8.3f} {c['se']:7.3f} {c['p']:9.4f}")
    # Joint Wald test: do ANY groups differ from Box-B2 in learning rate?
    names = list(res.params.index)
    idx = [names.index(xterm(g)) for g in OTHERS]
    R = np.zeros((len(idx), len(names)))
    for i, j in enumerate(idx):
        R[i, j] = 1.0
    jt = res.wald_test(R, scalar=True)
    print(f"\nJoint test (all group x day interactions = 0): "
          f"chi2={float(jt.statistic):.2f}, df={len(idx)}, p={float(jt.pvalue):.4f}")
    print("  -> small p = groups improve at DIFFERENT rates; large p = parallel learning.")


def report_anchor_only(df):
    """A2-vs-B2 ONLY: refit on just the two anchors so nothing else influences
    the shared day terms / dispersion. B2 is the reference, so the single group
    term is the A2-vs-B2 comparison."""
    print("\n" + "=" * 78)
    print("ANCHORS ONLY — Box-A2 vs Box-B2 (two-group models)")
    print("=" * 78)
    d = df[df["group"].isin([ANCHOR_LOW, ANCHOR_HIGH])].copy()
    d["day_c"] = d["day"] - d["day"].mean()
    d["day_c2"] = d["day_c"] ** 2
    d["group"] = pd.Categorical(d["group"], categories=[ANCHOR_HIGH, ANCHOR_LOW])
    n_sub = d.groupby("group", observed=True)["subject"].nunique().to_dict()
    print(f"subjects: Box-B2={n_sub.get(ANCHOR_HIGH)}, Box-A2={n_sub.get(ANCHOR_LOW)}"
          f"   sessions: {len(d)}")
    print(f"formula: success ~ C(group, Treatment('{REF}')) + day_c + day_c2")
    for label, res in [("count/level (Poisson GEE)", fit_count_level(d)),
                       ("rate/level  (Binomial cluster-robust)", fit_rate_level(d))]:
        c = contrast(res, gterm(ANCHOR_LOW))
        p1 = one_sided_p_below(c)
        tag = "SUPPORTED" if (c["coef"] < 0 and p1 < 0.05) else "not supported"
        print(f"  [{label}] A2/B2 = {c['eff']:.2f} [{c['lo']:.2f}, {c['hi']:.2f}]  "
              f"=> B2 = {1.0 / c['eff']:.2f}x A2, one-sided p={p1:.4f}  -> H2 {tag}")


def report_mixed(df):
    """Subject random-intercept Poisson mixed model (sensitivity vs GEE)."""
    print("\n" + "=" * 78)
    print("SENSITIVITY — Poisson MIXED model, per-subject random intercept (MAP/Laplace)")
    print("=" * 78)
    try:
        res = fit_mixed_count(df)
    except Exception as exc:                       # VB can be finicky; degrade gracefully
        print(f"  (mixed model skipped: {type(exc).__name__}: {exc})")
        return
    print("Group effect vs Box-B2 (IRR; each subject given its own baseline):")
    h = f"{'group':20s} {'IRR':>7s} {'~95% CI':>16s} {'mean':>8s} {'SD':>7s} {'~p':>9s}"
    print(h)
    print("-" * len(h))
    for g in OTHERS:
        c = mixed_contrast(res, gterm(g))
        ci = f"[{c['lo']:.2f}, {c['hi']:.2f}]"
        print(f"{g:20s} {c['eff']:7.3f} {ci:>16s} {c['coef']:8.3f} {c['se']:7.3f} {c['p']:9.4f}")
    print(f"  subject random-intercept SD = {float(np.exp(res.vcp_mean[0])):.3f} "
          "(log scale); Laplace SEs, treat p-values as approximate.")
    print("  Compare directions/magnitudes with the GEE table (they should agree).")
    print("  Single-subject groups (Box-A, Right-Electrode) are partly confounded")
    print("  with their own random intercept here, so their effects are shrunk.")


def plot_curves(df, plot_path):
    fig, ax = plt.subplots(1, 2, figsize=(12, 4.8))
    for g in MAIN:
        s = df[df["group"] == g]
        d = s.groupby("day").apply(
            lambda x: pd.Series({"succ": x["success"].mean(),
                                 "rate": x["success"].sum() / x["total"].sum()}),
            include_groups=False)
        ax[0].plot(d.index, d["succ"], marker="o", ms=3, label=g)
        ax[1].plot(d.index, d["rate"], marker="o", ms=3, label=g)
    ax[0].set(xlabel="# Days Reach (training day)", ylabel="mean successful reaches",
              title="(A) Success count learning curves")
    ax[1].set(xlabel="# Days Reach (training day)", ylabel="success rate (success/attempts)",
              title="(B) Success-rate learning curves")
    for a in ax:
        a.legend(fontsize=8)
        a.grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig(plot_path, dpi=120)
    print(f"\nSaved learning-curve plot -> {plot_path}")


def classify_one(res, unknown):
    """Compare an unknown group against BOTH anchors; return a label + detail."""
    vlo = diff_contrast(res, unknown, ANCHOR_LOW)   # unknown vs Box-A2
    vhi = diff_contrast(res, unknown, ANCHOR_HIGH)  # unknown vs Box-B2
    like_lo = vlo["p"] >= 0.05   # indistinguishable from Box-A2
    like_hi = vhi["p"] >= 0.05   # indistinguishable from Box-B2
    if like_hi and not like_lo:
        verdict = "B2-like (differs from A2, matches B2)"
    elif like_lo and not like_hi:
        verdict = "A2-like (differs from B2, matches A2)"
    elif like_lo and like_hi:
        # tie: point toward whichever ratio is closer to 1 on the log scale
        nearer = ANCHOR_HIGH if abs(vhi["coef"]) < abs(vlo["coef"]) else ANCHOR_LOW
        verdict = f"AMBIGUOUS (matches both; numerically nearer {nearer})"
    else:
        verdict = "UNLIKE BOTH (differs from A2 and B2)"
    return vlo, vhi, verdict


def verdicts(count, rate):
    print("\n" + "=" * 78)
    print("VERDICTS")
    print("=" * 78)

    print("\nAnchor check — H2: Box-B2 BETTER than Box-A2 (the two anchors must differ)")
    for label, res in [("count/level", count), ("rate/level ", rate)]:
        c = contrast(res, gterm(ANCHOR_LOW))
        p1 = one_sided_p_below(c)
        tag = "SUPPORTED" if (c["coef"] < 0 and p1 < 0.05) else "not supported"
        print(f"  [{label}] Box-B2 = {1.0 / c['eff']:.2f}x Box-A2  "
              f"(one-sided p={p1:.4f})  -> {tag}")

    if not UNKNOWNS:
        print("\n(No unknown groups — they were merged into the anchors; "
              "only the H2 anchor contrast applies.)")
        return
    print("\nClassification — is each UNKNOWN condition A2-like or B2-like?")
    print("(ratio >1 = above that anchor; p = differs from that anchor)")
    for u in UNKNOWNS:
        print(f"\n  {u}:")
        for label, res in [("count/level", count), ("rate/level ", rate)]:
            vlo, vhi, verdict = classify_one(res, u)
            print(f"    [{label}] vs Box-A2: {vlo['eff']:.2f}x "
                  f"[{vlo['lo']:.2f},{vlo['hi']:.2f}] p={vlo['p']:.3f}   "
                  f"vs Box-B2: {vhi['eff']:.2f}x [{vhi['lo']:.2f},{vhi['hi']:.2f}] "
                  f"p={vhi['p']:.3f}")
            print(f"                 -> {verdict}")
    print("\n  NOTE: each unknown has ONLY 1 subject. 'Matches' means 'not statistically")
    print("  distinguishable' — weak evidence at n=1, not proof of equivalence.")


def main():
    ap = argparse.ArgumentParser(description="tDCS GLM hypothesis tests")
    ap.add_argument("--max-day", type=int, default=None,
                    help="restrict analysis to training days 0..MAX_DAY (default: all)")
    ap.add_argument("--merge", action="store_true",
                    help="preset: Electrode-Box-A->Box-A2 and Right-Electrode->Box-B2")
    ap.add_argument("--assign", type=str, default=None,
                    help="custom merges 'UNKNOWN=TARGET,...', e.g. "
                         "'Electrode-Box-A=Electrode-Box-B2,Right-Electrode=Electrode-Box-B2'")
    args = ap.parse_args()

    mapping = {}
    if args.merge:
        mapping.update(MERGE_MAP)
    if args.assign:
        for pair in args.assign.split(","):
            k, v = pair.split("=")
            mapping[k.strip()] = v.strip()

    df = load(args.max_day, mapping)

    global MAIN, OTHERS, UNKNOWNS
    present = set(df["group"])
    MAIN = [g for g in ALL_GROUPS if g in present]
    OTHERS = [g for g in MAIN if g != REF]
    UNKNOWNS = [g for g in CANON_UNKNOWNS if g in present]

    window = f"days 0-{args.max_day}" if args.max_day is not None else "all training days (0+)"
    tag = f"_d0-{args.max_day}" if args.max_day is not None else ""
    tag += "_merged" if mapping else ""
    plot_path = os.path.join(HERE, f"tdcs_learning_curves{tag}.png")
    mode = ("assign " + ", ".join(f"{k.replace('Electrode-','')}->{v.replace('Electrode-','')}"
                                  for k, v in mapping.items())) if mapping \
        else "5 groups (anchors + unknowns)"
    print(f"ANALYSIS WINDOW: {window}   |   MODE: {mode}   |   observations: {len(df)}")
    describe(df)
    count = fit_count_level(df)
    rate = fit_rate_level(df)
    learn = fit_learning_rate(df)
    report_level("(A) LEVEL / COUNT — Poisson GEE on successful reaches",
                 count, "incidence-rate ratio")
    report_level("(B) LEVEL / RATE — Binomial GLM on success/attempts (cluster-robust)",
                 rate, "odds ratio")
    report_learning(learn)
    report_anchor_only(df)
    report_mixed(df)
    verdicts(count, rate)
    plot_curves(df, plot_path)
    print("\nDone. See the module docstring for modeling choices and caveats.")


if __name__ == "__main__":
    main()
