"""
Extended Data Figure 6 — Multi-scale helix vs radial-accordion test on real Dip-C 3D coordinates.

REAL DATA ONLY. No rng outside permutation null. Reviewer-runnable.

Triggered by Wolf 2026-05-02: the harmonic-coupling findings (cross-scale amplitude
correlation r=+0.55 L1↔L2, +0.77 L2↔L3 in K562 phase) make the helix model worth
testing more carefully, since a real helix would imprint scale-invariantly on the
3D structure. Existing helix proposals in the literature (e.g., Boettiger 2016
super-resolution; Halverson 2014 polymer simulation; Dixon-lab structural
proposals) further motivate a rigorous test rather than a single-scale one.

Source data:
  - 14 GM12878 Dip-C 3D structures (Tan 2018 GSE117874): per-bead (chrom_hap, pos, x, y, z)
    at 100-kb resolution. (`polychrom simulation/dipc_data/extracted/gm12878_*.3dg.gz`)
  - GM12878 hg19 master annotation (`go_rerun_v2_2026-04-30/master_annotation_GM12878_hg19.csv`)
    with cwt_phase + dom_period_Mbp at 100-kb resolution (5,776 bins, 22 autosomes).

Pipeline (no rng except permutation null):
  1. Per cell: cos(φ)-weighted cell-axis â (cell-global speckle-lamina direction;
     positive â = speckle pole). Validates radial-accordion claim with no circularity.
  2. Per cell: PCA basis (b̂, ĉ) in the plane orthogonal to â. b̂ is the SLINKY-test
     direction — uncorrelated with cos(φ) by construction, so any z_b oscillation
     vs cos(φ) is geometric, not tautological.
  3. Per chromosome × per scale (L1, L2 = L1/3, L3 = L1/5): fit
     signal ~ A·cos(2π·pos/L_Mb + ψ) + c separately for signal ∈ {r, z_perp, θ_unwrapped}.
     Read amplitude A and fit R² as the per-chromosome × per-scale diagnostic.

Diagnostics interpretation:
  - Radial accordion: r-amplitude high at L1, drops at L2/L3; z and θ amplitudes
    near zero in well-fit cosine sense.
  - Pure helix: θ-amplitude high at all scales with consistent ω = (Δθ along chrom)/(2π).
  - Slinky: z_perp-amplitude rivals r-amplitude at L1.
  - Cross-scale invariance: helix would preserve amplitude RATIOS across scales;
    radial accordion concentrates at L1.

Headline real-data result (from `helix_multiscale_data/per_chrom_per_scale_FULL.csv`,
2,772 rows = 14 cells × 22 chr × 3 scales × 3 signals):
  - At L1: median r-amplitude 0.040 vs median z_perp-amplitude 0.034 → similar magnitude;
    median R² of r-cosine = 0.008 (per-chromosome; cell-aggregate R² = 0.87 across 14 cells).
  - θ has high raw amplitude (0.32) but fits cosine poorly (R² < 0.005) → no helical winding.
  - Cross-scale: r-amplitude collapses 4× from L1 → L3 (0.040 → 0.009), θ-amplitude
    collapses 3× (0.32 → 0.10) — ratios are NOT scale-invariant.
  - Cell-level: 14/14 cells positive R²_radial (mean 0.866 ± 0.072), r_cl < 0.05
    in 14/14 cells (mean 0.009), R²_slinky_perp < 0.10 in 12/14 (mean 0.10, two outliers).

Conclusion (subject to Wolf review): the data is consistent with a radial accordion
concentrated at the fundamental scale, NOT with a scale-invariant helix. The helix is
rejected at all three scales by both per-cell r_cl and per-chromosome cosine-fit
amplitude tests.
"""
from __future__ import annotations
import sys, os
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

sys.path.insert(0, str(Path(__file__).resolve().parent))
from _style import apply_style, panel_letter, savefig_both, WONG, SEMANTIC

BUNDLE_ROOT = Path(__file__).resolve().parents[2]
DATA = BUNDLE_ROOT / '04_supplementary_data' / 'figure_data'
PER_CELL = DATA / 'per_cell_L1_results.csv'
PER_CHROM_SCALE = DATA / 'per_chrom_per_scale_FULL.csv'

OUT_BASENAMES = [
    # canonical name + alias matching the manuscript .tex \includegraphics path
    # (.tex retains the legacy `extended_data_figure_6_model_allelic` filename
    # because v40 ED Fig 6 was a model+allelic-asymmetry panel; the v95 → v98
    # regen replaced its content with the multi-scale helix-rejection test
    # documented in the v98 caption.)
    BUNDLE_ROOT / '02_figures' / 'extended_data' / 'extended_data_figure_6_helix_multiscale',
    BUNDLE_ROOT / '02_figures' / 'extended_data' / 'extended_data_figure_6_model_allelic',
    BUNDLE_ROOT / '03_figures_to_edit' / 'extended_data' / 'extended_data_figure_6_helix_multiscale',
    BUNDLE_ROOT / '03_figures_to_edit' / 'extended_data' / 'extended_data_figure_6_model_allelic',
]
DRAFT_BASENAME = BUNDLE_ROOT / '05_code' / 'figure_regen' / '_drafts' / 'ed6_helix_multiscale_v1'

# ---------- Load ----------
df_cell = pd.read_csv(PER_CELL)
df_pcs = pd.read_csv(PER_CHROM_SCALE)
print(f"per-cell L1: {len(df_cell)} cells")
print(f"per-chrom × per-scale: {len(df_pcs)} rows")

scales = ['L1', 'L2', 'L3']
signals = ['r', 'z_perp', 'theta']
sig_label = {'r': 'radial r', 'z_perp': 'axial-perp z_b', 'theta': 'azimuthal θ'}
sig_color = {'r': WONG['blue'], 'z_perp': WONG['vermil'], 'theta': WONG['gold']}

# ---------- Figure ----------
apply_style()
fig = plt.figure(figsize=(15.0, 11.5))
gs = fig.add_gridspec(2, 3, height_ratios=[1.0, 1.0], width_ratios=[1.0, 1.0, 1.05],
                      hspace=0.55, wspace=0.40, left=0.055, right=0.985, top=0.93, bottom=0.065)
ax_a = fig.add_subplot(gs[0, 0])  # per-cell L1 diagnostics
ax_b = fig.add_subplot(gs[0, 1])  # cross-scale amplitude collapse
ax_c = fig.add_subplot(gs[0, 2])  # per-chrom R² distribution by scale × signal
ax_d = fig.add_subplot(gs[1, 0])  # PCA of per-cell wave-geometry profiles
ax_e = fig.add_subplot(gs[1, 1])  # cross-scale amplitude correlation (harmonic coupling)
ax_f = fig.add_subplot(gs[1, 2])  # cell-aggregate vs per-chromosome R² comparison

# === Panel a: per-cell L1 diagnostics ===
xa = np.arange(len(df_cell))
df_cell_sort = df_cell.sort_values('R2_radial').reset_index(drop=True)
xa = np.arange(len(df_cell_sort))
w = 0.27
ax_a.bar(xa - w, df_cell_sort['R2_radial'], width=w, color=WONG['blue'], label='R² radial', edgecolor='white', lw=0.4)
ax_a.bar(xa,     df_cell_sort['R2_slinky_perp'], width=w, color=WONG['vermil'], label='R² slinky_perp', edgecolor='white', lw=0.4)
ax_a.bar(xa + w, df_cell_sort['r_cl'], width=w, color=WONG['gold'], label='r_cl (helix)', edgecolor='white', lw=0.4)
ax_a.axhline(0.5, color='0.5', lw=0.5, ls=':')
ax_a.set_xticks(xa)
ax_a.set_xticklabels([c.replace('gm12878_','c') for c in df_cell_sort['cell']], rotation=0, fontsize=7.5)
ax_a.set_ylabel('per-cell diagnostic value')
ax_a.legend(loc='upper left', fontsize=8)
ax_a.set_title(f'L1: 14/14 cells radial; helix rejected',
               loc='left', pad=8, fontsize=10)
ax_a.set_ylim(0, 1.0)
panel_letter(fig, ax_a, 'a')

# === Panel b: cross-scale amplitude collapse ===
amp_data = {sig: [] for sig in signals}
for sc in scales:
    for sig in signals:
        sub = df_pcs[(df_pcs['scale']==sc)&(df_pcs['signal']==sig)]
        amp_data[sig].append(sub['amplitude'].values)
positions_b = np.arange(len(scales)) * 4.0
for i, sig in enumerate(signals):
    bp = ax_b.boxplot(amp_data[sig], positions=positions_b + i*0.85, widths=0.7,
                      patch_artist=True, showfliers=False, zorder=3)
    for box in bp['boxes']:
        box.set_facecolor(sig_color[sig]); box.set_alpha(0.6); box.set_edgecolor('0.2'); box.set_linewidth(0.5)
    for med in bp['medians']:
        med.set_color('0.1'); med.set_linewidth(1.0)
ax_b.set_xticks(positions_b + 0.85)
ax_b.set_xticklabels(scales, fontsize=9)
ax_b.set_xlabel('scale (T₁, T₁/3, T₁/5)')
ax_b.set_ylabel('per-chromosome cosine-fit amplitude')
ax_b.set_yscale('log')
# Legend
handles_b = [mpatches.Patch(color=sig_color[s], label=sig_label[s]) for s in signals]
ax_b.legend(handles=handles_b, loc='upper left', fontsize=8, framealpha=0.92)
ax_b.set_title('Cross-scale amplitude collapse',
               loc='left', pad=8, fontsize=10)
panel_letter(fig, ax_b, 'b')

# === Panel c: per-chrom R² distribution by scale × signal ===
positions_c = np.arange(len(scales)) * 4.0
for i, sig in enumerate(signals):
    R2s = [df_pcs[(df_pcs['scale']==sc)&(df_pcs['signal']==sig)]['R2'].values for sc in scales]
    bp = ax_c.boxplot(R2s, positions=positions_c + i*0.85, widths=0.7,
                      patch_artist=True, showfliers=False)
    for box in bp['boxes']:
        box.set_facecolor(sig_color[sig]); box.set_alpha(0.6); box.set_edgecolor('0.2'); box.set_linewidth(0.5)
    for med in bp['medians']:
        med.set_color('0.1'); med.set_linewidth(1.0)
ax_c.set_xticks(positions_c + 0.85)
ax_c.set_xticklabels(scales, fontsize=9)
ax_c.set_xlabel('scale')
ax_c.set_ylabel('per-chromosome cosine-fit R²')
ax_c.legend(handles=handles_b, loc='upper right', fontsize=8)
ax_c.set_title('Cosine-fit R² per chrom × scale',
               loc='left', pad=8, fontsize=10)
panel_letter(fig, ax_c, 'c')

# === Panel d: PCA of per-cell wave-geometry profiles ===
# Build feature matrix: per cell, [R²_radial, R²_slinky_perp, r_cl,
#  median amp_r at L1, L2, L3, median amp_θ at L1, L2, L3, median amp_z at L1, L2, L3] = 12 features
features = []
cell_order = []
for cell in df_cell['cell']:
    sub = df_pcs[df_pcs['cell']==cell]
    feat = [
        df_cell.loc[df_cell['cell']==cell, 'R2_radial'].iloc[0],
        df_cell.loc[df_cell['cell']==cell, 'R2_slinky_perp'].iloc[0],
        df_cell.loc[df_cell['cell']==cell, 'r_cl'].iloc[0],
    ]
    for sc in scales:
        for sig in signals:
            v = sub[(sub['scale']==sc)&(sub['signal']==sig)]['amplitude'].median()
            feat.append(v if not np.isnan(v) else 0.0)
    features.append(feat)
    cell_order.append(cell)
features = np.array(features)
# z-score columns
features_z = (features - features.mean(axis=0))/features.std(axis=0, ddof=1)
features_z = np.nan_to_num(features_z, 0)
# PCA via SVD
u, s, vh = np.linalg.svd(features_z, full_matrices=False)
pcs = features_z @ vh.T  # PC scores
# Variance explained
var_explained = (s**2) / (s**2).sum() * 100
ax_d.scatter(pcs[:,0], pcs[:,1], s=80, c=df_cell['R2_radial'].values, cmap='viridis',
             edgecolor='black', linewidth=0.5, zorder=3)
for i, cell in enumerate(cell_order):
    ax_d.annotate(cell.replace('gm12878_','c'), (pcs[i,0], pcs[i,1]),
                  xytext=(5, 5), textcoords='offset points', fontsize=7, color='0.2')
ax_d.set_xlabel(f'PC1 ({var_explained[0]:.1f}% var)')
ax_d.set_ylabel(f'PC2 ({var_explained[1]:.1f}% var)')
ax_d.axhline(0, color='0.6', lw=0.5)
ax_d.axvline(0, color='0.6', lw=0.5)
ax_d.set_title('PCA on per-cell wave-geometry feature vector (12 features)',
               loc='left', pad=8, fontsize=10)
panel_letter(fig, ax_d, 'd')

# === Panel e: cross-scale amplitude correlation (harmonic coupling tie-in) ===
# For each chromosome (averaged across cells), compute amp_r at L1 vs L2 vs L3
chrom_amp = df_pcs[df_pcs['signal']=='r'].groupby(['chrom','scale'])['amplitude'].mean().unstack()
chrom_amp = chrom_amp[['L1','L2','L3']]
# Plot L1 vs L2 and L2 vs L3
ax_e.scatter(chrom_amp['L1'], chrom_amp['L2'], s=60, color=WONG['blue'], edgecolor='white', linewidth=0.4,
             label='L1 → L2', zorder=3)
ax_e.scatter(chrom_amp['L2'], chrom_amp['L3'], s=60, color=WONG['vermil'], marker='s', edgecolor='white', linewidth=0.4,
             label='L2 → L3', zorder=3)
# trend lines
from scipy.stats import pearsonr
r12, p12 = pearsonr(chrom_amp['L1'], chrom_amp['L2'])
r23, p23 = pearsonr(chrom_amp['L2'], chrom_amp['L3'])
# Annotate
ax_e.text(0.04, 0.97,
          f'L1→L2 amp coupling:  r = {r12:+.2f},  p = {p12:.0e}\n'
          f'L2→L3 amp coupling:  r = {r23:+.2f},  p = {p23:.0e}',
          transform=ax_e.transAxes, fontsize=8, va='top', ha='left',
          bbox=dict(boxstyle='round,pad=0.3', facecolor='white', edgecolor='0.7', lw=0.4, alpha=0.95))
# y=x reference
xs_e = np.linspace(chrom_amp.values.min(), chrom_amp.values.max(), 100)
ax_e.plot(xs_e, xs_e, color='0.6', ls='--', lw=0.7, label='y = x')
ax_e.set_xlabel('amp at coarser scale')
ax_e.set_ylabel('amp at finer scale')
ax_e.legend(loc='lower right', fontsize=8)
ax_e.set_title('Cross-scale amplitude coupling on r-signal (harmonic-coupling tie-in)',
               loc='left', pad=8, fontsize=10)
panel_letter(fig, ax_e, 'e')

# === Panel f: cell-aggregate vs per-chromosome R² ===
# show that the R²_radial=0.87 at cell-aggregate is much higher than per-chromosome R²s
agg_R2 = df_cell['R2_radial'].values
# per-chrom R² for r-signal at L1 averaged per cell
per_chrom_R2 = df_pcs[(df_pcs['scale']=='L1')&(df_pcs['signal']=='r')].groupby('cell')['R2'].median().values
# Match order
cell_to_chr = df_pcs[(df_pcs['scale']=='L1')&(df_pcs['signal']=='r')].groupby('cell')['R2'].median().to_dict()
per_chrom_R2_ordered = np.array([cell_to_chr.get(c, 0) for c in df_cell['cell']])
xf = np.arange(len(df_cell))
df_cell_sort_f = df_cell.sort_values('R2_radial').reset_index(drop=True)
order = df_cell.sort_values('R2_radial').index.tolist()
agg_sorted = agg_R2[order]
per_chr_sorted = per_chrom_R2_ordered[order]
ax_f.bar(xf - 0.20, agg_sorted, width=0.4, color=WONG['gold'],
         label='cell-aggregate R²_radial (whole-cell phi-binned fit)', edgecolor='white', lw=0.4)
ax_f.bar(xf + 0.20, per_chr_sorted, width=0.4, color=WONG['grey'],
         label='per-chrom L1 r-cosine median R² (genomic-position fit)', edgecolor='white', lw=0.4)
ax_f.set_xticks(xf)
ax_f.set_xticklabels([c.replace('gm12878_','c') for c in df_cell_sort_f['cell']], rotation=90, fontsize=6.5)
ax_f.set_ylabel('R²')
ax_f.legend(loc='upper left', fontsize=7.5)
ax_f.set_title('Two-frame view: cell-aggregate vs per-chromosome — different tests, both real',
               loc='left', pad=8, fontsize=9.5)
panel_letter(fig, ax_f, 'f')

# Save
DRAFT_BASENAME.parent.mkdir(parents=True, exist_ok=True)
savefig_both(fig, DRAFT_BASENAME)
print(f"[wrote draft] {DRAFT_BASENAME}.png + .svg")
if os.environ.get('PROMOTE') == '1':
    for base in OUT_BASENAMES:
        savefig_both(fig, base)
        print(f"[promoted to canonical] {base}.png + .svg")
plt.close(fig)
print("\nKey real-data findings:")
print(f"  L1 radial r median amp = {df_pcs[(df_pcs['scale']=='L1')&(df_pcs['signal']=='r')]['amplitude'].median():.4f}")
print(f"  L1 slinky z_perp median amp = {df_pcs[(df_pcs['scale']=='L1')&(df_pcs['signal']=='z_perp')]['amplitude'].median():.4f}")
print(f"  L1 theta amp = {df_pcs[(df_pcs['scale']=='L1')&(df_pcs['signal']=='theta')]['amplitude'].median():.4f}")
print(f"  Cross-scale r amp ratio L1/L2 = {df_pcs[(df_pcs['scale']=='L1')&(df_pcs['signal']=='r')]['amplitude'].median()/df_pcs[(df_pcs['scale']=='L2')&(df_pcs['signal']=='r')]['amplitude'].median():.2f}")
print(f"  Cross-scale r amp ratio L2/L3 = {df_pcs[(df_pcs['scale']=='L2')&(df_pcs['signal']=='r')]['amplitude'].median()/df_pcs[(df_pcs['scale']=='L3')&(df_pcs['signal']=='r')]['amplitude'].median():.2f}")
