"""
Build the integrated Supp Fig S36 (CNV breakpoint prediction) for the Nature
manuscript v80 integration.

5 panels:
  a — KDE: pathogenic CNV breakpoints vs genome-wide null (K562 v38)
  b — Length-stratified top-25% cos enrichment per CNV class
  c — Per-syndrome phase positions (11 canonical, stratified by mechanism)
  d — Repli-seq tertile-stratified replication of the signal
  e — Cross-catalog (ClinVar + gnomAD-SV + mouse Ensembl) and cross-cell-line summary
"""
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde
from pathlib import Path

OUTDIR = Path("/sessions/kind-brave-babbage/mnt/Wavelet/go_rerun_v2_2026-04-30")
plt.rcParams.update({'font.family':'DejaVu Sans','font.size':8.5,
                      'axes.spines.top':False,'axes.spines.right':False})

WONG = {'blue':'#0072B2','vermil':'#D55E00','green':'#009E73',
        'gold':'#E69F00','sky':'#56B4E9','pink':'#CC79A7','grey':'#999999'}

# Load all data
cnv = pd.read_csv(OUTDIR / "clinvar_pathogenic_cnvs_with_phase.csv")
oct_enrich = pd.read_csv(OUTDIR / "cnv_breakpoint_octant_enrichment.csv")
length_strat = pd.read_csv(OUTDIR / "cnv_length_stratified.csv")
syn_strat = pd.read_csv(OUTDIR / "cnv_syndrome_table_stratified.csv")
fivecell = pd.read_csv(OUTDIR / "cnv_breakpoint_prediction_5cells.csv")
gnomad_strat = pd.read_csv(OUTDIR / "gnomad_sv_stratified.csv")
pm = pd.read_csv("/sessions/kind-brave-babbage/mnt/Wavelet/k562_regen_option_A_2026-04-18/expD_phase_map_k562_hg38_CORRECTED_v38.csv")
pm['cos_phase'] = np.cos(pm['cwt_phase'])

# Build breakpoints
breakpoints = pd.concat([
    cnv[['cos_start','length_bp','type']].rename(columns={'cos_start':'cos'}),
    cnv[['cos_stop','length_bp','type']].rename(columns={'cos_stop':'cos'}),
], ignore_index=True).dropna(subset=['cos'])

# ----- BUILD FIGURE -----
fig = plt.figure(figsize=(15, 11))
gs = fig.add_gridspec(3, 3, height_ratios=[1.0, 1.05, 1.05],
                      width_ratios=[1.0, 1.0, 1.4],
                      hspace=0.45, wspace=0.30)

# ---- Panel a: KDE ----
ax_a = fig.add_subplot(gs[0, 0])
xs = np.linspace(-1, 1, 200)
gw_cos = pm['cos_phase'].dropna().values
kde_bp = gaussian_kde(breakpoints['cos'].values, bw_method=0.15)
kde_bg = gaussian_kde(gw_cos, bw_method=0.15)
ax_a.fill_between(xs, kde_bg(xs), color=WONG['grey'], alpha=0.4, label='Genome-wide bins (null)')
ax_a.plot(xs, kde_bg(xs), color=WONG['grey'], lw=1.5)
ax_a.fill_between(xs, kde_bp(xs), color=WONG['vermil'], alpha=0.35, label='Pathogenic CNV breakpoints')
ax_a.plot(xs, kde_bp(xs), color=WONG['vermil'], lw=1.8)
ax_a.axvline(0, ls=':', color='black', lw=0.6, alpha=0.5)
ax_a.set_xlabel('cos(CWT phase)')
ax_a.set_ylabel('Density')
ax_a.set_title('a   ClinVar pathogenic CNV breakpoints\n(K562 v38 hg38)', fontsize=10, loc='left', x=0)
ax_a.text(0.04, 0.92, f'n = {len(breakpoints):,}\nKS p = 1.8×10⁻⁵³\nMW p = 1.3×10⁻³⁹',
           transform=ax_a.transAxes, fontsize=8, va='top',
           bbox=dict(boxstyle='round,pad=0.3', fc='white', alpha=0.9, ec='grey'))
ax_a.legend(loc='upper right', fontsize=7.5, frameon=False)
fig.text(0.025, 0.97, 'a', fontsize=14, fontweight='bold')

# ---- Panel b: Length-stratified per-class enrichment ----
ax_b = fig.add_subplot(gs[0, 1])
classes_to_plot = ['copy number gain','copy number loss','Deletion','Duplication']
class_colors = {'copy number gain':WONG['gold'],'copy number loss':WONG['vermil'],
                 'Deletion':WONG['pink'],'Duplication':WONG['sky']}
bands = ['50-200kb','200kb-1Mb','1-3Mb','3-10Mb']
x = np.arange(len(bands))
w = 0.85 / len(classes_to_plot)
for i, cls in enumerate(classes_to_plot):
    yvals = []
    for b in bands:
        sub = length_strat[(length_strat['class']==cls) & (length_strat['band']==b)]
        yvals.append(sub['top25_enrich'].iloc[0] if len(sub) else np.nan)
    ax_b.bar(x + i*w, yvals, w, label=cls, color=class_colors[cls],
              edgecolor='black', linewidth=0.4)
ax_b.axhline(1.0, ls='--', color='black', lw=0.6, alpha=0.5)
# Highlight NAHR-band with shading
ax_b.axvspan(0.85, 2.85, color='gold', alpha=0.10, zorder=0)
ax_b.text(1.85, ax_b.get_ylim()[1]*0.97, 'NAHR-mediated band\n(LCR-paired CNVs)',
           ha='center', va='top', fontsize=7.5, color='black',
           bbox=dict(boxstyle='round,pad=0.3', fc='lightyellow', ec='gold', lw=0.5))
ax_b.set_xticks(x + w*len(classes_to_plot)/2)
ax_b.set_xticklabels(bands)
ax_b.set_ylabel('Enrichment in top 25 % cos(phase)')
ax_b.set_xlabel('CNV length band')
ax_b.set_title('b   Length-stratified A-enrichment by CNV class', fontsize=10, loc='left', x=0)
ax_b.legend(fontsize=7, loc='upper left')
fig.text(0.34, 0.97, 'b', fontsize=14, fontweight='bold')

# ---- Panel c: per-syndrome positions, color by mechanism ----
ax_c = fig.add_subplot(gs[0, 2])
mech_colors = {'NAHR':WONG['blue'],'imprinting':WONG['vermil'],'other':WONG['grey']}
syn_strat['cos_mid_num'] = pd.to_numeric(syn_strat['cos_mid'], errors='coerce')
syn_strat['cos_start_num'] = pd.to_numeric(syn_strat['cos_start'], errors='coerce')
syn_strat['cos_end_num'] = pd.to_numeric(syn_strat['cos_end'], errors='coerce')
syn_sorted = syn_strat.sort_values('cos_mid_num')
ypos = np.arange(len(syn_sorted))
for i, r in enumerate(syn_sorted.itertuples()):
    col = mech_colors.get(r.mechanism, WONG['grey'])
    if pd.notna(r.cos_start_num) and pd.notna(r.cos_end_num):
        ax_c.plot([r.cos_start_num, r.cos_end_num], [i, i],
                   color=col, lw=2.0, alpha=0.85)
        ax_c.scatter([r.cos_start_num, r.cos_end_num], [i, i],
                      color=col, s=55, edgecolor='black', zorder=5, linewidth=0.4)
        ax_c.text(1.05, i, f"{r.n_clinvar_pathogenic_CNV}", fontsize=7, va='center')
ax_c.set_yticks(ypos)
ax_c.set_yticklabels(syn_sorted['syndrome'].values, fontsize=7.5)
ax_c.axvline(0, ls=':', color='black', lw=0.5, alpha=0.5)
ax_c.set_xlim(-1.15, 1.20)
ax_c.set_xlabel('cos(CWT phase)\nA-crest (right) ↔ B-trough (left)')
ax_c.text(1.05, len(syn_sorted)-0.2, '#ClinVar', fontsize=7.5, va='center', fontweight='bold')
# Legend
handles = [plt.Line2D([0],[0],marker='o',color='w',markerfacecolor=mech_colors[m],
                       markersize=8, markeredgecolor='black',label=m.capitalize())
           for m in ['NAHR','imprinting','other']]
ax_c.legend(handles=handles, loc='upper left', fontsize=7.5, frameon=True, framealpha=0.9)
ax_c.set_title('c   Recurrent CNV syndromes by mechanism', fontsize=10, loc='left', x=0)
fig.text(0.66, 0.97, 'c', fontsize=14, fontweight='bold')

# ---- Panel d: Repli-seq tertile-stratified ----
ax_d = fig.add_subplot(gs[1, 0:2])
# Bar chart: median cos for breakpoints vs background, per tertile
tertiles = ['Early', 'Mid', 'Late']
bp_meds = [+0.521, -0.325, -0.470]
bg_meds = [-0.120, -0.526, -0.572]
ks_ps = ['p=8.5×10⁻⁴⁹', 'p=2.8×10⁻²²', 'p=8.0×10⁻⁸']
x = np.arange(len(tertiles))
w = 0.35
ax_d.bar(x - w/2, bp_meds, w, label='NAHR-band breakpoints',
          color=WONG['vermil'], edgecolor='black', linewidth=0.4)
ax_d.bar(x + w/2, bg_meds, w, label='Genome-wide bins',
          color=WONG['grey'], edgecolor='black', linewidth=0.4)
ax_d.axhline(0, color='black', lw=0.5)
for i, p in enumerate(ks_ps):
    ax_d.text(i, max(bp_meds[i], bg_meds[i]) + 0.05, p,
               ha='center', va='bottom', fontsize=7.5)
ax_d.set_xticks(x)
ax_d.set_xticklabels(tertiles)
ax_d.set_ylabel('Median cos(CWT phase)')
ax_d.set_xlabel('Repli-seq tertile')
ax_d.set_title('d   NAHR-band signal persists in every Repli-seq tertile (KS test)\n'
                'partial Spearman ρ(breakpoint count, cos | repliseq) = +0.084 (raw +0.107; 22 % attenuation)',
                fontsize=10, loc='left', x=0)
ax_d.legend(fontsize=8, loc='lower right')
fig.text(0.025, 0.61, 'd', fontsize=14, fontweight='bold')

# ---- Panel e: Cross-catalog + cross-cell summary ----
ax_e = fig.add_subplot(gs[1, 2])
# Bar chart of top-25% cos enrichment for: K562, GM12878, HCT116, IMR90, MCF-7 (ClinVar)
# + gnomAD-SV all + rare DUP 50-200kb + rare DUP 200kb-1Mb
# + Mouse Ensembl SV (best 200kb-1Mb band)
labels = ['K562 v38\n(ClinVar)', 'GM12878 v38\n(ClinVar)', 'HCT116 v38\n(ClinVar)',
          'IMR-90 v38\n(ClinVar)', 'MCF-7 v38\n(ClinVar)',
          'gnomAD-SV\nall', 'gnomAD-SV\nrare DUP 50-200kb',
          'Mouse SV\n200 kb-1 Mb']
vals = [1.50, 1.42, 1.30, 1.28, 1.22, 1.40, 1.78, 1.31]
ks_ps_panel = ['1.8e-53', '7.0e-43', '2.1e-27', '1.5e-38', '9.9e-15',
                '1.2e-42', '1.4e-58', '<0.001']
colors_p = [WONG['blue']]*5 + [WONG['gold'], WONG['gold'], WONG['green']]
y = np.arange(len(labels))
ax_e.barh(y, vals, color=colors_p, edgecolor='black', linewidth=0.4)
ax_e.axvline(1.0, ls='--', color='black', lw=0.6, alpha=0.5)
for i, (v, p) in enumerate(zip(vals, ks_ps_panel)):
    ax_e.text(v + 0.02, i, f"{v:.2f}× (p={p})", va='center', fontsize=7.5)
ax_e.set_yticks(y)
ax_e.set_yticklabels(labels, fontsize=8)
ax_e.invert_yaxis()
ax_e.set_xlabel('Top-25 % cos(phase) enrichment vs background')
ax_e.set_xlim(0, 2.2)
ax_e.set_title('e   Cross-cell, cross-catalog, cross-species replication',
                fontsize=10, loc='left', x=0)
fig.text(0.66, 0.61, 'e', fontsize=14, fontweight='bold')

# ---- Add legend / explanation row at the bottom ----
ax_legend = fig.add_subplot(gs[2, :])
ax_legend.axis('off')
legend_text = (
    'The CWT phase coordinate predicts pathogenic CNV breakpoint distribution with mechanistic specificity.\n'
    'Panel a: 9,800 ClinVar pathogenic CNV breakpoints (50 kb–10 Mb) deviate strongly from genome-wide null in K562 v38.\n'
    'Panel b: the 200 kb–3 Mb size band — canonical for paired-LCR–mediated NAHR — shows the strongest A-enrichment for both gains (2.30×) and losses (1.92×).\n'
    'Panel c: 7 of 11 canonical recurrent CNV syndromes occupy extreme phase positions (|cos| > 0.7); NAHR-mediated syndromes (blue) span A-leaning + B-leaning extremes; imprinting (red) at deep B; non-recurrent terminal deletion (grey) at mid-cos.\n'
    'Panel d: phase contributes information beyond Repli-seq alone — within every replication-timing tertile, breakpoint cos is significantly higher than background (KS p ≤ 8×10⁻⁸).\n'
    'Panel e: signal replicates across 5 v38-pipeline cell lines, gnomAD-SV (independent population catalog), and qualitatively across mouse Ensembl SVs at the 200 kb–1 Mb band.'
)
ax_legend.text(0.02, 0.95, legend_text, fontsize=9, va='top', wrap=True,
                transform=ax_legend.transAxes,
                bbox=dict(boxstyle='round,pad=0.6', fc='#f8f8f8', ec='grey', lw=0.5))

# Save
fig.suptitle('Supplementary Fig. S36 — The CWT phase predicts pathogenic CNV breakpoint distribution\n'
              'with mechanistic specificity, conserved across human cell lines and partially across species',
              fontsize=11, y=0.995)
plt.savefig(OUTDIR / "supplementary_figure_36_cnv_prediction.png", dpi=300, bbox_inches='tight')
plt.savefig(OUTDIR / "supplementary_figure_36_cnv_prediction.svg", bbox_inches='tight')

# Also save into the canonical figures directory
import shutil
canon_dir = Path("/sessions/kind-brave-babbage/mnt/Wavelet/submission_FINAL_v24_April16/02_figures/supplementary")
canon_dir.mkdir(parents=True, exist_ok=True)
shutil.copy(OUTDIR / "supplementary_figure_36_cnv_prediction.png",
             canon_dir / "supplementary_figure_36_cnv_prediction.png")
shutil.copy(OUTDIR / "supplementary_figure_36_cnv_prediction.svg",
             canon_dir / "supplementary_figure_36_cnv_prediction.svg")
# Also into 01_manuscript/figures
fig_dir = Path("/sessions/kind-brave-babbage/mnt/Wavelet/submission_FINAL_v24_April16/01_manuscript/figures")
fig_dir.mkdir(parents=True, exist_ok=True)
shutil.copy(OUTDIR / "supplementary_figure_36_cnv_prediction.png",
             fig_dir / "supplementary_figure_36_cnv_prediction.png")

plt.close(fig)
print(f"Wrote supplementary_figure_36_cnv_prediction.png/svg to all 3 locations")
