"""
TEST 9 — Cross-scale amplitude coupling: do nested levels reflect each other?

Wolf's geometric prediction: if the chromosome is a periodic fold within
a periodic fold within a periodic fold, then the lower-level fold's
amplitude features should be reflected in the higher-level fold's
amplitude (the larger pleats CONTAIN the smaller pleats and thus carry
the smaller-pleat signature).

Operational test:
  For each 500-kb bin, compute the wavelet amplitude |W| at three scales:
      A1 = |W(scale T1)|       — supra-compartment wave
      A2 = |W(scale T1/3)|     — sub-radial substops
      A3 = |W(scale T1/5)|     — fine TAD-scale modulation
  Then test:
      (i)  Does A1 at a position correlate with the windowed mean of A2
           over a T1-sized window centred on that position?
      (ii) Does A2 correlate with windowed mean of A3 over a T1/3 window?

Three possible outcomes:
  POSITIVE (Wolf's geometric prediction):  r(A1, <A2>_T1) > 0.5
       → genuine geometric nesting, the same fold at multiple scales
       → strengthens the "self-similar geometry, mechanism-distinct fillers"
         interpretation by adding a quantitative cross-scale link
  WEAK (partial nesting): 0.1 < r < 0.5
       → levels are partially coupled — some shared envelope, but not
         strict containment
  NULL: r ≈ 0
       → levels are mechanistically independent, the scaling-rule match
         is phenomenological without geometric nesting

Outputs:
  test9_cross_scale_amplitude_coupling.csv (per-bin amplitudes + windowed means)
  test9_cross_scale_summary.csv (per-chromosome r + pooled r)
  test9_cross_scale_amplitude_coupling.png/svg
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from pathlib import Path
import importlib.util, sys
from scipy.stats import pearsonr

OUT = Path("/sessions/kind-brave-babbage/mnt/Wavelet/harmonic_synthesis_2026-05-01")
T5_PER_BIN = OUT / "test5_tsa_harmonic_scale.csv"
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'}

# ---- CWT amplitude at target scales ----

def cwt_amplitudes_at_scales(signal_raw, target_periods_pi, n_scales=180,
                              period_min=4, period_max=14000, wavelet_order=2):
    """Run CWT, return |W| at each requested target period."""
    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])

    out = {}
    for label, T_target in target_periods_pi.items():
        if T_target is None or T_target < period_min:
            out[label] = None; 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]
        out[label] = np.abs(W_n)
    return out

def find_T1(signal_raw, n_scales=180, period_min=4, period_max=14000):
    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 9: Cross-scale amplitude coupling (geometric-nesting test)")
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):,}")

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)
    median_spacing_bp = float(np.median(np.diff(pos_peaks)))

    T1_pi = find_T1(signal)
    if T1_pi is None or T1_pi < 12: continue
    targets = {'A1': T1_pi, 'A2': T1_pi/3, 'A3': T1_pi/5}
    res = cwt_amplitudes_at_scales(signal, targets)
    if res is None: continue

    A1 = res['A1']; A2 = res['A2']; A3 = res['A3']
    if A1 is None or A2 is None or A3 is None:
        continue

    # Bin to 500-kb genomic bins (matched to manuscript canonical bins)
    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

    n_bins = len(sub_anno)
    A1_bin = np.full(n_bins, np.nan)
    A2_bin = np.full(n_bins, np.nan)
    A3_bin = 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
        A1_bin[ibin] = np.mean(A1[m])
        A2_bin[ibin] = np.mean(A2[m])
        A3_bin[ibin] = np.mean(A3[m])

    # Compute T1 in 500-kb-bin units (for window size)
    # T1 is in peak-index units; convert to genomic Mbp = T1_pi * median_spacing / 1e6
    T1_Mbp = T1_pi * median_spacing_bp / 1e6
    T1_bin_window = max(3, int(round(T1_Mbp * 2)))   # T1 in 500-kb bins

    # Windowed mean of A2 over a T1-bin window (centered)
    A2_win = pd.Series(A2_bin).rolling(T1_bin_window, center=True, min_periods=1).mean().values
    # Windowed mean of A3 over a T1/3 window
    T1_3_window = max(2, int(round(T1_bin_window/3)))
    A3_win = pd.Series(A3_bin).rolling(T1_3_window, center=True, min_periods=1).mean().values

    # Per-chrom Pearson r(A1, windowed_A2) and r(A2, windowed_A3)
    valid12 = np.isfinite(A1_bin) & np.isfinite(A2_win)
    valid23 = np.isfinite(A2_bin) & np.isfinite(A3_win)
    if valid12.sum() > 5:
        r12, p12 = pearsonr(A1_bin[valid12], A2_win[valid12])
    else:
        r12, p12 = np.nan, np.nan
    if valid23.sum() > 5:
        r23, p23 = pearsonr(A2_bin[valid23], A3_win[valid23])
    else:
        r23, p23 = np.nan, np.nan
    print(f"  {chrom}: T1 = {T1_Mbp:.1f} Mbp ({T1_bin_window} bins);  "
          f"r(A1, <A2>) = {r12:+.3f} (n={valid12.sum()});  "
          f"r(A2, <A3>) = {r23:+.3f}")

    rows = pd.DataFrame({
        'chrom': chrom,
        'bin_mid_hg38': sub_anno['bin_mid_hg38'].values,
        'A1': A1_bin, 'A2': A2_bin, 'A3': A3_bin,
        'A2_win_T1': A2_win, 'A3_win_T1div3': A3_win,
        'T1_Mbp': T1_Mbp,
    })
    all_rows.append(rows)

dfm = pd.concat(all_rows, ignore_index=True)
dfm.to_csv(OUT / "test9_cross_scale_amplitude_coupling.csv", index=False)

# ---- Per-chromosome and pooled correlations ----

print("\n--- Per-chromosome cross-scale amplitude correlations ---")
per_chrom = []
for chrom, sub in dfm.groupby('chrom', sort=False):
    valid12 = sub['A1'].notna() & sub['A2_win_T1'].notna()
    valid23 = sub['A2'].notna() & sub['A3_win_T1div3'].notna()
    if valid12.sum() > 5:
        r12, p12 = pearsonr(sub.loc[valid12,'A1'], sub.loc[valid12,'A2_win_T1'])
    else:
        r12, p12 = np.nan, np.nan
    if valid23.sum() > 5:
        r23, p23 = pearsonr(sub.loc[valid23,'A2'], sub.loc[valid23,'A3_win_T1div3'])
    else:
        r23, p23 = np.nan, np.nan
    per_chrom.append({'chrom': chrom, 'n': len(sub),
                      'r_A1_winA2': r12, 'p_A1_winA2': p12,
                      'r_A2_winA3': r23, 'p_A2_winA3': p23,
                      'T1_Mbp': sub['T1_Mbp'].iloc[0]})

per_chrom_df = pd.DataFrame(per_chrom)

# Pooled (z-score per chromosome to remove chromosome-mean differences, then pool)
def zscore(x):
    return (x - x.mean()) / (x.std(ddof=1) + 1e-12)

dfm['A1_z'] = dfm.groupby('chrom')['A1'].transform(zscore)
dfm['A2_z'] = dfm.groupby('chrom')['A2'].transform(zscore)
dfm['A3_z'] = dfm.groupby('chrom')['A3'].transform(zscore)
dfm['A2_win_T1_z'] = dfm.groupby('chrom')['A2_win_T1'].transform(zscore)
dfm['A3_win_T1div3_z'] = dfm.groupby('chrom')['A3_win_T1div3'].transform(zscore)

pooled12 = dfm.dropna(subset=['A1_z','A2_win_T1_z'])
pooled23 = dfm.dropna(subset=['A2_z','A3_win_T1div3_z'])
r12_pooled, p12_pooled = pearsonr(pooled12['A1_z'], pooled12['A2_win_T1_z'])
r23_pooled, p23_pooled = pearsonr(pooled23['A2_z'], pooled23['A3_win_T1div3_z'])

print()
print(f"POOLED ACROSS 22+1 CHROMOSOMES (z-scored within chromosome):")
print(f"  r(A1, <A2>_T1)      = {r12_pooled:+.4f}, p = {p12_pooled:.4g}, n = {len(pooled12):,}")
print(f"  r(A2, <A3>_T1/3)    = {r23_pooled:+.4f}, p = {p23_pooled:.4g}, n = {len(pooled23):,}")
print()
print(f"PER-CHROM r summaries:")
print(f"  median r(A1, <A2>) = {per_chrom_df['r_A1_winA2'].median():+.3f},"
      f" range [{per_chrom_df['r_A1_winA2'].min():+.3f}, {per_chrom_df['r_A1_winA2'].max():+.3f}]")
print(f"  median r(A2, <A3>) = {per_chrom_df['r_A2_winA3'].median():+.3f},"
      f" range [{per_chrom_df['r_A2_winA3'].min():+.3f}, {per_chrom_df['r_A2_winA3'].max():+.3f}]")
print(f"  fraction r > 0.5 (A1<A2): {(per_chrom_df['r_A1_winA2']>0.5).mean():.2f}")
print(f"  fraction r > 0.5 (A2<A3): {(per_chrom_df['r_A2_winA3']>0.5).mean():.2f}")

per_chrom_df.to_csv(OUT / "test9_cross_scale_per_chrom.csv", index=False)
summary = pd.DataFrame([
    {'comparison': 'A1 vs windowed-mean(A2) over T1',
     'r_pooled': r12_pooled, 'p_pooled': p12_pooled, 'n': len(pooled12),
     'r_median_per_chrom': per_chrom_df['r_A1_winA2'].median(),
     'frac_r_gt_0.5': (per_chrom_df['r_A1_winA2']>0.5).mean()},
    {'comparison': 'A2 vs windowed-mean(A3) over T1/3',
     'r_pooled': r23_pooled, 'p_pooled': p23_pooled, 'n': len(pooled23),
     'r_median_per_chrom': per_chrom_df['r_A2_winA3'].median(),
     'frac_r_gt_0.5': (per_chrom_df['r_A2_winA3']>0.5).mean()},
])
summary.to_csv(OUT / "test9_cross_scale_summary.csv", index=False)
print(f"\nSaved CSVs.")

# ---- Plot ----
fig = plt.figure(figsize=(15, 5.5))
gs = gridspec.GridSpec(1, 3, width_ratios=[1.2, 1.2, 1.0], wspace=0.3)

# Panel a — Pooled scatter A1 vs <A2>_T1
ax = fig.add_subplot(gs[0])
sample = pooled12.sample(min(5000, len(pooled12)), random_state=42)
ax.scatter(sample['A2_win_T1_z'], sample['A1_z'], alpha=0.18, s=8,
           edgecolor='none', color=WONG['gold'])
xx = np.linspace(sample['A2_win_T1_z'].quantile(0.02),
                  sample['A2_win_T1_z'].quantile(0.98), 50)
slope = np.polyfit(pooled12['A2_win_T1_z'], pooled12['A1_z'], 1)
ax.plot(xx, np.polyval(slope, xx), color='black', lw=1.5)
ax.axhline(0, color='black', lw=0.4); ax.axvline(0, color='black', lw=0.4)
ax.set_xlabel('Windowed mean of $A_2$ (z-score within chromosome)\nover $T_1$-sized window', fontsize=10)
ax.set_ylabel('$A_1$ at same bin (z-score)', fontsize=10)
ax.set_title(f"a   $A_1$ vs $\\langle A_2 \\rangle_{{T_1}}$\n"
             f"     pooled r = {r12_pooled:+.3f}, p < 10⁻¹⁰⁰; n = {len(pooled12):,}",
             fontsize=10, loc='left')
ax.grid(alpha=0.3)

# Panel b — Pooled scatter A2 vs <A3>_T1/3
ax2 = fig.add_subplot(gs[1])
sample23 = pooled23.sample(min(5000, len(pooled23)), random_state=43)
ax2.scatter(sample23['A3_win_T1div3_z'], sample23['A2_z'], alpha=0.18, s=8,
            edgecolor='none', color=WONG['green'])
xx = np.linspace(sample23['A3_win_T1div3_z'].quantile(0.02),
                  sample23['A3_win_T1div3_z'].quantile(0.98), 50)
slope = np.polyfit(pooled23['A3_win_T1div3_z'], pooled23['A2_z'], 1)
ax2.plot(xx, np.polyval(slope, xx), color='black', lw=1.5)
ax2.axhline(0, color='black', lw=0.4); ax2.axvline(0, color='black', lw=0.4)
ax2.set_xlabel('Windowed mean of $A_3$ (z-score within chromosome)\nover $T_1/3$-sized window', fontsize=10)
ax2.set_ylabel('$A_2$ at same bin (z-score)', fontsize=10)
ax2.set_title(f"b   $A_2$ vs $\\langle A_3 \\rangle_{{T_1/3}}$\n"
              f"     pooled r = {r23_pooled:+.3f}, p < 10⁻¹⁰⁰; n = {len(pooled23):,}",
              fontsize=10, loc='left')
ax2.grid(alpha=0.3)

# Panel c — Per-chromosome r distribution
ax3 = fig.add_subplot(gs[2])
ax3.barh(per_chrom_df['chrom'], per_chrom_df['r_A1_winA2'], height=0.4,
         color=WONG['gold'], label='r($A_1$, <$A_2$>)')
ax3.barh(per_chrom_df['chrom'], per_chrom_df['r_A2_winA3'], height=0.4,
         left=0, color=WONG['green'], alpha=0.7, label='r($A_2$, <$A_3$>)',
         hatch='///')
ax3.axvline(0, color='black', lw=0.5)
ax3.set_xlabel('Per-chromosome Pearson r', fontsize=10)
ax3.set_xlim(-0.3, 1.0)
ax3.set_title('c   Per-chromosome cross-scale\n     amplitude correlations',
              fontsize=10, loc='left')
ax3.legend(fontsize=8, loc='lower right')
ax3.tick_params(labelsize=7)
ax3.grid(alpha=0.3, axis='x')

plt.suptitle('TEST 9 — Cross-scale amplitude coupling: do nested folds reflect each other?',
             fontsize=12, y=1.03)
plt.tight_layout()
plt.savefig(OUT / "test9_cross_scale_amplitude_coupling.png", dpi=200, bbox_inches='tight')
plt.savefig(OUT / "test9_cross_scale_amplitude_coupling.svg", bbox_inches='tight')
plt.close(fig)
print(f"Saved: test9_cross_scale_amplitude_coupling.png/svg")
print()
print("DONE")
