"""
TEST 5 — TSA-seq harmonic-SCALE test (one critical attribute).

Question: Does phase at T1/3 (the rank-2 GWS harmonic SCALE) explain
   SON-LMNB1 radial position INDEPENDENTLY of phase at T1?

This is distinct from Test 4 (v83): Test 4 added cos(2φ_T1), sin(2φ_T1) —
the doubled-frequency Fourier harmonic of the T1 phase. Test 5 instead
extracts phase at the T1/3 SCALE — a separate CWT scale with its own
phase angle — and asks whether THAT carries new radial-geometry info.

If yes:  the radial accordion has scale-specific substructure — 3 sub-radial
         layers per supra-compartment cycle — and the rank-2 harmonic in
         the GWS is not just a square-wave artefact but a genuine geometric
         feature.
If no:   the radial geometry is fully captured by phase at T1; the rank-2
         harmonic is a within-compartment 1D structure with no radial
         readout.

Method:
- For each K562 v38 autosome, run CWT on consecutive-peak signalValue
- Extract dominant period T1 from GWS rank-1
- Extract phase at scale T1 (already in canonical phase map)
- Extract phase at scale T1/3 (new) and at T1/5 (also new, for completeness)
- Bin to 500-kb genomic bins (same as canonical phase map)
- Match to TSA-seq SON_z - LMNB1_z axis (master annotation matrix)
- Fit nested OLS:
    M0: y = a0 + a1·cos(φ_T1) + b1·sin(φ_T1)               [baseline rank-1]
    M1: M0 + a2·cos(φ_T1/3) + b2·sin(φ_T1/3)               [add T1/3 SCALE]
    M2: M1 + a3·cos(φ_T1/5) + b3·sin(φ_T1/5)               [add T1/5 SCALE]
- Report octant-mean R² + partial F-test M1 vs M0 (the critical comparison)

Outputs:
- test5_tsa_harmonic_scale.csv (per-bin phase at T1, T1/3, T1/5 + TSA axis)
- test5_tsa_harmonic_scale_summary.csv (model R² + F-test)
- test5_tsa_harmonic_scale.png/svg
"""
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
import importlib.util, sys

OUT = Path("/sessions/kind-brave-babbage/mnt/Wavelet/harmonic_synthesis_2026-05-01")
PEAKS = "/sessions/kind-brave-babbage/mnt/Wavelet/k562_regen_option_A_2026-04-18/ENCFF361VGY.bed"
ANNO = "/sessions/kind-brave-babbage/mnt/Wavelet/go_rerun_v2_2026-04-30/master_annotation_K562_v38_hg38.csv"

SRC = "/sessions/kind-brave-babbage/mnt/Wavelet/k562_regen_option_A_2026-04-18/regen_k562_phase_maps.py"
spec = importlib.util.spec_from_file_location('chromperiod', SRC)
cp = importlib.util.module_from_spec(spec)
sys.modules['chromperiod'] = cp
spec.loader.exec_module(cp)

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

def cwt_at_scales(signal_raw, target_periods_pi, n_scales=180,
                   period_min=4, period_max=14000, wavelet_order=2):
    """Run CWT, return phase at each requested target period (peak-index units)."""
    N = len(signal_raw)
    sigma = np.std(signal_raw, ddof=1)
    if sigma == 0: return None
    signal = (signal_raw - np.mean(signal_raw)) / sigma
    fourier_factor, coi_factor = cp._get_paul_params(wavelet_order)
    s0 = period_min / fourier_factor
    smax = period_max / fourier_factor
    scales = np.geomspace(s0, smax, n_scales)
    periods = scales * fourier_factor
    N_pad = int(2 ** np.ceil(np.log2(N)))
    x_hat = np.fft.fft(signal, n=N_pad)
    omega_pos = 2.0 * np.pi * np.arange(1, N_pad // 2 + 1) / N_pad
    omega_neg = -2.0 * np.pi * np.arange(N_pad // 2 - 1, 0, -1) / N_pad
    omega = np.concatenate([[0.0], omega_pos, omega_neg])
    coi = cp.compute_coi(N, 1.0, coi_factor)

    # Compute phase at each target period (find nearest scale)
    out = {}
    for label, T_target in target_periods_pi.items():
        if T_target is None or T_target < period_min:
            out[label] = {'phi': None, 'in_coi': None, 'period_pi': T_target}
            continue
        j = int(np.argmin(np.abs(periods - T_target)))
        s = scales[j]
        psi_vals = np.conj(cp._paul_psi_hat(s, omega, m=wavelet_order))
        norm = np.sqrt(2.0 * np.pi * s)
        W_hat = norm * x_hat * psi_vals
        W_n = np.fft.ifft(W_hat)[:N]
        phi = np.angle(W_n)
        in_coi = periods[j] >= coi   # True where COI excludes this position
        out[label] = {'phi': phi, 'in_coi': in_coi, 'period_pi': periods[j]}
    return out, periods, scales, coi

def find_T1(signal_raw, n_scales=180, period_min=4, period_max=14000):
    """Return dominant period in peak-index units (rank-1 GWS)."""
    N = len(signal_raw)
    sigma = np.std(signal_raw, ddof=1)
    if sigma == 0: return None
    signal = (signal_raw - np.mean(signal_raw)) / sigma
    fourier_factor, coi_factor = cp._get_paul_params(2)
    s0 = period_min / fourier_factor
    smax = period_max / fourier_factor
    scales = np.geomspace(s0, smax, n_scales)
    periods = scales * fourier_factor
    N_pad = int(2 ** np.ceil(np.log2(N)))
    x_hat = np.fft.fft(signal, n=N_pad)
    omega_pos = 2.0 * np.pi * np.arange(1, N_pad // 2 + 1) / N_pad
    omega_neg = -2.0 * np.pi * np.arange(N_pad // 2 - 1, 0, -1) / N_pad
    omega = np.concatenate([[0.0], omega_pos, omega_neg])
    coi = cp.compute_coi(N, 1.0, coi_factor)
    gws = np.zeros(n_scales)
    for j, s in enumerate(scales):
        psi_vals = np.conj(cp._paul_psi_hat(s, omega, m=2))
        norm = np.sqrt(2.0 * np.pi * s)
        W_hat = norm * x_hat * psi_vals
        W_n = np.fft.ifft(W_hat)[:N]
        m = periods[j] < coi
        gws[j] = np.mean(np.abs(W_n[m])**2) if m.sum() > 0 else np.nan
    valid = np.isfinite(gws)
    return periods[np.argmax(np.where(valid, gws, -np.inf))]


# ---- Run ----

print("="*70)
print("TEST 5: Does phase at T1/3 SCALE explain TSA-seq beyond phase at T1?")
print("="*70)

print("\nLoading peaks + annotation...")
df_peaks = cp.load_narrowpeak(PEAKS)
df_anno = pd.read_csv(ANNO)
print(f"Peaks: {len(df_peaks):,}; annotation bins: {len(df_anno):,}")

# Per-bin SON-LMNB1 axis (z-score per chromosome)
def zscore_per_chrom(df, col):
    return df.groupby('chrom')[col].transform(lambda x: (x - x.mean()) / x.std(ddof=1))
df_anno['son_z']  = zscore_per_chrom(df_anno, 'son_tsa')
df_anno['lmnb_z'] = zscore_per_chrom(df_anno, 'lmnb_tsa')
df_anno['axis']   = df_anno['son_z'] - df_anno['lmnb_z']

CHROMS = [f'chr{i}' for i in range(1, 23)] + ['chrX']
all_rows = []

for chrom in CHROMS:
    sub_peaks = df_peaks[df_peaks['chrom']==chrom].sort_values('start').reset_index(drop=True)
    if len(sub_peaks) < 100: continue
    signal = sub_peaks['signalValue'].values.astype(float)
    pos_peaks = ((sub_peaks['start'].values + sub_peaks['end'].values)/2.0).astype(int)

    T1_pi = find_T1(signal)
    if T1_pi is None or T1_pi < 12: continue

    # Extract phase at T1, T1/3, T1/5 (three scales)
    targets = {'T1': T1_pi, 'T1_div3': T1_pi/3, 'T1_div5': T1_pi/5}
    res, periods, scales, coi = cwt_at_scales(signal, targets)

    # Bin peak-level phases to 500-kb genomic bins (matched to annotation)
    sub_anno = df_anno[df_anno['chrom']==chrom].copy().sort_values('bin_mid_hg38').reset_index(drop=True)
    bin_starts = sub_anno['bin_start_hg38'].values
    bin_ends   = sub_anno['bin_end_hg38'].values

    # For each bin, find peaks falling inside, take circular mean of their phase
    # (and amplitude-weighted circular mean for cleaner estimate)
    n_bins = len(sub_anno)
    phi_T1     = np.full(n_bins, np.nan)
    phi_T1div3 = np.full(n_bins, np.nan)
    phi_T1div5 = np.full(n_bins, np.nan)

    for ibin, (bs, be) in enumerate(zip(bin_starts, bin_ends)):
        m = (pos_peaks >= bs) & (pos_peaks < be)
        if m.sum() == 0: continue
        for label, key in [('T1','phi_T1'),('T1_div3','phi_T1div3'),('T1_div5','phi_T1div5')]:
            phi_arr = res[label]['phi']
            if phi_arr is None: continue
            phi_local = phi_arr[m]
            if len(phi_local) == 0: continue
            # circular mean
            cmean = np.arctan2(np.sin(phi_local).mean(), np.cos(phi_local).mean())
            if key == 'phi_T1':     phi_T1[ibin]     = cmean
            elif key == 'phi_T1div3': phi_T1div3[ibin] = cmean
            elif key == 'phi_T1div5': phi_T1div5[ibin] = cmean

    rows = pd.DataFrame({
        'chrom': chrom,
        'bin_mid_hg38': sub_anno['bin_mid_hg38'].values,
        'axis': sub_anno['axis'].values,
        'phi_T1':     phi_T1,
        'phi_T1div3': phi_T1div3,
        'phi_T1div5': phi_T1div5,
    })
    all_rows.append(rows)
    print(f"  {chrom}: T1 = {T1_pi:.0f} pi, "
          f"valid bins (T1, T1/3, T1/5): "
          f"{(~np.isnan(phi_T1)).sum()}, "
          f"{(~np.isnan(phi_T1div3)).sum()}, "
          f"{(~np.isnan(phi_T1div5)).sum()}")

dfm = pd.concat(all_rows, ignore_index=True)
dfm.to_csv(OUT / "test5_tsa_harmonic_scale.csv", index=False)
print(f"\nSaved: test5_tsa_harmonic_scale.csv ({len(dfm)} rows)")

# Filter to bins with all three phases + valid TSA axis
mask = (dfm['axis'].notna() & dfm['phi_T1'].notna() &
        dfm['phi_T1div3'].notna() & dfm['phi_T1div5'].notna())
df_f = dfm[mask].copy()
print(f"Bins with all 3 phases + axis: {len(df_f):,}")

# ---- Build design matrices ----
def design(phi, intercept_done=True):
    return np.column_stack([np.cos(phi), np.sin(phi)])

X0 = np.column_stack([np.ones(len(df_f)), design(df_f['phi_T1'].values)])
X1 = np.column_stack([X0, design(df_f['phi_T1div3'].values)])
X2 = np.column_stack([X1, design(df_f['phi_T1div5'].values)])
y = df_f['axis'].values

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)
    R2 = 1 - SSE / SS_tot
    df_resid = X.shape[0] - X.shape[1]
    return {'beta': beta, 'yhat': yhat, 'SSE': SSE, 'R2_pop': R2,
            'df_resid': df_resid, 'p': X.shape[1]}

m0 = fit_ols(X0, y)
m1 = fit_ols(X1, y)
m2 = fit_ols(X2, y)

# Octant-mean R²
n_oct = 8
edges = np.linspace(-np.pi, np.pi, n_oct+1)
oct_idx = np.digitize(df_f['phi_T1'].values, edges) - 1
oct_idx = np.clip(oct_idx, 0, n_oct-1)
oct_y = np.array([y[oct_idx==k].mean() for k in range(n_oct)])

def oct_R2(yhat):
    yhat_oct = 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 - yhat_oct)**2)
    return 1 - SSE/SS if SS > 0 else np.nan

R2_oct_0 = oct_R2(m0['yhat'])
R2_oct_1 = oct_R2(m1['yhat'])
R2_oct_2 = oct_R2(m2['yhat'])

print()
print(f"  Model M0 (rank-1 phase only):              R²_pop = {m0['R2_pop']:.4f}, R²_oct = {R2_oct_0:.4f}")
print(f"  Model M1 (M0 + phase at T1/3 scale):       R²_pop = {m1['R2_pop']:.4f}, R²_oct = {R2_oct_1:.4f}")
print(f"  Model M2 (M1 + phase at T1/5 scale):       R²_pop = {m2['R2_pop']:.4f}, R²_oct = {R2_oct_2:.4f}")

# Nested F-tests
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

F1, df_n1, df_d1, p1 = nested_F(m0, m1)
F2, df_n2, df_d2, p2 = nested_F(m1, m2)

print()
print(f"  M1 vs M0 (does T1/3 add to T1?):  F = {F1:.2f}, df = ({df_n1}, {df_d1}), p = {p1:.4g}")
print(f"  M2 vs M1 (does T1/5 add further?): F = {F2:.2f}, df = ({df_n2}, {df_d2}), p = {p2:.4g}")
print()
print(f"  ΔR²_pop M0→M1: {m1['R2_pop'] - m0['R2_pop']:+.4f}")
print(f"  ΔR²_oct M0→M1: {R2_oct_1 - R2_oct_0:+.4f}")
print(f"  ΔR²_pop M1→M2: {m2['R2_pop'] - m1['R2_pop']:+.4f}")
print(f"  ΔR²_oct M1→M2: {R2_oct_2 - R2_oct_1:+.4f}")

summary = pd.DataFrame([
    {'model':'M0', 'description':'rank-1 phase only', 'p':m0['p'],
     'R2_pop':m0['R2_pop'], 'R2_oct':R2_oct_0, 'SSE':m0['SSE']},
    {'model':'M1', 'description':'+ phase at T1/3 SCALE', 'p':m1['p'],
     'R2_pop':m1['R2_pop'], 'R2_oct':R2_oct_1, 'SSE':m1['SSE'],
     'F_vs_prev':F1, 'p_F_vs_prev':p1},
    {'model':'M2', 'description':'+ phase at T1/5 SCALE', 'p':m2['p'],
     'R2_pop':m2['R2_pop'], 'R2_oct':R2_oct_2, 'SSE':m2['SSE'],
     'F_vs_prev':F2, 'p_F_vs_prev':p2},
])
summary.to_csv(OUT / "test5_tsa_harmonic_scale_summary.csv", index=False)

# ---- Plot: bar chart of R² + scatter ----
fig, axes = plt.subplots(1, 2, figsize=(13, 5))

# Panel a — bar chart of R²
ax = axes[0]
xs = [0, 1, 2]
labels = ['M0\nrank-1\nphase', 'M1\n+ T₁/3\nSCALE', 'M2\n+ T₁/5\nSCALE']
R2_pop = [m0['R2_pop'], m1['R2_pop'], m2['R2_pop']]
R2_oct = [R2_oct_0, R2_oct_1, R2_oct_2]
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)
# Annotate F-test results
ax.annotate(f"F = {F1:.1f}\np = {p1:.2g}",
            xy=(0.5, max(R2_oct_0, R2_oct_1)+0.06), ha='center',
            fontsize=9, color=WONG['vermil'])
ax.annotate(f"F = {F2:.1f}\np = {p2:.2g}",
            xy=(1.5, max(R2_oct_1, R2_oct_2)+0.06), ha='center',
            fontsize=9, color=WONG['vermil'])
ax.set_xticks(xs); ax.set_xticklabels(labels, fontsize=10)
ax.set_ylabel('R² (TSA-seq SON−LMNB1 axis)', fontsize=11)
ax.set_ylim(0, 1.15)
ax.set_title('a   Adding phase at the T₁/3 scale to the rank-1 fit\n'
             '    tests for radial-geometric content beyond the dominant period',
             fontsize=10, loc='left')
ax.legend(fontsize=9, loc='upper left')
ax.grid(alpha=0.3, axis='y')

# Panel b — predicted vs observed octant means for M0 vs M1
ax2 = axes[1]
oct_phi = (edges[:-1] + edges[1:])/2
oct_yhat_0 = np.array([m0['yhat'][oct_idx==k].mean() for k in range(n_oct)])
oct_yhat_1 = np.array([m1['yhat'][oct_idx==k].mean() for k in range(n_oct)])
ax2.errorbar(oct_phi, oct_y, fmt='o', color='black', ms=10, mfc='white', mew=2, capsize=4,
             label='Observed octant means')
phi_grid = np.linspace(-np.pi, np.pi, 200)
def predict_M0(phi_grid):
    return m0['beta'][0] + m0['beta'][1]*np.cos(phi_grid) + m0['beta'][2]*np.sin(phi_grid)
ax2.plot(phi_grid, predict_M0(phi_grid), color=WONG['blue'], lw=2, ls='-',
         label=f"M0 (rank-1, R²_oct={R2_oct_0:.3f})")
# For M1 we'd need 2D phase grid; just plot the M1 octant predictions as another marker
ax2.plot(oct_phi, oct_yhat_1, 'D', color=WONG['gold'], ms=8, mew=1.2, mec='black',
         label=f"M1 octant predictions (R²_oct={R2_oct_1:.3f})")
ax2.axhline(0, color='black', lw=0.5)
ax2.set_xticks(np.linspace(-np.pi, np.pi, 5))
ax2.set_xticklabels(['−π','−π/2','0','+π/2','+π'])
ax2.set_xlabel('CWT phase φ at T₁', fontsize=11)
ax2.set_ylabel(r'SON−LMNB1 axis (z; A-side $> 0$)', fontsize=11)
ax2.set_title('b   Adding T₁/3-scale phase shifts the M1 predictions\n'
              '    away from the smooth M0 sinusoid (octant view)',
              fontsize=10, loc='left')
ax2.legend(fontsize=9, loc='best')
ax2.grid(alpha=0.3)

plt.suptitle('TEST 5 — TSA-seq harmonic-SCALE test: does the rank-2 GWS scale carry radial-geometric information?',
             fontsize=11, y=1.02)
plt.tight_layout()
plt.savefig(OUT / "test5_tsa_harmonic_scale.png", dpi=200, bbox_inches='tight')
plt.savefig(OUT / "test5_tsa_harmonic_scale.svg", bbox_inches='tight')
plt.close(fig)

print(f"\nSaved: test5_tsa_harmonic_scale.png/svg")
print(f"Saved: test5_tsa_harmonic_scale_summary.csv")
print()
print("DONE")
