"""
TEST 6 — Extend harmonic-scale test to 3 critical attributes:
  (a) Repli-seq (replication timing — RIF1-coupled mechanism)
  (b) Gene density (functional readout)
  (c) ClinVar pathogenic CNV breakpoint density (disease vulnerability)

Same nested-OLS design as Test 5: M0 = phase at T1 only;
M1 = M0 + phase at T1/3 SCALE; M2 = M1 + phase at T1/5 SCALE.

For continuous attributes (a, b): direct OLS on per-bin attribute.
For event attribute (c): convert to per-bin breakpoint density (count of
breakpoints per 500-kb bin), then OLS on log(density+1) so a sinusoidal
A/B prediction is testable in the same framework as the continuous ones.
"""
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"   # already has phi_T1, phi_T1div3, phi_T1div5 per bin
ANNO = "/sessions/kind-brave-babbage/mnt/Wavelet/go_rerun_v2_2026-04-30/master_annotation_K562_v38_hg38.csv"
CLINVAR = "/sessions/kind-brave-babbage/mnt/Wavelet/go_rerun_v2_2026-04-30/clinvar_pathogenic_cnvs_with_phase.csv"

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

print("="*70)
print("TEST 6: Harmonic-SCALE test on three more critical attributes")
print("="*70)

# Load Test 5's per-bin phase data (already has phi_T1, phi_T1div3, phi_T1div5)
print("\nLoading Test 5 per-bin phases...")
df_phi = pd.read_csv(T5_PER_BIN)
print(f"Bins: {len(df_phi):,}")

# Load master annotation (has Repli-seq, gene_count, etc.)
print("Loading master annotation...")
df_anno = pd.read_csv(ANNO)
print(f"Annotation bins: {len(df_anno):,}")

# Merge on (chrom, bin_mid_hg38) — use 'bin_mid_hg38' from both
m = df_anno[['chrom','bin_mid_hg38','repliseq_K562','gene_count','log_gene_density','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: {len(m):,} bins")

# Load ClinVar CNV breakpoints, count per bin
print("Loading ClinVar pathogenic CNV breakpoints...")
clinvar = pd.read_csv(CLINVAR)
print(f"ClinVar pathogenic CNVs: {len(clinvar):,}")
# Count breakpoints per 500-kb bin (each CNV contributes 2 breakpoints: start and stop)
bp_records = []
for _, r in clinvar.iterrows():
    bp_records.append({'chrom': r['chrom'], 'pos': r['start']})
    bp_records.append({'chrom': r['chrom'], 'pos': r['stop']})
df_bp = pd.DataFrame(bp_records)
print(f"Total breakpoints: {len(df_bp):,}")

# For each row in m, count breakpoints inside its 500-kb bin
bp_count = np.zeros(len(m), dtype=int)
for ic in m['chrom'].unique():
    rows_chrom = m[m['chrom']==ic].copy()
    bps_chrom = df_bp[df_bp['chrom']==ic].copy()
    if not len(bps_chrom): continue
    bp_pos = bps_chrom['pos'].values
    rmask = (m['chrom']==ic).values
    starts = rows_chrom['bin_start_hg38'].values
    ends   = rows_chrom['bin_end_hg38'].values
    for ridx, (s, e) in zip(np.where(rmask)[0], zip(starts, ends)):
        bp_count[ridx] = int(np.sum((bp_pos >= s) & (bp_pos < e)))
m['cnv_breakpoint_density'] = bp_count
print(f"Bins with ≥1 CNV breakpoint: {(m['cnv_breakpoint_density']>0).sum():,}")
print(f"Total breakpoints assigned: {m['cnv_breakpoint_density'].sum():,}")

# Also z-score Repli-seq + gene density per chromosome
def zscore_per_chrom(df, col):
    return df.groupby('chrom')[col].transform(lambda x: (x - x.mean()) / (x.std(ddof=1) + 1e-12))

m['repliseq_z'] = zscore_per_chrom(m, 'repliseq_K562')
m['gene_density_z'] = zscore_per_chrom(m, 'log_gene_density')
m['log_cnv_bp_density'] = np.log1p(m['cnv_breakpoint_density'])

# Filter to bins with all 3 phases + valid attributes
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()
    if 'breakpoint' in y_col:
        use = use & (df[y_col].notna())   # log1p handled above
    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}

print()
results = []
print("\n--- (a) REPLICATION TIMING (Repli-seq, mechanism) ---")
r = run_test('repliseq_z', 'Repli-seq (RT)', m, mask_phase)
if r: results.append(r)

print("\n--- (b) GENE DENSITY (function) ---")
r = run_test('gene_density_z', 'log gene density', m, mask_phase)
if r: results.append(r)

print("\n--- (c) CNV BREAKPOINT DENSITY (disease) ---")
r = run_test('log_cnv_bp_density', 'log CNV breakpoint density', m, mask_phase)
if r: results.append(r)

res_df = pd.DataFrame(results)
res_df.to_csv(OUT / "test6_three_attributes_harmonic_summary.csv", index=False)
print(f"\nSaved: test6_three_attributes_harmonic_summary.csv")

# ---- Plot: 3-panel R² gain per attribute ----
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
xs = [0, 1, 2]
labels = ['M0\nrank-1', 'M1\n+T₁/3', 'M2\n+T₁/5']
panel_titles = [
    ('a   Repli-seq (replication timing — mechanism)',  WONG['green']),
    ('b   Gene density (functional readout)',            WONG['pink']),
    ('c   ClinVar pathogenic CNV breakpoint density\n    (disease vulnerability)', WONG['vermil']),
]
for ax, r, (title, col) in zip(axes, results, panel_titles):
    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² 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*np.sign(R2_pop[i]+0.001), f"{R2_pop[i]:+.3f}",
                ha='center', fontsize=8.5)
        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.5)
    # Annotations
    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.5, 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.5, color=WONG['vermil'])
    ax.set_xticks(xs); ax.set_xticklabels(labels, fontsize=9.5)
    ax.set_ylabel('R²', fontsize=10)
    ax.set_title(title, fontsize=10, loc='left')
    ax.set_ylim(min(0, min(R2_pop)*1.2 - 0.05), max(1.05, max(R2_oct)*1.15))
    ax.axhline(0, color='black', lw=0.5)
    ax.grid(alpha=0.3, axis='y')
    ax.text(2.4, 0.05, f"n = {r['n']:,}", fontsize=9, color='grey', ha='right')
axes[0].legend(fontsize=9, loc='upper left')

plt.suptitle('TEST 6 — Phase at the rank-2 GWS scale ($T_1/3$) carries scale-specific information '
             'beyond phase at $T_1$ for three additional critical attributes',
             fontsize=11, y=1.02)
plt.tight_layout()
plt.savefig(OUT / "test6_three_attributes_harmonic_scale.png", dpi=200, bbox_inches='tight')
plt.savefig(OUT / "test6_three_attributes_harmonic_scale.svg", bbox_inches='tight')
plt.close(fig)
print("\nSaved: test6_three_attributes_harmonic_scale.png/svg")
print("\nDONE")
