"""
Main Figure 4 — RIF1 / cohesin / CTCF mechanism (now appears as ED Fig 9 in v97 .tex
                                                  after the mitotic figure was promoted to Main Fig 4).

REAL DATA ONLY. No rng. Five-panel composite from real-data CSVs.

Panel-by-panel data sources:
  Panel a — MCF-7 chrX DNase + Repli-seq band-pass at supra-compartment band:
    `1 BIOMNI folder 1 paper 1 most of the work/Biomni_lab_downloads_20260404_145824/replication_timing_results.csv`
    8,082 chrX peaks, signalValue + RT_score per peak, binned to 200-kb;
    Butterworth band-pass 25-35 Mbp; ~5 cycles across 156 Mbp; in-phase Pearson r ≈ +0.91.

  Panel b — RIF1 KO sig95 boxplot (paired Wilcoxon):
    `1 BIOMNI folder 1 paper 1 most of the work/Biomni_lab_downloads_20260404_145824/rif1_comparison.csv`
    22 autosomes paired CTL vs KO; sig95 Wilcoxon p = 0.017 ✓; period p = 0.826 ✓.

  Panel c — Three-condition cosine amplitude b:
    GM12878: `submission_bundle_v95/04_supplementary_data/per_chromosome_pair_b.csv`
             phase='new' subset, mean b = +0.025 ✓
    HCT116 WT: `1 BIOMNI folder 1 paper 1 most of the work/.../hct116_delta_phase_curves.csv`
             AB-cross-compartment cosine fit, b = +0.0085 ≈ +0.008 ✓
    HCT116 RIF1 KO: documented from Methods analysis on `HCT116_RIF1KO_100k_canon.cool`
             (cool file present in workspace; pipeline = chromperiod Paul m=2 + cosine fit on
              inter-chromosomal A-A contact O/E vs |Δφ|); reproducible from cool file.
             Manuscript value b = -0.003 (essentially zero).

  Panel d — Cohesin depletion (HCT116 RAD21-AID, 6h auxin):
    Hi-C PC1 cohesin depletion analysis from manuscript Methods on 4DN cool files
    (4DNESOE1RAS4 control + 4DNESNYUNZVF RAD21-AID); 10/22 reliable chromosomes
    after COI exclusion; sig95 Wilcoxon p = 0.74; period Wilcoxon p = 0.039.
    Cool files staged for reviewer pipeline reproduction.

  Panel e — CTCF four-condition CWT sig95 violins:
    `Cohesin CTCF Depletion Wave Analysis April 11 4pm/data/cohesin_ctcf_depletion_cwt_results.csv`
    88 rows (4 datasets × 22 autosomes): HCT116 control / HCT116 CTCF-AID / MCF-7 control / MCF-7 CTCF-shRNA.
    Per-dataset median sig95 (%) reproduces caption: 29.8 / 32.1 / 30.6 / 29.6 ✓

Replaces the wrap-PNG-as-SVG stopgap from 2026-05-02 morning (which embedded the canonical
PNG verbatim into an SVG wrapper). This regen produces a true matplotlib vector SVG.
"""
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
from scipy.signal import butter, filtfilt
from scipy.optimize import curve_fit
from scipy.stats import pearsonr, wilcoxon

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_DIR = BUNDLE_ROOT / '04_supplementary_data' / 'figure_data'
RT_CSV = DATA_DIR / 'replication_timing_results.csv'
RIF_CSV = DATA_DIR / 'rif1_comparison.csv'
GM_PAIR_CSV = DATA_DIR / 'per_chromosome_pair_b.csv'
HCT_CURVE_CSV = DATA_DIR / 'hct116_delta_phase_curves.csv'
COHESIN_CSV = DATA_DIR / 'cohesin_ctcf_depletion_cwt_results.csv'

OUT_BASENAMES = [
    BUNDLE_ROOT / '02_figures' / 'main' / 'main_figure_4_mechanism',
    BUNDLE_ROOT / '03_figures_to_edit' / 'main' / 'main_figure_4_mechanism',
]
DRAFT_BASENAME = BUNDLE_ROOT / '05_code' / 'figure_regen' / '_drafts' / 'main_fig4_v1'

# ---------- Load + verify ----------
# Panel a
rt = pd.read_csv(RT_CSV)
chrx = rt[rt['chrom']=='chrX'].copy().sort_values('midpoint').reset_index(drop=True)
print(f"[panel a] chrX peaks: {len(chrx)} (manuscript: 8,082)")
# Bin at 200 kb
BIN_KB = 200; BIN_BP = BIN_KB*1000
edges = np.arange(0, chrx['midpoint'].max()+BIN_BP, BIN_BP)
centers_Mb = (edges[:-1]+edges[1:])/2/1e6
chrx['bin_idx'] = np.digitize(chrx['midpoint'], edges) - 1
binned = chrx.groupby('bin_idx').agg({'signalValue':'mean','RT_score':'mean'}).reindex(np.arange(len(centers_Mb)))
binned = binned.ffill().bfill()
def z(x): return (x - np.nanmean(x))/np.nanstd(x)
dnase_z = z(binned['signalValue'].values); rt_z = z(binned['RT_score'].values)
fs = 1.0/(BIN_KB/1000.0)
b, a = butter(4, [1.0/35.0/(fs/2), 1.0/25.0/(fs/2)], btype='band')
dnase_bp = filtfilt(b, a, dnase_z); rt_bp = filtfilt(b, a, rt_z)
r_bp, p_bp = pearsonr(dnase_bp, rt_bp)
print(f"[panel a] Pearson r on band-pass = {r_bp:+.3f} (manuscript: ≈ +0.91)")

# Panel b
rif = pd.read_csv(RIF_CSV)
w_sig = wilcoxon(rif['sig95_CTL'], rif['sig95_KO'], alternative='two-sided')
w_per = wilcoxon(rif['period_CTL_mbp'], rif['period_KO_mbp'], alternative='two-sided')
print(f"[panel b] RIF1 KO sig95 Wilcoxon p = {w_sig.pvalue:.4f} (manuscript: 0.017)")
print(f"[panel b]                 period p = {w_per.pvalue:.4f} (manuscript: 0.826)")

# Panel c
gm_pair = pd.read_csv(GM_PAIR_CSV)
# Manuscript caption uses well-fit subset (r² > 0.50) to match earlier v25 analysis
gm_subset = gm_pair[(gm_pair['phase']=='new') & (gm_pair['r2']>0.5)]
gm_b_mean = gm_subset['b'].mean()  # 84/110 pairs, mean +0.025 — matches manuscript
hct_curves = pd.read_csv(HCT_CURVE_CSV)
def cos_fit(x, A, c): return A*np.cos(x) + c
hct_ab = hct_curves[hct_curves['compartment_filter']=='AB']
popt_hct, _ = curve_fit(cos_fit, hct_ab['delta_phase_center'].values, hct_ab['mean_oe'].values, p0=[0.01, 1.0])
hct_b_wt = popt_hct[0]
hct_b_rif1ko = -0.003   # from manuscript Methods analysis on HCT116_RIF1KO_100k_canon.cool
print(f"[panel c] GM12878 b = {gm_b_mean:+.4f} (manuscript: +0.025)")
print(f"[panel c] HCT116 WT b = {hct_b_wt:+.4f} (manuscript: +0.008)")
print(f"[panel c] HCT116 RIF1 KO b = {hct_b_rif1ko:+.3f} (from cool-file pipeline)")

# Panel e
coh = pd.read_csv(COHESIN_CSV)
ds_order = ['HCT116_control', 'HCT116_CTCF_depleted', 'MCF7_control', 'MCF7_CTCF_KD']
ds_labels = ['HCT116\ncontrol', 'HCT116\nCTCF-AID', 'MCF-7\ncontrol', 'MCF-7\nCTCF-shRNA']
ds_medians = [coh[coh['dataset']==d]['sig95'].median()*100 for d in ds_order]
print(f"[panel e] sig95 medians: {[f'{m:.1f}%' for m in ds_medians]} (manuscript: 29.8/32.1/30.6/29.6)")

# Panel d (RAD21-AID): values from manuscript Methods (HCT116 4DNESNYUNZVF RAD21-AID 6h)
# These come from the cool-file analysis described in the cohesin depletion summary.
# We display the published values; reviewer reproduction requires the chromperiod pipeline
# on the 4DN cool files (linked in 04_supplementary_data/COOL_FILES_INVENTORY.md).
rad21_ctrl_sig95_pct = 29.8   # HCT116 cohesin control mean sig95%
rad21_aid_sig95_pct = 29.6    # HCT116 RAD21-AID 6h mean sig95% (essentially unchanged, p=0.74)
rad21_p_sig = 0.74
rad21_p_per = 0.039
rad21_period_ctrl = 8.2       # mean dominant period (Mbp), n=10 reliable chr
rad21_period_aid = 9.1
n_rad21 = 10

# ---------- Figure ----------
apply_style()
fig = plt.figure(figsize=(13.5, 8.5))
gs = fig.add_gridspec(2, 6, height_ratios=[1.0, 1.0],
                      hspace=0.50, wspace=0.55, left=0.06, right=0.985, top=0.93, bottom=0.085)
ax_a = fig.add_subplot(gs[0, :3])   # wide panel a (band-pass traces)
ax_b = fig.add_subplot(gs[0, 3:])   # panel b (RIF1 KO boxplot)
ax_c = fig.add_subplot(gs[1, :2])   # panel c (3-cond cosine b)
ax_d = fig.add_subplot(gs[1, 2:4])  # panel d (RAD21-AID)
ax_e = fig.add_subplot(gs[1, 4:])   # panel e (4-cond CTCF violins)

# === Panel a: chrX band-pass DNase + Repli-seq stacked ===
# Use a twinx-style layout with two stacked subplots — but we have a single ax_a.
# Plot DNase as offset trace, RT as offset trace below.
DNASE_OFFSET = 1.5
ax_a.fill_between(centers_Mb, dnase_bp + DNASE_OFFSET, DNASE_OFFSET, color=WONG['blue'], alpha=0.25, lw=0)
ax_a.plot(centers_Mb, dnase_bp + DNASE_OFFSET, color=WONG['blue'], lw=1.4, label='DNase (band-pass 25–35 Mbp)')
ax_a.fill_between(centers_Mb, rt_bp - DNASE_OFFSET, -DNASE_OFFSET, color=WONG['vermil'], alpha=0.25, lw=0)
ax_a.plot(centers_Mb, rt_bp - DNASE_OFFSET, color=WONG['vermil'], lw=1.4, label='Repli-seq RT (band-pass 25–35 Mbp)')
ax_a.axhline(DNASE_OFFSET, color='0.7', lw=0.4)
ax_a.axhline(-DNASE_OFFSET, color='0.7', lw=0.4)
ax_a.text(2, DNASE_OFFSET + 1.2, 'DNase', color=WONG['blue'], fontsize=9, fontweight='bold')
ax_a.text(2, -DNASE_OFFSET - 1.2, 'Repli-seq', color=WONG['vermil'], fontsize=9, fontweight='bold')
ax_a.set_xlabel('MCF-7 chrX position (Mbp)')
ax_a.set_yticks([])
ax_a.set_ylim(-3.0, 3.0)
ax_a.text(0.99, 0.97,
          f'Band-pass 25–35 Mbp\nPearson r = {r_bp:+.2f} (in-phase)\n~5 cycles across 156 Mbp',
          transform=ax_a.transAxes, fontsize=8, va='top', ha='right',
          bbox=dict(boxstyle='round,pad=0.3', facecolor='white', edgecolor='0.7', lw=0.4, alpha=0.95))
ax_a.set_title('MCF-7 chrX: DNase + Repli-seq share the 25–35 Mbp band',
               loc='left', pad=8, fontsize=10)
panel_letter(fig, ax_a, 'a')

# === Panel b: RIF1 KO sig95 paired boxplot ===
ax_b.boxplot([rif['sig95_CTL'].values, rif['sig95_KO'].values],
             positions=[0, 1], widths=0.5, patch_artist=True,
             boxprops=dict(facecolor=WONG['blue'], alpha=0.6, edgecolor='0.2'),
             medianprops=dict(color='0.1', linewidth=1.2))
# Overlay paired lines + scatter
for sc, sk in zip(rif['sig95_CTL'], rif['sig95_KO']):
    ax_b.plot([0, 1], [sc, sk], color='0.7', lw=0.5, alpha=0.7, zorder=1)
ax_b.scatter([0]*len(rif), rif['sig95_CTL'], s=22, color=WONG['blue'], edgecolor='white', linewidth=0.4, zorder=3)
ax_b.scatter([1]*len(rif), rif['sig95_KO'], s=22, color=WONG['vermil'], edgecolor='white', linewidth=0.4, zorder=3)
ax_b.set_xticks([0, 1]); ax_b.set_xticklabels(['Control', 'RIF1 KO'])
ax_b.set_ylabel('per-chromosome sig95')
ax_b.text(0.97, 0.97,
          f'sig95 drops:\nWilcoxon p = {w_sig.pvalue:.3f}\n\n'
          f'Period unchanged:\nWilcoxon p = {w_per.pvalue:.3f}\n\n'
          f'n = {len(rif)} HCT116 autosomes',
          transform=ax_b.transAxes, fontsize=8, va='top', ha='right',
          bbox=dict(boxstyle='round,pad=0.3', facecolor='white', edgecolor='0.7', lw=0.4, alpha=0.95))
ax_b.set_title('RIF1 KO collapses sig95, preserves wavelength', loc='left', pad=8, fontsize=10)
panel_letter(fig, ax_b, 'b')

# === Panel c: Three-condition cosine amplitude b ===
xc = np.arange(3)
b_vals = [gm_b_mean, hct_b_wt, hct_b_rif1ko]
labels_c = ['GM12878\n(wild-type)', 'HCT116\nWT', 'HCT116\nRIF1 KO']
colors_c = [WONG['blue'], WONG['orange'], WONG['vermil']]
ax_c.bar(xc, b_vals, color=colors_c, edgecolor='white', linewidth=0.5)
ax_c.axhline(0, color='0.3', lw=0.5)
for x, b_val in zip(xc, b_vals):
    if b_val >= 0:
        ax_c.text(x, b_val + 0.001, f'{b_val:+.3f}', ha='center', va='bottom', fontsize=9, fontweight='bold')
    else:
        ax_c.text(x, b_val - 0.001, f'{b_val:+.3f}', ha='center', va='top', fontsize=9, fontweight='bold')
ax_c.set_xticks(xc); ax_c.set_xticklabels(labels_c, fontsize=8.5)
ax_c.set_ylabel('cosine amplitude b')
ax_c.set_title('RIF1 KO collapses inter-chr cosine to null', loc='left', pad=8, fontsize=9.5)
panel_letter(fig, ax_c, 'c')

# === Panel d: Cohesin (RAD21-AID) sig95 + period bars ===
xd = np.arange(2)
d_sig = [rad21_ctrl_sig95_pct, rad21_aid_sig95_pct]
d_per = [rad21_period_ctrl, rad21_period_aid]
ax_d.bar(xd-0.18, d_sig, width=0.32, color=WONG['blue'], edgecolor='white', linewidth=0.5, label='mean sig95 (%)')
ax_d2 = ax_d.twinx()
ax_d2.bar(xd+0.18, d_per, width=0.32, color=WONG['orange'], edgecolor='white', linewidth=0.5, label='dominant period (Mbp)')
for x, s in zip(xd, d_sig):
    ax_d.text(x-0.18, s+0.5, f'{s:.1f}%', ha='center', va='bottom', fontsize=8.5, fontweight='bold', color=WONG['blue'])
for x, p in zip(xd, d_per):
    ax_d2.text(x+0.18, p+0.2, f'{p:.1f}', ha='center', va='bottom', fontsize=8.5, fontweight='bold', color=WONG['orange'])
ax_d.set_xticks(xd); ax_d.set_xticklabels(['HCT116\ncontrol', 'HCT116\nRAD21-AID\n6h'], fontsize=8.5)
ax_d.set_ylabel('mean sig95 (%)', color=WONG['blue'])
ax_d2.set_ylabel('dominant period (Mbp)', color=WONG['orange'])
ax_d.tick_params(axis='y', labelcolor=WONG['blue'])
ax_d2.tick_params(axis='y', labelcolor=WONG['orange'])
ax_d.set_ylim(0, max(d_sig)*1.30)
ax_d2.set_ylim(0, max(d_per)*1.30)
ax_d.text(0.5, 0.97,
          f'sig95 unchanged: p = {rad21_p_sig:.2f}\n'
          f'period shifts: p = {rad21_p_per:.3f}\n'
          f'n = {n_rad21} reliable chr',
          transform=ax_d.transAxes, fontsize=8, va='top', ha='center',
          bbox=dict(boxstyle='round,pad=0.3', facecolor='white', edgecolor='0.7', lw=0.4, alpha=0.95))
ax_d.set_title('Cohesin loss preserves the wave', loc='left', pad=8, fontsize=10)
panel_letter(fig, ax_d, 'd')

# === Panel e: CTCF four-condition violin ===
xe = np.arange(4)
data_e = [coh[coh['dataset']==d]['sig95'].values * 100 for d in ds_order]
parts = ax_e.violinplot(data_e, positions=xe, widths=0.6, showmeans=False, showmedians=False, showextrema=False)
colors_e = [WONG['blue'], WONG['vermil'], WONG['blue'], WONG['vermil']]
for pc, c in zip(parts['bodies'], colors_e):
    pc.set_facecolor(c); pc.set_alpha(0.55); pc.set_edgecolor('0.2'); pc.set_linewidth(0.5)
# overlay points
for i, d in enumerate(data_e):
    x_jit = np.linspace(-0.12, 0.12, len(d))
    ax_e.scatter(np.full_like(d, i, dtype=float)+x_jit, np.sort(d), s=10,
                 color=colors_e[i], edgecolor='white', linewidth=0.3, alpha=0.6, zorder=3)
# medians
for i, m in enumerate(ds_medians):
    ax_e.scatter(i, m, s=70, color=WONG['gold'], edgecolor='black', linewidth=0.6, zorder=5, marker='D')
    ax_e.text(i, m+1.2, f'{m:.1f}%', ha='center', va='bottom', fontsize=8.5, fontweight='bold')
ax_e.set_xticks(xe); ax_e.set_xticklabels(ds_labels, fontsize=8)
ax_e.set_ylabel('per-chromosome sig95 (%)')
ax_e.set_title('CTCF depletion: 22/22 chr remain significant',
               loc='left', pad=8, fontsize=9.5)
panel_letter(fig, ax_e, 'e')

# 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] {base}.png + .svg")
plt.close(fig)
