"""
03_aggregate_and_test.py
=========================

Aggregate per-sample chromperiod summaries (output of 02_run_chromperiod_per_sample.py)
into a single per-line, per-chromosome long table and execute the pre-committed
HGPS-vs-control hypothesis tests.

Pre-committed analysis (locked in PRE_COMMITTED_DISPOSITION.md, NOT modifiable
post-hoc):

  Primary endpoint:  sig95
    Mixed-effects model:  sig95 ~ disease + (1 | line) + (1 | chrom)
    One-sided test:       HGPS sig95 < control sig95
                          (consistent with v102 prediction that lamin loss
                           degrades the periodic surface coordinate)

  Secondary endpoint A:  dom_period_Mbp
    Mixed-effects model:  dom_period_Mbp ~ disease + (1 | line) + (1 | chrom)
    Two-sided test:       wavelength change (direction not pre-specified)

  Secondary endpoint B:  N_cycles  (= L_peaks / dominant period in peaks)
    Mixed-effects model:  N_cycles ~ disease + (1 | line) + (1 | chrom)
    Two-sided test:       cycle-count perturbation

Disposition rules (Outcome A / B / C):
  A.  Primary p < 0.05 one-sided AND mean sig95 attenuation > 15% in HGPS.
      => Lamin A/C maintains the supra-compartment wave; HGPS attenuates it.
      => Headline result: 1 paragraph in Discussion + Supp Fig.

  B.  Primary p < 0.05 AND attenuation 5-15%, OR primary 0.05<=p<0.10 AND
      attenuation >15%.  => Methods footnote + Supp Fig only.

  C.  Primary p >= 0.10 OR attenuation <5%.
      => Methods sentence noting the negative; figure deferred to revision.

Pseudo-replication discipline:
  * Statistical unit = (line, chrom) pair, NOT individual peaks or bins.
  * Random intercepts on `line` and `chrom` absorb donor- and chromosome-level
    baselines; effect of `disease` is what we test.
  * For lines with multiple replicates (none here -- GSE150136 is N=1 per line)
    we would average replicates first.
  * 6 HGPS lines x 22 autosomes = 132 obs; 2 control lines x 22 = 44 obs.
    Total n=176, between-group n_lines=8.  Power note: at this sample size
    we have power 0.80 to detect Cohen's d ~ 0.85 in sig95 (large effect).
    A null result therefore bounds the effect from above; it does not refute
    the hypothesis at smaller magnitudes.
"""

import argparse
import os
import sys

import numpy as np
import pandas as pd

try:
    import statsmodels.api as sm
    import statsmodels.formula.api as smf
    HAVE_SM = True
except ImportError:
    HAVE_SM = False


def aggregate_summaries(summary_dir):
    """Read all per-sample summaries and stitch them together."""
    files = sorted([f for f in os.listdir(summary_dir)
                    if f.endswith('_per_chrom_summary.csv')])
    if not files:
        raise FileNotFoundError(
            f"No *_per_chrom_summary.csv files in {summary_dir}")
    dfs = []
    for f in files:
        d = pd.read_csv(os.path.join(summary_dir, f))
        dfs.append(d)
    out = pd.concat(dfs, ignore_index=True)
    return out


def fit_mixed_one_sided_lt(df, response, group='line', second='chrom'):
    """Mixed-model with random intercepts on `group` AND `second` via crossed
    grouping. statsmodels' MixedLM only supports a single grouping variable
    natively; we approximate by using `group` as the primary random effect and
    treat `chrom` as a fixed effect block.

    Returns (effect_HGPS_minus_control, two_sided_p, one_sided_p_lt).
    """
    if not HAVE_SM:
        raise RuntimeError("statsmodels not installed; install with "
                           "`pip3 install --user statsmodels`")
    df = df.copy()
    df['disease_HGPS'] = (df['disease'] == 'HGPS').astype(int)
    # chrom as fixed-effect block (22 levels, treated as factor):
    formula = f"{response} ~ disease_HGPS + C(chrom)"
    model = smf.mixedlm(formula, df, groups=df[group])
    fit = model.fit(method='lbfgs', reml=False)
    coef = fit.params['disease_HGPS']
    pval_two = fit.pvalues['disease_HGPS']
    tval = fit.tvalues['disease_HGPS']
    # one-sided: H1: HGPS < control => coef < 0
    if coef < 0:
        one_sided_p = pval_two / 2.0
    else:
        one_sided_p = 1.0 - pval_two / 2.0
    return float(coef), float(pval_two), float(one_sided_p), fit


def disposition(primary_one_sided_p, attenuation_pct):
    if primary_one_sided_p < 0.05 and attenuation_pct > 15.0:
        return 'A — Headline result (Discussion paragraph + Supp Fig)'
    if (primary_one_sided_p < 0.05 and 5.0 <= attenuation_pct <= 15.0) or \
       (0.05 <= primary_one_sided_p < 0.10 and attenuation_pct > 15.0):
        return 'B — Methods footnote + Supp Fig'
    return 'C — Methods sentence; figure deferred to revision'


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument('--summary-dir', required=True,
                   help='Dir with *_per_chrom_summary.csv files')
    p.add_argument('--out', required=True,
                   help='Output stats results CSV path')
    args = p.parse_args()

    df = aggregate_summaries(args.summary_dir)
    print(f"[aggregate] Loaded {len(df)} (line, chrom) rows from "
          f"{df['line'].nunique()} lines.", flush=True)
    print(f"[aggregate] Lines: {sorted(df['line'].unique())}", flush=True)
    print(f"[aggregate] Diseases: "
          f"{df.groupby('disease').size().to_dict()}", flush=True)

    # Per-line means
    per_line = df.groupby(['line', 'disease']).agg(
        mean_sig95=('sig95', 'mean'),
        mean_dom_period_Mbp=('dom_period_Mbp', 'mean'),
        median_N_cycles=('N_cycles', 'median'),
        n_chrom=('chrom', 'count'),
    ).reset_index()
    print(f"\n=== Per-line summary ===", flush=True)
    print(per_line.to_string(index=False), flush=True)

    # Group means for attenuation pct
    grp = df.groupby('disease')['sig95'].mean()
    if 'control' in grp.index and 'HGPS' in grp.index:
        ctrl_mean = float(grp.loc['control'])
        hgps_mean = float(grp.loc['HGPS'])
        atten = (ctrl_mean - hgps_mean) / ctrl_mean * 100 if ctrl_mean else 0.0
    else:
        ctrl_mean = float('nan'); hgps_mean = float('nan'); atten = float('nan')

    print(f"\n=== Group means (sig95) ===", flush=True)
    print(f"  control: {ctrl_mean:.4f}", flush=True)
    print(f"  HGPS   : {hgps_mean:.4f}", flush=True)
    print(f"  attenuation: {atten:+.1f}%", flush=True)

    if not HAVE_SM:
        print("\n[stats] statsmodels not available; skipping mixed-model.",
              flush=True)
        sys.exit(2)

    print("\n=== Mixed model: sig95 ~ disease + C(chrom), groups=line ===",
          flush=True)
    eff_p, p_two_p, p_one_p, fit_p = fit_mixed_one_sided_lt(df, 'sig95')
    print(f"  effect (HGPS - control): {eff_p:+.4f}", flush=True)
    print(f"  two-sided p           : {p_two_p:.4g}", flush=True)
    print(f"  one-sided p (HGPS<ctrl): {p_one_p:.4g}", flush=True)

    print("\n=== Mixed model: dom_period_Mbp ~ disease + C(chrom) ===",
          flush=True)
    eff_T, p_two_T, _, fit_T = fit_mixed_one_sided_lt(df, 'dom_period_Mbp')
    print(f"  effect (HGPS - control): {eff_T:+.3f} Mbp", flush=True)
    print(f"  two-sided p           : {p_two_T:.4g}", flush=True)

    print("\n=== Mixed model: N_cycles ~ disease + C(chrom) ===", flush=True)
    eff_N, p_two_N, _, fit_N = fit_mixed_one_sided_lt(df, 'N_cycles')
    print(f"  effect (HGPS - control): {eff_N:+.3f} cycles", flush=True)
    print(f"  two-sided p           : {p_two_N:.4g}", flush=True)

    disp = disposition(p_one_p, atten)
    print(f"\n=== DISPOSITION ===\n  {disp}\n", flush=True)

    # Save results table
    results = pd.DataFrame([
        {'response': 'sig95', 'effect': eff_p, 'two_sided_p': p_two_p,
         'one_sided_p_lt': p_one_p, 'control_mean': ctrl_mean,
         'hgps_mean': hgps_mean, 'attenuation_pct': atten,
         'disposition': disp},
        {'response': 'dom_period_Mbp', 'effect': eff_T,
         'two_sided_p': p_two_T, 'one_sided_p_lt': float('nan'),
         'control_mean': float(df[df['disease'] == 'control']['dom_period_Mbp'].mean()),
         'hgps_mean': float(df[df['disease'] == 'HGPS']['dom_period_Mbp'].mean()),
         'attenuation_pct': float('nan'), 'disposition': ''},
        {'response': 'N_cycles', 'effect': eff_N,
         'two_sided_p': p_two_N, 'one_sided_p_lt': float('nan'),
         'control_mean': float(df[df['disease'] == 'control']['N_cycles'].mean()),
         'hgps_mean': float(df[df['disease'] == 'HGPS']['N_cycles'].mean()),
         'attenuation_pct': float('nan'), 'disposition': ''},
    ])
    results.to_csv(args.out, index=False)
    per_line.to_csv(args.out.replace('.csv', '_per_line.csv'), index=False)
    df.to_csv(args.out.replace('.csv', '_long.csv'), index=False)
    print(f"[stats] Saved: {args.out}", flush=True)


if __name__ == '__main__':
    main()
