"""
Supplementary Figure 14 — RIF1 KO scRS T_width broadening: chromosome- and
                          compartment-stratified triangulation.

REAL DATA ONLY. No rng. Reviewer-runnable.

Source CSVs (BIOMNI 2026-04-15 RIF1 triangulation deliverable):
  - RIF1 triangulation/rif1_triangulation/scrs_twidth_per_chrom.csv
       22 autosomes, paired WT vs RIF1-KO median T_width, fold change, padj.
  - RIF1 triangulation/rif1_triangulation/twidth_by_compartment_partition.csv
       4 partitions (stable_A, stable_B, A_to_B, B_to_A) × bin counts + medians.
  - RIF1 triangulation/rif1_triangulation/twidth_by_partition_delta.csv
       Per-partition ΔT_width means for contribution accounting.

Headline numbers reproduced (manuscript caption + Section IV text):
  - 22/22 autosomes show RIF1-KO > WT T_width at padj < 0.05 (median +16.1%).
  - Stable A + Stable B partitions contribute 58% of the total ΔT_width;
    B→A compartment-switching contributes 38%, A→B contributes 5%.
  - Effect-size r_rb is highest in B→A switch bins (0.95) — the wave's broadening
    is driven by bins that change compartment identity, not by stable bins alone.

Replaces 2026-04-21 v73 synthetic placeholder
(`figure_regen_v73/scripts/supplementary_figure_14_rif1_triangulation.py`,
 rng seed 173, fully synthetic 4-panel figure).
"""
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_DIR = BUNDLE_ROOT / '04_supplementary_data' / 'figure_data'
PER_CHR = DATA_DIR / 'scrs_twidth_per_chrom.csv'
PARTS_FULL = DATA_DIR / 'twidth_by_compartment_partition.csv'
PARTS_DELTA = DATA_DIR / 'twidth_by_partition_delta.csv'

OUT_BASENAMES = [
    BUNDLE_ROOT / '02_figures' / 'supplementary' / 'supplementary_figure_14_rif1_triangulation',
    BUNDLE_ROOT / '03_figures_to_edit' / 'supplementary' / 'supplementary_figure_14_rif1_triangulation',
]
DRAFT_BASENAME = BUNDLE_ROOT / '05_code' / 'figure_regen' / '_drafts' / 'supp14_v1'

# ---------- Load real data ----------
per_chr = pd.read_csv(PER_CHR)
# Sort chromosomes naturally
per_chr['_k'] = per_chr['chrom'].str.replace('chr','').astype(int)
per_chr = per_chr.sort_values('_k').drop(columns='_k').reset_index(drop=True)

parts_full = pd.read_csv(PARTS_FULL)
parts_delta = pd.read_csv(PARTS_DELTA)
parts_delta['contribution'] = parts_delta['n_bins'] * parts_delta['delta_mean']
parts_delta['contribution_pct'] = 100 * parts_delta['contribution'] / parts_delta['contribution'].sum()
# Merge for unified ordering
PARTITION_ORDER = ['stable_A', 'stable_B', 'B_to_A', 'A_to_B']
PARTITION_LABEL = {
    'stable_A': 'Stable A',
    'stable_B': 'Stable B',
    'B_to_A':   'B → A switch',
    'A_to_B':   'A → B switch',
}
PARTITION_COLOUR = {
    'stable_A': WONG['blue'],
    'stable_B': WONG['vermil'],
    'B_to_A':   WONG['gold'],
    'A_to_B':   WONG['grey'],
}
parts_full = parts_full.set_index('partition').loc[PARTITION_ORDER].reset_index()
parts_delta = parts_delta.set_index('partition').loc[PARTITION_ORDER].reset_index()

# ---------- Verify ----------
n_sig = (per_chr['padj']<0.05).sum()
median_delta = per_chr['delta_pct'].median()
mean_delta = per_chr['delta_pct'].mean()
stable_pct = parts_delta[parts_delta['partition'].isin(['stable_A','stable_B'])]['contribution_pct'].sum()
b_to_a_pct = float(parts_delta[parts_delta['partition']=='B_to_A']['contribution_pct'].iloc[0])
a_to_b_pct = float(parts_delta[parts_delta['partition']=='A_to_B']['contribution_pct'].iloc[0])
print(f"[verify] {n_sig}/{len(per_chr)} autosomes RIF1-KO > WT at padj<0.05 (manuscript: 22/22)")
print(f"[verify] median ΔT_width = +{median_delta:.1f}% (manuscript: +16.1%)")
print(f"[verify] Stable A+B contribution = {stable_pct:.0f}% (manuscript: 58%)")
print(f"[verify] B→A switching contribution = {b_to_a_pct:.0f}% (manuscript: 38%)")
print(f"[verify] A→B switching contribution = {a_to_b_pct:.0f}% (manuscript: 5%)")

# ---------- Figure ----------
apply_style()
fig = plt.figure(figsize=(12.0, 8.5))
gs = fig.add_gridspec(2, 2, height_ratios=[1.0, 1.0], width_ratios=[1.3, 1.0],
                      hspace=0.45, wspace=0.30, left=0.07, right=0.985, top=0.93, bottom=0.085)
ax_a = fig.add_subplot(gs[0, 0])  # per-chr ΔT_width
ax_b = fig.add_subplot(gs[0, 1])  # WT vs KO paired scatter
ax_c = fig.add_subplot(gs[1, 0])  # contribution share by partition
ax_d = fig.add_subplot(gs[1, 1])  # per-partition Δmean with effect size

# === Panel a: per-chromosome ΔT_width sorted ===
order = per_chr.sort_values('delta_pct').reset_index(drop=True)
xa = np.arange(len(order))
ax_a.bar(xa, order['delta_pct'], color=WONG['blue'], edgecolor='white', linewidth=0.4)
ax_a.axhline(median_delta, color=WONG['gold'], lw=1.6, ls='--',
             label=f'median = +{median_delta:.1f}%')
ax_a.axhline(0, color='0.3', lw=0.5)
ax_a.set_xticks(xa)
ax_a.set_xticklabels(order['chrom'], rotation=90, fontsize=7)
ax_a.set_ylabel('ΔT_width (RIF1-KO vs WT, %)')
ax_a.legend(loc='upper left', fontsize=8)
ax_a.set_title(f'RIF1-KO broadens replication-front coherence on every autosome — {n_sig}/{len(per_chr)} sig (padj < 0.05)',
               loc='left', pad=8, fontsize=10)
panel_letter(fig, ax_a, 'a')

# === Panel b: WT vs RIF1-KO median T_width paired scatter ===
ax_b.plot([0.2, 0.55], [0.2, 0.55], color='0.5', lw=0.8, ls='--', label='y = x (no change)')
for _, row in per_chr.iterrows():
    ax_b.plot([row['wt_median_twidth'], row['ko_median_twidth']], [row['wt_median_twidth']]*2,
              color='0.85', lw=0.4, alpha=0.5, zorder=1)  # hint for delta direction
ax_b.scatter(per_chr['wt_median_twidth'], per_chr['ko_median_twidth'],
             s=42, c=per_chr['delta_pct'], cmap='viridis',
             edgecolor='white', linewidth=0.6, zorder=3)
# label key chromosomes
top3_delta = per_chr.nlargest(3, 'delta_pct')
for _, row in top3_delta.iterrows():
    ax_b.annotate(row['chrom'],
                  (row['wt_median_twidth'], row['ko_median_twidth']),
                  xytext=(4, 4), textcoords='offset points', fontsize=7, color='0.2')
ax_b.set_xlabel('WT median T_width')
ax_b.set_ylabel('RIF1-KO median T_width')
ax_b.legend(loc='upper left', fontsize=8)
ax_b.set_title(f'All 22 chromosomes lie above y = x  (median +{median_delta:.1f}%)', loc='left', pad=8, fontsize=10)
panel_letter(fig, ax_b, 'b')

# === Panel c: per-partition contribution share ===
xc = np.arange(len(parts_delta))
colours_c = [PARTITION_COLOUR[p] for p in parts_delta['partition']]
labels_c = [PARTITION_LABEL[p] for p in parts_delta['partition']]
bars1 = ax_c.bar(xc - 0.21, parts_delta['contribution_pct'], width=0.40,
                  color=colours_c, edgecolor='white', linewidth=0.5,
                  label='% of total ΔT contribution')
bars2 = ax_c.bar(xc + 0.21, 100*parts_delta['n_bins']/parts_delta['n_bins'].sum(),
                  width=0.40, color=colours_c, edgecolor='white', linewidth=0.5,
                  alpha=0.40, hatch='//', label='% of total bins')
# Annotate values
for x, p in zip(xc, parts_delta['contribution_pct'].values):
    ax_c.text(x-0.21, p+1.5, f'{p:.0f}%', ha='center', va='bottom', fontsize=8, fontweight='bold')
for x, n_pct in zip(xc, 100*parts_delta['n_bins'].values/parts_delta['n_bins'].sum()):
    ax_c.text(x+0.21, n_pct+1.5, f'{n_pct:.0f}%', ha='center', va='bottom', fontsize=7, color='0.35')
ax_c.set_xticks(xc)
ax_c.set_xticklabels(labels_c, fontsize=8.5)
ax_c.set_ylabel('% of total')
ax_c.legend(loc='upper right', fontsize=8)
ax_c.set_title(f'B → A switch is the disproportionate contributor: 17% of bins, 38% of broadening',
               loc='left', pad=8, fontsize=10)
ax_c.set_ylim(0, max(parts_delta['contribution_pct'].max(), 100*parts_delta['n_bins'].max()/parts_delta['n_bins'].sum())*1.20)
panel_letter(fig, ax_c, 'c')

# === Panel d: per-partition Δmean + r_rb effect size ===
xd = np.arange(len(parts_full))
labels_d = [PARTITION_LABEL[p] for p in parts_full['partition']]
colours_d = [PARTITION_COLOUR[p] for p in parts_full['partition']]
ax_d.bar(xd, parts_full['delta_pct'], color=colours_d, edgecolor='white', linewidth=0.5,
         label='ΔT_width (%)')
# Twin axis for r_rb
ax_d2 = ax_d.twinx()
ax_d2.plot(xd, parts_full['r_rb'], 'o-', color=WONG['black'], lw=1.4, markersize=8,
           markerfacecolor='white', markeredgecolor=WONG['black'], markeredgewidth=1.6,
           label='effect size r_rb')
ax_d2.set_ylabel('rank-biserial effect size  r_rb', fontsize=9)
ax_d2.set_ylim(0, 1.0)
ax_d2.tick_params(axis='y', labelsize=8)
# combine legends
lines1, lab1 = ax_d.get_legend_handles_labels()
lines2, lab2 = ax_d2.get_legend_handles_labels()
ax_d.legend(lines1+lines2, lab1+lab2, loc='upper left', fontsize=8)
# value annotations
for x, p in zip(xd, parts_full['delta_pct'].values):
    ax_d.text(x, p+1.0, f'+{p:.1f}%', ha='center', va='bottom', fontsize=7.5, fontweight='bold')
for x, r in zip(xd, parts_full['r_rb'].values):
    ax_d2.text(x, r-0.04, f'{r:.2f}', ha='center', va='top', fontsize=7.5, color='0.15')
ax_d.set_xticks(xd)
ax_d.set_xticklabels(labels_d, fontsize=8.5)
ax_d.set_ylabel('ΔT_width (% increase, median)')
ax_d.set_ylim(0, parts_full['delta_pct'].max()*1.25)
ax_d.set_title('Per-partition broadening: B → A bins largest (Δ +50%, r_rb 0.95)', loc='left', pad=8, fontsize=10)
panel_letter(fig, ax_d, 'd')

# 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")
else:
    print("[draft only — set PROMOTE=1 to write canonical paths]")
plt.close(fig)
