"""
Extended Data Figure 11 — Mitotic compartment dissolution collapses the supra-compartment wave 426×.

REAL DATA ONLY. No rng. Reviewer-runnable end-to-end.

Source CSVs (all under <bundle>/04_supplementary_data/figure_data/, copied from
  Wavelet/mitotic_wave_abolishment/ at bundle assembly):
  - sig95_chrom_x_stage_chromperiod.csv  (24 DT40 autosomes × 7 stages)
  - alpha_chrom_x_stage_chromperiod.csv  (24 × 7; AR(1) alpha proxy for period stability)
  - stage_summary_chromperiod.csv         (per-stage means, n_sig95_gt_005, including HeLa)

Bundle-relative paths: this script is intended to be reviewer-runnable as
    cd <bundle>; python 05_code/figure_regen/extended_data_figure_11_mitotic_abolishment.py

Pipeline: Gibcus 2018 DT40 CDK1-as time course (synchronisation purity >95%) processed via the
  manuscript-canonical chromperiod Paul m=2 pipeline; Hi-C PC1 eigenvector → consecutive-peak
  CWT → sig95 fraction (proportion of bins above AR(1) 5% null) per chromosome per stage.

Headline numbers reproduced (must match manuscript abstract, Fig 4 caption, ED11 caption):
  - G2_R1 median sig95 = 0.1018
  - Prophase median sig95 = 0.000239
  - G2 → Prophase median collapse ratio = 426×
  - 21/24 autosomes above AR(1) 5% null in G2_R1, 1/24 in Prophase
  - HeLa async vs metaphase comparator: mean 0.087 → 0.057 (n=22; 20/22 → 11/22 above null)
  - AR(1) α stable across mitosis (~0.74 across all stages)

Figure layout (4 panels, hero a, 3 bottom):
  a — per-chromosome sig95 trajectory across the full DT40 cell-cycle time course, one line
       per chromosome, gold median overlay; the 426× G2→Prophase collapse is the visual climax.
  b — stage-wise mean sig95 with per-stage above-null count annotation (21/24 → 1/24).
  c — AR(1) α stable across stages: wavelength preserved while sig95 collapses.
  d — HeLa async → metaphase cross-cell-line corroboration.

Provenance recorded in FIGURE_PROVENANCE.md by the audit step that calls this script.
"""
from __future__ import annotations
import sys
from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import wilcoxon

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

# Bundle-relative paths. Script lives at <bundle>/05_code/figure_regen/<name>.py
BUNDLE_ROOT = Path(__file__).resolve().parents[2]
DATA_DIR = BUNDLE_ROOT / '04_supplementary_data' / 'figure_data'
SIG95_CSV = DATA_DIR / 'sig95_chrom_x_stage_chromperiod.csv'
ALPHA_CSV = DATA_DIR / 'alpha_chrom_x_stage_chromperiod.csv'
STAGE_CSV = DATA_DIR / 'stage_summary_chromperiod.csv'

# Canonical PNG + editable SVG written inside the bundle only
OUT_BASENAMES = [
    BUNDLE_ROOT / '02_figures' / 'extended_data' / 'extended_data_figure_11_mitotic_abolishment',
    BUNDLE_ROOT / '03_figures_to_edit' / 'extended_data' / 'extended_data_figure_11_mitotic_abolishment',
]
DRAFT_BASENAME = BUNDLE_ROOT / '05_code' / 'figure_regen' / '_drafts' / 'ed11_v2'

# ---------- Load real data ----------
df_sig = pd.read_csv(SIG95_CSV)
df_alpha = pd.read_csv(ALPHA_CSV)
df_stage = pd.read_csv(STAGE_CSV)

# Sort chromosomes naturally
def chrom_key(c):
    n = c.replace('chr', '')
    return int(n) if n.isdigit() else 99
df_sig = df_sig.assign(_k=df_sig['chrom'].map(chrom_key)).sort_values('_k').drop(columns='_k').reset_index(drop=True)
df_alpha = df_alpha.assign(_k=df_alpha['chrom'].map(chrom_key)).sort_values('_k').drop(columns='_k').reset_index(drop=True)

stages_per_rep = ['NS', 'G2_R1', 'G2_R2', 'Prophase', 'PM30_R1', 'PM30_R2', 'PM60']
stage_labels   = ['NS',  'G2\nrep 1', 'G2\nrep 2', 'Prophase', 'PM30\nrep 1', 'PM30\nrep 2', 'PM60']

# Stage colours (semantic)
stage_colour = {
    'NS':       SEMANTIC['NS'],
    'G2_R1':    SEMANTIC['G2'],
    'G2_R2':    SEMANTIC['G2'],
    'Prophase': SEMANTIC['Prophase'],
    'PM30_R1':  SEMANTIC['PM'],
    'PM30_R2':  SEMANTIC['PM'],
    'PM60':     SEMANTIC['PM'],
}

# Build stage-summary lookup (DT40 only)
sample_to_stat = df_stage.set_index('sample')

# ---------- Verify ----------
g2_r1_median = df_sig['G2_R1'].median()
prophase_median = df_sig['Prophase'].median()
collapse_ratio = g2_r1_median / max(prophase_median, 1e-12)
n_chr = len(df_sig)
n_g2_above = int(sample_to_stat.loc['G2_R1', 'n_sig95_gt_005'])
n_prophase_above = int(sample_to_stat.loc['Prophase', 'n_sig95_gt_005'])
print(f"[verify] G2_R1 median sig95 = {g2_r1_median:.4f}        manuscript = 0.1018")
print(f"[verify] Prophase median sig95 = {prophase_median:.6f}    manuscript = 0.000239")
print(f"[verify] G2_R1/Prophase collapse = {collapse_ratio:.1f}×    manuscript = 426×")
print(f"[verify] 24 DT40 autosomes: n = {n_chr}")
print(f"[verify] G2_R1 above null: {n_g2_above}/{n_chr}            manuscript = 21/24")
print(f"[verify] Prophase above null: {n_prophase_above}/{n_chr}    manuscript = 1/24")
g2_avg = df_sig[['G2_R1','G2_R2']].mean(axis=1).values
wstat = wilcoxon(g2_avg, df_sig['Prophase'].values, alternative='greater')
print(f"[verify] paired Wilcoxon G2_avg > Prophase: stat={wstat.statistic:.0f}, p={wstat.pvalue:.2e}")
hela_ns_mean = float(sample_to_stat.loc['HeLa_NS', 'mean_sig95'])
hela_m_mean  = float(sample_to_stat.loc['HeLa_M',  'mean_sig95'])
hela_ns_count = int(sample_to_stat.loc['HeLa_NS', 'n_sig95_gt_005'])
hela_m_count  = int(sample_to_stat.loc['HeLa_M',  'n_sig95_gt_005'])
hela_ratio = hela_ns_mean / max(hela_m_mean, 1e-12)
print(f"[verify] HeLa NS mean sig95 = {hela_ns_mean:.3f}, Metaphase = {hela_m_mean:.3f}, ratio = {hela_ratio:.2f}×")
print(f"[verify] HeLa above null: NS {hela_ns_count}/22, Metaphase {hela_m_count}/22")

# ---------- Figure ----------
apply_style()
fig = plt.figure(figsize=(11.5, 9.2))
# 2 rows. Top row = panel a (full width). Bottom row = b/c/d (3 cols).
gs = fig.add_gridspec(
    2, 3,
    height_ratios=[1.55, 1.0],
    hspace=0.50, wspace=0.42,
    left=0.06, right=0.985, top=0.94, bottom=0.075,
)
ax_a = fig.add_subplot(gs[0, :])
ax_b = fig.add_subplot(gs[1, 0])
ax_c = fig.add_subplot(gs[1, 1])
ax_d = fig.add_subplot(gs[1, 2])

# === Panel a: per-chr sig95 trajectory, full time-course ===
xs = np.arange(len(stages_per_rep))
chrom_lines_color = '0.65'
for i in range(n_chr):
    yvals = df_sig.iloc[i][stages_per_rep].values.astype(float)
    ax_a.plot(xs, yvals, color=chrom_lines_color, lw=0.6, alpha=0.4, zorder=2)
    ax_a.scatter(xs, yvals, color=chrom_lines_color, s=8, alpha=0.5, zorder=3, edgecolor='white', linewidth=0.25)

median_traj = [df_sig[s].median() for s in stages_per_rep]
ax_a.plot(xs, median_traj, color=WONG['gold'], lw=2.8, zorder=6,
          marker='o', markersize=9,
          markerfacecolor='white', markeredgecolor=WONG['gold'], markeredgewidth=2.2,
          label=f'median across {n_chr} autosomes')

# AR(1) 5% reference line
ax_a.axhline(0.05, color=WONG['vermil'], lw=0.7, ls=':', zorder=1, alpha=0.85)
ax_a.text(len(stages_per_rep)-0.3, 0.054, 'AR(1) 5% null', color=WONG['vermil'],
          fontsize=8, style='italic', va='bottom', ha='right')

# 426× collapse callout — placed clear of the median line
g2_x, g2_y = 1, g2_r1_median
prophase_x, prophase_y = 3, prophase_median
ax_a.annotate('',
              xy=(prophase_x, prophase_y + 0.012), xytext=(g2_x, g2_y + 0.022),
              arrowprops=dict(arrowstyle='-|>,head_width=0.4,head_length=0.6',
                              color=WONG['gold'], lw=1.7, alpha=0.95))
callout_x = (g2_x + prophase_x) / 2
callout_y = max(g2_y, prophase_y) + 0.10
ax_a.text(callout_x, callout_y,
          f'{collapse_ratio:.0f}× collapse\n(median 0.102 → 0.0002)',
          fontsize=9.5, fontweight='bold', color=WONG['gold'],
          ha='center', va='bottom',
          bbox=dict(boxstyle='round,pad=0.35', facecolor='white', edgecolor=WONG['gold'], lw=0.8, alpha=0.95))

ax_a.set_xticks(xs)
ax_a.set_xticklabels(stage_labels, fontsize=8)
ax_a.set_ylabel('sig95 fraction')
ax_a.set_xlabel('DT40 CDK1-as cell-cycle stage  (Gibcus 2018; >95% synchronisation purity)', labelpad=8)
ax_a.set_title(
    f'Wave coherence collapses {collapse_ratio:.0f}× at mitosis — DT40 G2 → Prophase, '
    f'{n_chr} autosomes',
    loc='left', pad=10, fontsize=10.5)
ax_a.legend(loc='upper right', fontsize=8, ncol=1)
ax_a.set_ylim(-0.02, 0.40)
ax_a.set_xlim(-0.4, len(stages_per_rep)-0.6)
panel_letter(fig, ax_a, 'a')

# === Panel b: stage means + counts above null ===
xb = np.arange(len(stages_per_rep))
means_b = [float(sample_to_stat.loc[s, 'mean_sig95']) for s in stages_per_rep]
sems_b  = [float(sample_to_stat.loc[s, 'sem_sig95']) for s in stages_per_rep]
counts_b = [int(sample_to_stat.loc[s, 'n_sig95_gt_005']) for s in stages_per_rep]
colors_b = [stage_colour[s] for s in stages_per_rep]

bars = ax_b.bar(xb, means_b, yerr=sems_b, color=colors_b,
                edgecolor='white', linewidth=0.5, capsize=2.5,
                error_kw={'lw': 0.7, 'ecolor': '0.3'})
# Count annotations ABOVE bars (offset above SEM whisker)
ymax_b = max(m+s for m,s in zip(means_b, sems_b)) * 1.18
for x, m, s, c in zip(xb, means_b, sems_b, counts_b):
    y_anno = m + s + 0.005
    ax_b.text(x, y_anno, f'{c}/{n_chr}', ha='center', va='bottom',
              fontsize=8, fontweight='bold', color='0.15')
ax_b.set_ylim(0, ymax_b * 1.05)
ax_b.set_xticks(xb)
ax_b.set_xticklabels(stage_labels, fontsize=7.5)
ax_b.set_ylabel('mean sig95')
# Wilcoxon annotation, top-right area inside panel
ax_b.text(0.97, 0.97,
          f'paired Wilcoxon\nG2 (avg) > Prophase\np = {wstat.pvalue:.1e}',
          transform=ax_b.transAxes, fontsize=7.5, 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(f'Above-null count: {n_g2_above}/{n_chr} → {n_prophase_above}/{n_chr}',
               loc='left', pad=6, fontsize=9.5)
panel_letter(fig, ax_b, 'b')

# === Panel c: AR(1) α invariance ===
alpha_means = [df_alpha[s].mean() for s in stages_per_rep]
alpha_sems  = [df_alpha[s].std(ddof=1)/np.sqrt(n_chr) for s in stages_per_rep]
ax_c.errorbar(xb, alpha_means, yerr=alpha_sems, fmt='o-', color=WONG['blue'],
              lw=1.5, markersize=7, capsize=3,
              markerfacecolor='white', markeredgewidth=1.6, markeredgecolor=WONG['blue'])
# Reference line: G2 α
g2_alpha_ref = df_alpha[['G2_R1','G2_R2']].mean(axis=1).mean()
ax_c.axhline(g2_alpha_ref, color=WONG['grey'], lw=0.6, ls='--', alpha=0.7)
ax_c.text(len(stages_per_rep)-0.3, g2_alpha_ref + 0.005,
          f'G2 baseline α = {g2_alpha_ref:.2f}',
          color=WONG['grey'], fontsize=7.5, ha='right', va='bottom')
ax_c.set_xticks(xb)
ax_c.set_xticklabels(stage_labels, fontsize=7.5)
ax_c.set_ylabel('AR(1) α')
ax_c.set_ylim(0.55, 0.92)
ax_c.set_title('Wavelength preserved — only sig95 collapses', loc='left', pad=6, fontsize=9.5)
panel_letter(fig, ax_c, 'c')

# === Panel d: HeLa cross-cell-line ===
xd = np.arange(2)
means_d = [hela_ns_mean, hela_m_mean]
sems_d  = [float(sample_to_stat.loc[s, 'sem_sig95']) for s in ['HeLa_NS','HeLa_M']]
colors_d = [SEMANTIC['NS'], SEMANTIC['Prophase']]
ax_d.bar(xd, means_d, yerr=sems_d, color=colors_d,
         edgecolor='white', linewidth=0.5, capsize=2.5,
         error_kw={'lw': 0.7, 'ecolor': '0.3'})
counts_d = [hela_ns_count, hela_m_count]
ymax_d = max(m+s for m,s in zip(means_d, sems_d)) * 1.30
for x, m, s, c in zip(xd, means_d, sems_d, counts_d):
    ax_d.text(x, m + s + 0.004, f'{c}/22', ha='center', va='bottom', fontsize=8, fontweight='bold', color='0.15')
ax_d.set_ylim(0, ymax_d * 1.05)
ax_d.set_xticks(xd)
ax_d.set_xticklabels(['NS', 'Metaphase\n(98% sync)'], fontsize=8)
ax_d.set_ylabel('mean sig95 (HeLa-S3, n=22)')
ax_d.text(0.5, 0.97,
          f'{hela_ratio:.2f}× mean drop\nsynchronisation 98%',
          transform=ax_d.transAxes, fontsize=7.5, 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(f'HeLa-S3 metaphase: {hela_ns_count}/22 → {hela_m_count}/22 above null',
               loc='left', pad=6, fontsize=9.5)
panel_letter(fig, ax_d, 'd')

# Save: draft first
DRAFT_BASENAME.parent.mkdir(parents=True, exist_ok=True)
savefig_both(fig, DRAFT_BASENAME)
print(f"[wrote draft] {DRAFT_BASENAME}.png + .svg")

# Promote-to-canonical only if PROMOTE env var set
import os
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)
