"""
Extended Data Figure 1 — Three orthogonal accessibility assays converge on the same
                          supra-compartment periodicity (MCF-7).

REAL DATA ONLY. No rng. Reviewer-runnable.

Source CSVs (BIOMNI 2026-04-04 cross-assay deliverable):
  - cross_assay_cwt_results.csv     5 chromosomes × 3 assays (DNase, ATAC, H3K4me3),
                                    per-chromosome dominant period + sig95 + AR(1) α
  - cross_assay_coherence_all.csv   5 chromosomes × 3 assay-pair combinations,
                                    cross-wavelet coherence at the dominant-period band

The cross-assay sweep is scoped to 5 representative MCF-7 chromosomes (chrX, chr1, chr2,
chr11, chr17) — selected for span of length, gene density, and HER2-amplicon disruption
on chr17 as a stress test.

Headline numbers reproduced (must match manuscript caption):
  - chrX dominant period: DNase 30.8 Mbp / ATAC 28.3 Mbp / H3K4me3 30.8 Mbp (28-31 Mbp band ✓)
  - chrX cross-wavelet coherence at dominant band: DNase-ATAC R² = 0.877; DNase-H3K4me3
    R² = 0.896; ATAC-H3K4me3 R² = 0.879 — all > 0.85 in matched-cell-line cross-assay.
  - chrX phase difference at dominant band: 6-15° (essentially zero, i.e. assays are
    in-phase at the supra-compartment scale).
  - chr1 dominant period heterogeneity (DNase 43.6 Mbp vs ATAC 113.1 Mbp vs H3K4me3 113.1 Mbp)
    illustrates that long chromosomes admit multiple dominant-period candidates above the
    AR(1) null; the rank-1 GWS peak depends on assay.

Replaces 2026-04-21 v73 synthetic placeholder
(`figure_regen_v73/scripts/extended_data_figure_1_cross_assay.py`, rng seed 7,
 baked per-assay sig95 means).
"""
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-relative paths. Reviewer-runnable as `cd <bundle>; python 05_code/figure_regen/<name>.py`.
BUNDLE_ROOT = Path(__file__).resolve().parents[2]
DATA_DIR = BUNDLE_ROOT / '04_supplementary_data' / 'figure_data'
RESULTS_CSV = DATA_DIR / 'cross_assay_cwt_results.csv'
COHERENCE_CSV = DATA_DIR / 'cross_assay_coherence_all.csv'

OUT_BASENAMES = [
    BUNDLE_ROOT / '02_figures' / 'extended_data' / 'extended_data_figure_1_cross_assay',
    BUNDLE_ROOT / '03_figures_to_edit' / 'extended_data' / 'extended_data_figure_1_cross_assay',
]
DRAFT_BASENAME = BUNDLE_ROOT / '05_code' / 'figure_regen' / '_drafts' / 'ed1_v1'

# ---------- Load ----------
results = pd.read_csv(RESULTS_CSV)
coherence = pd.read_csv(COHERENCE_CSV)

# Verify
chrX = results[results['chr']=='chrX']
chrX_periods = dict(zip(chrX['assay'], chrX['dominant_period_Mbp']))
print(f"[verify] chrX periods: DNase {chrX_periods['dnase']:.2f} | ATAC {chrX_periods['atac']:.2f} | H3K4me3 {chrX_periods['h3k4me3']:.2f} Mbp")
print(f"[verify] chrX coherence R²:")
for _, row in coherence[coherence['chr']=='chrX'].iterrows():
    print(f"  {row['pair']}: R² = {row['mean_r2_at_assay1_band']:.3f}, phase diff = {row['phase_at_assay1_band_deg']:.1f}°")
print(f"[verify] {results['chr'].nunique()} chromosomes covered: {sorted(results['chr'].unique())}")

# ---------- Figure ----------
apply_style()
fig = plt.figure(figsize=(13.0, 4.6))
gs = fig.add_gridspec(1, 3, width_ratios=[1.0, 1.0, 1.45],
                      wspace=0.38, left=0.06, right=0.985, top=0.86, bottom=0.16)
ax_a = fig.add_subplot(gs[0, 0])
ax_b = fig.add_subplot(gs[0, 1])
ax_c = fig.add_subplot(gs[0, 2])

ASSAYS = ['dnase', 'atac', 'h3k4me3']
LABELS = ['DNase', 'ATAC', 'H3K4me3']
COLORS = [WONG['blue'], WONG['green'], WONG['vermil']]

# === Panel a: chrX three-assay convergence ===
periods = [chrX_periods[a] for a in ASSAYS]
xa = np.arange(3)
ax_a.bar(xa, periods, color=COLORS, edgecolor='white', linewidth=0.5)
# 28-31 Mbp band shaded
ax_a.axhspan(28, 31, alpha=0.18, color=WONG['gold'], zorder=0,
             label='28–31 Mbp band')
for x, p in zip(xa, periods):
    ax_a.text(x, p + 0.5, f'{p:.1f}', ha='center', va='bottom',
              fontsize=9, fontweight='bold')
ax_a.set_xticks(xa); ax_a.set_xticklabels(LABELS, fontsize=9)
ax_a.set_ylabel('dominant period (Mbp)')
ax_a.set_ylim(0, max(periods)*1.18)
ax_a.legend(loc='upper right', fontsize=8)
ax_a.set_title('chrX: 3 assays converge at ~30 Mbp',
               loc='left', pad=8, fontsize=10)
panel_letter(fig, ax_a, 'a')

# === Panel b: cross-wavelet coherence on chrX ===
chrX_coh = coherence[coherence['chr']=='chrX'].copy()
xb = np.arange(len(chrX_coh))
r2_vals = chrX_coh['mean_r2_at_assay1_band'].values
phase_vals = chrX_coh['phase_at_assay1_band_deg'].values
pair_labels = [p.replace('-', ' ↔ ').replace('dnase','DNase').replace('atac','ATAC').replace('h3k4me3','H3K4me3')
               for p in chrX_coh['pair'].values]
ax_b.bar(xb, r2_vals, color=WONG['orange'], edgecolor='white', linewidth=0.5)
for x, r, ph in zip(xb, r2_vals, phase_vals):
    ax_b.text(x, r + 0.012, f'R² = {r:.3f}', ha='center', va='bottom',
              fontsize=8, fontweight='bold')
    ax_b.text(x, r/2, f'Δφ = {ph:.1f}°', ha='center', va='center',
              fontsize=7, color='white', fontweight='bold')
ax_b.set_xticks(xb); ax_b.set_xticklabels(pair_labels, rotation=15, fontsize=8.5)
ax_b.set_ylabel('cross-wavelet R² at dominant band')
ax_b.set_ylim(0, 1.05)
ax_b.set_title('Cross-wavelet coherence: R² > 0.87, near-zero Δφ',
               loc='left', pad=8, fontsize=10)
panel_letter(fig, ax_b, 'b')

# === Panel c: per-chromosome dominant period across assays ===
chrs = sorted(results['chr'].unique(),
              key=lambda x: 0 if x == 'chrX' else int(x.replace('chr','')))
xc = np.arange(len(chrs))
width = 0.27
for i, (assay, lbl, color) in enumerate(zip(ASSAYS, LABELS, COLORS)):
    pers = []
    for ch in chrs:
        row = results[(results['chr']==ch) & (results['assay']==assay)]
        pers.append(row['dominant_period_Mbp'].iloc[0] if len(row) else np.nan)
    ax_c.bar(xc + (i-1)*width, pers, width=width, color=color,
             edgecolor='white', linewidth=0.4, label=lbl)
ax_c.set_xticks(xc)
ax_c.set_xticklabels([c.replace('chr','') for c in chrs], fontsize=9)
ax_c.set_xlabel('MCF-7 chromosome')
ax_c.set_ylabel('dominant period (Mbp)')
ax_c.legend(loc='upper right', fontsize=8.5)
ax_c.set_title('5-chromosome cross-assay sweep (chrX, chr1, chr2, chr11, chr17)',
               loc='left', pad=8, fontsize=10)
panel_letter(fig, ax_c, 'c')

# 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)
