"""
TEST 8 — PCAWG mutation density × harmonic-SCALE.

The manuscript already reports a sinusoidal fit on PCAWG pan-cancer mutation
density at the dominant period: R² = 0.951 (octant-mean) at T1.

This test asks: does phase at the rank-2 GWS scale (T1/3) carry additional
mutation-density-relevant information beyond phase at T1?

Procedure:
  - Load PCAWG 1-Mbp per-bin mutation density (mut_density column)
  - For each PCAWG bin, find the matching K562 v38 500-kb bins from Test 5
    and average their phi_T1, phi_T1div3, phi_T1div5 to get 1-Mbp-scale phases
  - Run the same nested OLS framework: M0 = phase at T1; M1 = + T1/3; M2 = + T1/5
  - Report R²_pop, R²_oct, nested F-tests
"""
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 = OUT / "test5_tsa_harmonic_scale.csv"
PCAWG = "/sessions/kind-brave-babbage/mnt/Wavelet/BIOMNI_UPLOAD_data_bundle/by_experiment/pcawg_mutations/data/bins_1mbp_phase_repli.csv"

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

print("="*70)
print("TEST 8: PCAWG mutation density × harmonic SCALE")
print("="*70)

# Load the canonical PCAWG mutation density data
print("\nLoading PCAWG 1-Mbp mutation density...")
pc = pd.read_csv(PCAWG)
print(f"PCAWG bins: {len(pc):,}")
print(f"Columns: {list(pc.columns)}")
# Normalise chrom prefix
pc['chrom'] = pc['chrom'].astype(str).apply(lambda x: f"chr{x}" if not x.startswith('chr') else x)

# Load Test 5 per-bin (500 kb) phases
print("Loading Test 5 per-bin phases (500 kb)...")
t5 = pd.read_csv(T5)
print(f"Test 5 bins: {len(t5):,}")

# For each 1-Mbp PCAWG bin, find the 500-kb bins whose mid falls inside [start, end)
# Aggregate by circular mean of phases
def circ_mean(angles):
    a = angles[~np.isnan(angles)]
    if len(a) == 0: return np.nan
    return np.arctan2(np.sin(a).sum(), np.cos(a).sum())

phi_T1 = np.full(len(pc), np.nan)
phi_T1div3 = np.full(len(pc), np.nan)
phi_T1div5 = np.full(len(pc), np.nan)

for i, row in pc.iterrows():
    chrom = row['chrom']; s = row['start']; e = row['end']
    sub = t5[(t5['chrom']==chrom) & (t5['bin_mid_hg38']>=s) & (t5['bin_mid_hg38']<e)]
    if len(sub) == 0: continue
    phi_T1[i]     = circ_mean(sub['phi_T1'].values)
    phi_T1div3[i] = circ_mean(sub['phi_T1div3'].values)
    phi_T1div5[i] = circ_mean(sub['phi_T1div5'].values)

pc['phi_T1']     = phi_T1
pc['phi_T1div3'] = phi_T1div3
pc['phi_T1div5'] = phi_T1div5

# Filter to bins with all phases + valid mutation density
mask = (pc['phi_T1'].notna() & pc['phi_T1div3'].notna() &
        pc['phi_T1div5'].notna() & pc['mut_density'].notna() & (pc['n_tumors_in_bin']>10))
print(f"Bins with all phases + mut_density (n_tumors>10): {mask.sum():,}")

X_phi = pc.loc[mask, ['phi_T1','phi_T1div3','phi_T1div5']].values
y = pc.loc[mask, 'mut_density'].values
# z-score y per chromosome to control for between-chromosome mut-rate differences
chroms = pc.loc[mask, 'chrom'].values
y_z = np.full_like(y, np.nan, dtype=float)
for c in np.unique(chroms):
    m_c = chroms == c
    yc = y[m_c]
    if len(yc) < 2: continue
    y_z[m_c] = (yc - yc.mean()) / (yc.std(ddof=1) + 1e-12)
y = y_z[~np.isnan(y_z)]
X_phi = X_phi[~np.isnan(y_z)]
n = len(y)
print(f"Bins with z-scored mut_density: {n:,}")

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

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()
print(f"PCAWG pan-cancer mutation density:  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}")
print()
print(f"  ΔR²_pop M0→M1: {m1['R2_pop']-m0['R2_pop']:+.4f}")
print(f"  ΔR²_oct M0→M1: {R2_o1-R2_o0:+.4f}")
print(f"  ΔR²_pop M1→M2: {m2['R2_pop']-m1['R2_pop']:+.4f}")
print(f"  ΔR²_oct M1→M2: {R2_o2-R2_o1:+.4f}")

# Save summary + augmented attribute table
result = {
    'label': 'PCAWG pan-cancer mutation density',
    '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,
}
pd.DataFrame([result]).to_csv(OUT / "test8_pcawg_mutation_density_summary.csv", index=False)
print(f"\nSaved: test8_pcawg_mutation_density_summary.csv")

# Also augment the 5-attribute summary into a 7-attribute summary
prev = pd.read_csv(OUT / "test7_five_attributes_harmonic_summary.csv")
prev = pd.concat([prev, pd.DataFrame([result])], ignore_index=True)
prev.to_csv(OUT / "test8_seven_attributes_harmonic_summary.csv", index=False)
print(f"Saved: test8_seven_attributes_harmonic_summary.csv ({len(prev)} attributes total)")

# Plot single-attribute panel (will be merged into S43 v3 in next step)
fig, ax = plt.subplots(1, 1, figsize=(5.5, 5))
xs = [0, 1, 2]
labels = ['M0\nrank-1', 'M1\n+T₁/3', 'M2\n+T₁/5']
R2_pop = [m0['R2_pop'], m1['R2_pop'], m2['R2_pop']]
R2_oct = [R2_o0, R2_o1, R2_o2]
w = 0.35
ax.bar([x-w/2 for x in xs], R2_pop, w, color=WONG['blue'], label='R² population')
ax.bar([x+w/2 for x in xs], R2_oct, w, color=WONG['gold'], label='R² octant-mean')
for i, x in enumerate(xs):
    ax.text(x-w/2, R2_pop[i] + 0.012, f"{R2_pop[i]:+.3f}", ha='center', fontsize=9)
    ax.text(x+w/2, R2_oct[i] + 0.012, f"{R2_oct[i]:+.3f}", ha='center', fontsize=9)
ax.annotate(f"F = {F1:.1f}\np = {p1:.2g}",
            xy=(0.5, max(R2_o0, R2_o1)+0.06), ha='center',
            fontsize=9, color=WONG['vermil'])
ax.annotate(f"F = {F2:.1f}\np = {p2:.2g}",
            xy=(1.5, max(R2_o1, R2_o2)+0.06), ha='center',
            fontsize=9, color=WONG['vermil'])
ax.set_xticks(xs); ax.set_xticklabels(labels, fontsize=10)
ax.set_ylabel('R²', fontsize=11)
ax.set_title(f"PCAWG pan-cancer mutation density\n(n = {n:,} 1-Mbp bins)",
             fontsize=11, loc='left')
ax.set_ylim(0, max(1.05, max(R2_oct)*1.15))
ax.legend(fontsize=10, loc='upper left')
ax.grid(alpha=0.3, axis='y')
plt.tight_layout()
plt.savefig(OUT / "test8_pcawg_mutation_density.png", dpi=200, bbox_inches='tight')
plt.savefig(OUT / "test8_pcawg_mutation_density.svg", bbox_inches='tight')
plt.close(fig)
print(f"Saved: test8_pcawg_mutation_density.png/svg")
print("\nDONE")
