"""
TEST 7 — Extend harmonic-SCALE test to two pathology attributes:
  (d) gnomAD-SV breakpoint density (population structural variation, 7,863 SVs)
  (e) PCAWG translocation breakpoints density (cancer somatic SVs, 63,285 TRAs)

Same nested OLS design as Test 5/6: M0 = phase at T1; M1 = + phase at T1/3;
M2 = + phase at T1/5. The two new pathology attributes broaden the
disease-domain coverage from rare-disease (ClinVar) to common-population
(gnomAD-SV) and cancer somatic (PCAWG translocations).
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import f as f_dist
from pathlib import Path

OUT = Path("/sessions/kind-brave-babbage/mnt/Wavelet/harmonic_synthesis_2026-05-01")
T5_PER_BIN = OUT / "test5_tsa_harmonic_scale.csv"
ANNO = "/sessions/kind-brave-babbage/mnt/Wavelet/go_rerun_v2_2026-04-30/master_annotation_K562_v38_hg38.csv"
GNOMAD = "/sessions/kind-brave-babbage/mnt/Wavelet/go_rerun_v2_2026-04-30/gnomad_sv_with_phase.csv"
PCAWG_TRA = ("/sessions/kind-brave-babbage/mnt/Wavelet/1 BIOMNI folder 1 paper 1 most of the work/"
             "Biomni_lab_downloads_20260404_145824/translocation/tra_breakpoints_with_phase.csv")

WONG = {'blue':'#0072B2','vermil':'#D55E00','green':'#009E73',
        'gold':'#E69F00','sky':'#56B4E9','pink':'#CC79A7','grey':'#999999'}

print("="*70)
print("TEST 7: Extend harmonic-SCALE test to gnomAD-SV + PCAWG translocations")
print("="*70)

print("\nLoading Test 5 per-bin phases + master annotation...")
df_phi = pd.read_csv(T5_PER_BIN)
df_anno = pd.read_csv(ANNO)
m = df_anno[['chrom','bin_mid_hg38','bin_start_hg38','bin_end_hg38']].merge(
    df_phi[['chrom','bin_mid_hg38','phi_T1','phi_T1div3','phi_T1div5']],
    on=['chrom','bin_mid_hg38'], how='inner')
print(f"Merged bins: {len(m):,}")

# ---- (d) gnomAD-SV ----
print("\n--- Loading gnomAD-SV ---")
gnomad = pd.read_csv(GNOMAD)
print(f"gnomAD-SV records: {len(gnomad):,}")
# Each SV has start (and end) — count breakpoints per bin (start + end as 2 events)
bp_gnomad = []
for _, r in gnomad.iterrows():
    bp_gnomad.append({'chrom': r['chrom'], 'pos': r['start']})
    bp_gnomad.append({'chrom': r['chrom'], 'pos': r['end']})
df_bp_gnomad = pd.DataFrame(bp_gnomad)
print(f"Total gnomAD-SV breakpoints: {len(df_bp_gnomad):,}")

# Count per bin
bp_count_gnomad = np.zeros(len(m), dtype=int)
for ic in m['chrom'].unique():
    rmask = (m['chrom']==ic).values
    bps_chrom = df_bp_gnomad[df_bp_gnomad['chrom']==ic]
    if not len(bps_chrom): continue
    bp_pos = bps_chrom['pos'].values
    starts = m.loc[rmask, 'bin_start_hg38'].values
    ends   = m.loc[rmask, 'bin_end_hg38'].values
    for ridx, (s, e) in zip(np.where(rmask)[0], zip(starts, ends)):
        bp_count_gnomad[ridx] = int(np.sum((bp_pos >= s) & (bp_pos < e)))
m['gnomad_sv_density'] = bp_count_gnomad
m['log_gnomad_sv'] = np.log1p(bp_count_gnomad)
print(f"Bins with ≥1 gnomAD-SV breakpoint: {(bp_count_gnomad>0).sum():,}")
print(f"Total gnomAD-SV breakpoints assigned: {bp_count_gnomad.sum():,}")

# ---- (e) PCAWG translocations ----
print("\n--- Loading PCAWG translocations ---")
tra = pd.read_csv(PCAWG_TRA)
print(f"PCAWG TRA pairs: {len(tra):,}")
# Each translocation has chrom1+pos1 and chrom2+pos2 — count both as breakpoints
bp_tra = []
for _, r in tra.iterrows():
    bp_tra.append({'chrom': r['chrom1'], 'pos': r['pos1']})
    bp_tra.append({'chrom': r['chrom2'], 'pos': r['pos2']})
df_bp_tra = pd.DataFrame(bp_tra)
print(f"Total PCAWG TRA breakpoints: {len(df_bp_tra):,}")

bp_count_tra = np.zeros(len(m), dtype=int)
for ic in m['chrom'].unique():
    rmask = (m['chrom']==ic).values
    bps_chrom = df_bp_tra[df_bp_tra['chrom']==ic]
    if not len(bps_chrom): continue
    bp_pos = bps_chrom['pos'].values
    starts = m.loc[rmask, 'bin_start_hg38'].values
    ends   = m.loc[rmask, 'bin_end_hg38'].values
    for ridx, (s, e) in zip(np.where(rmask)[0], zip(starts, ends)):
        bp_count_tra[ridx] = int(np.sum((bp_pos >= s) & (bp_pos < e)))
m['pcawg_tra_density'] = bp_count_tra
m['log_pcawg_tra'] = np.log1p(bp_count_tra)
print(f"Bins with ≥1 PCAWG TRA breakpoint: {(bp_count_tra>0).sum():,}")
print(f"Total PCAWG TRA breakpoints assigned: {bp_count_tra.sum():,}")

# ---- Run nested OLS for each ----
mask_phase = (m['phi_T1'].notna() & m['phi_T1div3'].notna() & m['phi_T1div5'].notna())

def design(phi):
    return np.column_stack([np.cos(phi), np.sin(phi)])

def fit_ols(X, y):
    beta, _, _, _ = np.linalg.lstsq(X, y, rcond=None)
    yhat = X @ beta
    SSE = np.sum((y - yhat)**2)
    SS_tot = np.sum((y - y.mean())**2)
    return {'beta': beta, 'yhat': yhat, 'SSE': SSE,
            'R2_pop': 1 - SSE/SS_tot,
            'df_resid': X.shape[0] - X.shape[1], 'p': X.shape[1]}

def nested_F(red, full):
    df_diff = red['df_resid'] - full['df_resid']
    F = ((red['SSE'] - full['SSE']) / df_diff) / (full['SSE'] / full['df_resid'])
    p_F = 1 - f_dist.cdf(F, df_diff, full['df_resid'])
    return F, df_diff, full['df_resid'], p_F

def oct_R2(yhat, oct_idx, y):
    n_oct = 8
    oct_y = np.array([y[oct_idx==k].mean() for k in range(n_oct)])
    oct_yhat = np.array([yhat[oct_idx==k].mean() for k in range(n_oct)])
    SS = np.sum((oct_y - oct_y.mean())**2)
    SSE = np.sum((oct_y - oct_yhat)**2)
    return 1 - SSE/SS if SS > 0 else np.nan

def run_test(y_col, label, df, mask):
    use = mask & df[y_col].notna()
    X_phi = df.loc[use, ['phi_T1','phi_T1div3','phi_T1div5']].values
    y = df.loc[use, y_col].values
    n = len(y)
    if n < 100:
        print(f"  {label}: n too small ({n})")
        return None
    X0 = np.column_stack([np.ones(n), design(X_phi[:,0])])
    X1 = np.column_stack([X0, design(X_phi[:,1])])
    X2 = np.column_stack([X1, design(X_phi[:,2])])
    m0_ = fit_ols(X0, y); m1_ = fit_ols(X1, y); m2_ = fit_ols(X2, y)
    edges = np.linspace(-np.pi, np.pi, 9)
    oct_idx = np.clip(np.digitize(X_phi[:,0], edges)-1, 0, 7)
    R2_o0 = oct_R2(m0_['yhat'], oct_idx, y)
    R2_o1 = oct_R2(m1_['yhat'], oct_idx, y)
    R2_o2 = oct_R2(m2_['yhat'], oct_idx, y)
    F1, _, _, p1 = nested_F(m0_, m1_)
    F2, _, _, p2 = nested_F(m1_, m2_)
    print(f"  {label}: n = {n:,}")
    print(f"    M0 (rank-1 only):     R²_pop = {m0_['R2_pop']:+.4f},  R²_oct = {R2_o0:+.4f}")
    print(f"    M1 (+ T1/3 SCALE):    R²_pop = {m1_['R2_pop']:+.4f},  R²_oct = {R2_o1:+.4f}   F = {F1:6.1f}, p = {p1:.3g}")
    print(f"    M2 (+ T1/5 SCALE):    R²_pop = {m2_['R2_pop']:+.4f},  R²_oct = {R2_o2:+.4f}   F = {F2:6.1f}, p = {p2:.3g}")
    return {'label': label, 'n': n,
            'M0_R2_pop': m0_['R2_pop'], 'M1_R2_pop': m1_['R2_pop'], 'M2_R2_pop': m2_['R2_pop'],
            'M0_R2_oct': R2_o0, 'M1_R2_oct': R2_o1, 'M2_R2_oct': R2_o2,
            'F_M0_M1': F1, 'p_M0_M1': p1,
            'F_M1_M2': F2, 'p_M1_M2': p2}

results = []
print("\n--- (d) gnomAD-SV (population structural variation) ---")
r = run_test('log_gnomad_sv', 'gnomAD-SV breakpoint density', m, mask_phase)
if r: results.append(r)

print("\n--- (e) PCAWG translocations (cancer somatic SVs) ---")
r = run_test('log_pcawg_tra', 'PCAWG TRA breakpoint density', m, mask_phase)
if r: results.append(r)

# Combine with previous Test 5 + Test 6 for a 5-attribute panel
# (rebuild from CSVs)
prev = pd.read_csv(OUT / "test6_three_attributes_harmonic_summary.csv")
# Test 5 (TSA) was R²_pop 0.146 → 0.337 → 0.384, R²_oct 0.921 → 0.960 → 0.965, F=758.78 etc.
test5_summary = pd.read_csv(OUT / "test5_tsa_harmonic_scale_summary.csv")

print("\n--- 5-attribute summary (combining Tests 5, 6, 7) ---")
all_attrs = []
# TSA
m5 = test5_summary
all_attrs.append({'label': 'TSA-seq (radial geometry)',
                   'n': 5275,
                   'M0_R2_pop': float(m5.iloc[0]['R2_pop']),
                   'M1_R2_pop': float(m5.iloc[1]['R2_pop']),
                   'M2_R2_pop': float(m5.iloc[2]['R2_pop']),
                   'M0_R2_oct': float(m5.iloc[0]['R2_oct']),
                   'M1_R2_oct': float(m5.iloc[1]['R2_oct']),
                   'M2_R2_oct': float(m5.iloc[2]['R2_oct']),
                   'F_M0_M1': float(m5.iloc[1]['F_vs_prev']),
                   'p_M0_M1': float(m5.iloc[1]['p_F_vs_prev']),
                   'F_M1_M2': float(m5.iloc[2]['F_vs_prev']),
                   'p_M1_M2': float(m5.iloc[2]['p_F_vs_prev'])})
# Test 6: Repli-seq, gene density, ClinVar CNV
for _, row in prev.iterrows():
    all_attrs.append({c: row[c] for c in row.index})
# Test 7 new
for r in results:
    all_attrs.append(r)

all_df = pd.DataFrame(all_attrs)
all_df.to_csv(OUT / "test7_five_attributes_harmonic_summary.csv", index=False)
print(all_df[['label','n','M0_R2_pop','M1_R2_pop','M2_R2_pop','F_M0_M1','p_M0_M1']].to_string(
    index=False, float_format=lambda x: f"{x:.4f}" if isinstance(x, float) else str(x)))

# ---- Plot: 6-attribute panel (4 from Test 5+6, 2 new from Test 7) ----
n_panels = len(all_df)
fig, axes = plt.subplots(2, 3, figsize=(15, 9))
axes_flat = axes.flatten()
for ax, (_, r) in zip(axes_flat, all_df.iterrows()):
    xs = [0, 1, 2]
    R2_pop = [r['M0_R2_pop'], r['M1_R2_pop'], r['M2_R2_pop']]
    R2_oct = [r['M0_R2_oct'], r['M1_R2_oct'], r['M2_R2_oct']]
    w = 0.35
    ax.bar([x-w/2 for x in xs], R2_pop, w, color=WONG['blue'], label='R² pop')
    ax.bar([x+w/2 for x in xs], R2_oct, w, color=WONG['gold'], label='R² oct')
    for i, x in enumerate(xs):
        ax.text(x-w/2, R2_pop[i] + 0.012*np.sign(R2_pop[i]+0.001),
                f"{R2_pop[i]:+.3f}", ha='center', fontsize=8)
        ax.text(x+w/2, R2_oct[i] + 0.012*np.sign(R2_oct[i]+0.001),
                f"{R2_oct[i]:+.3f}", ha='center', fontsize=8)
    ax.annotate(f"F = {r['F_M0_M1']:.1f}\np = {r['p_M0_M1']:.2g}",
                xy=(0.5, max(R2_oct[0], R2_oct[1])+0.06), ha='center',
                fontsize=8, color=WONG['vermil'])
    ax.annotate(f"F = {r['F_M1_M2']:.1f}\np = {r['p_M1_M2']:.2g}",
                xy=(1.5, max(R2_oct[1], R2_oct[2])+0.06), ha='center',
                fontsize=8, color=WONG['vermil'])
    ax.set_xticks(xs); ax.set_xticklabels(['M0\nrank-1','M1\n+T₁/3','M2\n+T₁/5'], fontsize=8.5)
    ax.set_ylabel('R²', fontsize=10)
    ax.set_title(f"{r['label']}\n(n = {int(r['n']):,})", fontsize=10, loc='left')
    y_low = min(0, min(R2_pop)*1.2 - 0.05)
    y_high = max(1.05, max(R2_oct)*1.15)
    ax.set_ylim(y_low, y_high)
    ax.axhline(0, color='black', lw=0.5)
    ax.grid(alpha=0.3, axis='y')
# Hide last empty panel if any
for k in range(len(all_df), len(axes_flat)):
    axes_flat[k].set_visible(False)
axes_flat[0].legend(fontsize=9, loc='upper left')

plt.suptitle('Phase at the rank-2 GWS scale ($T_1/3$) generalises across critical attributes:\n'
             'radial geometry, mechanism, function, rare disease, population variation, cancer somatic',
             fontsize=11, y=1.005)
plt.tight_layout()
plt.savefig(OUT / "test7_five_attributes_harmonic_scale.png", dpi=200, bbox_inches='tight')
plt.savefig(OUT / "test7_five_attributes_harmonic_scale.svg", bbox_inches='tight')
plt.close(fig)
print("\nSaved: test7_five_attributes_harmonic_scale.png/svg")
print("\nDONE")
