"""
Main Fig 7: Triple validation of the supra-compartment wave.

Six-panel figure (2x3) combining:
  a — Schematic: cell-population × single-cell × imaging triple-validation logic
  b — Pooled phase × radial-position scatter (Test 10, all 14 GM12878 Dip-C cells)
  c — Per-cell sinusoidal envelopes overlaid (Test 10 panel b)
  d — Per-allele wave: maternal vs paternal phase (Test 11)
  e — MERFISH × TSA-seq: cos(φ_bulk) vs median d_lamina at 1,041 loci (Test 12)
  f — Octant-mean sinusoidal fit on MERFISH median d_lamina (Test 12 panel b)
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.patches import Circle, FancyBboxPatch, FancyArrowPatch
from pathlib import Path
from scipy import stats

OUT = Path("/sessions/kind-brave-babbage/mnt/Wavelet/harmonic_synthesis_2026-05-01")

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

# ---- Load data ----
print("Loading inputs...")
dfm10 = pd.read_csv(OUT / "test10_phase_radial_joint.csv")
fit10 = pd.read_csv(OUT / "test10_per_cell_fit.csv")
piv11 = pd.read_csv(OUT / "test11_per_allele_paired.csv")
dfm12 = pd.read_csv(OUT / "test12_merfish_per_locus.csv")
sum12 = pd.read_csv(OUT / "test12_summary.csv").iloc[0]
pcdf = pd.read_csv(OUT / "test12_per_chrom_r.csv")
print(f"  Test 10: {len(dfm10):,} bead-phase pairs, {len(fit10)} cells")
print(f"  Test 11: {len(piv11)} mat/pat allele pairs")
print(f"  Test 12: {len(dfm12):,} MERFISH loci")

# ---- Pooled fits ----
phi_t10 = dfm10['cwt_phase'].values
r_t10 = dfm10['r_norm'].values
X = np.column_stack([np.ones_like(phi_t10), np.cos(phi_t10), np.sin(phi_t10)])
beta10, *_ = np.linalg.lstsq(X, r_t10, rcond=None)
yhat10 = X @ beta10
n_oct = 8
edges = np.linspace(-np.pi, np.pi, n_oct+1)
oct10 = np.clip(np.digitize(phi_t10, edges)-1, 0, 7)
oct_phi = (edges[:-1]+edges[1:])/2
oct_y10 = np.array([r_t10[oct10==k].mean() for k in range(n_oct)])
oct_yhat10 = np.array([yhat10[oct10==k].mean() for k in range(n_oct)])
R2_oct10 = 1 - np.sum((oct_y10 - oct_yhat10)**2) / np.sum((oct_y10 - oct_y10.mean())**2)

# Test 12 stats
r12_pearson = float(sum12['r_pearson']); R2_oct12 = float(sum12['R2_oct'])
amp12 = float(sum12['amp']); phase_off12 = float(sum12['phase_offset'])
beta12 = [dfm12['d_lamina_z'].mean() - 0,  # intercept fit on z-scored later
          amp12 * np.cos(phase_off12),
          amp12 * np.sin(phase_off12)]
phi_t12 = dfm12['cwt_phase'].values
y_t12 = dfm12['d_lamina_z'].values
X12 = np.column_stack([np.ones_like(phi_t12), np.cos(phi_t12), np.sin(phi_t12)])
beta12_actual, *_ = np.linalg.lstsq(X12, y_t12, rcond=None)
yhat12 = X12 @ beta12_actual
oct12 = np.clip(np.digitize(phi_t12, edges)-1, 0, 7)
oct_y12 = np.array([y_t12[oct12==k].mean() for k in range(n_oct)])
oct_yhat12 = np.array([yhat12[oct12==k].mean() for k in range(n_oct)])

# Test 11 stats
median_dphase = float(np.median(piv11['phase_diff_abs']))
W_t11, p_t11 = stats.wilcoxon(piv11['phase_diff_abs'].values - np.pi/2, alternative='less')

# ---- Build 2×3 figure ----
fig = plt.figure(figsize=(15, 9.5))
gs = gridspec.GridSpec(2, 3, hspace=0.40, wspace=0.30,
                       height_ratios=[1.0, 1.0])

# ============================================================
# Panel a — Schematic of the triple-validation logic
# ============================================================
ax_a = fig.add_subplot(gs[0, 0])
ax_a.set_xlim(0, 10); ax_a.set_ylim(0, 8)
ax_a.axis('off')
ax_a.text(0, 7.7, 'a', fontsize=14, fontweight='bold')
ax_a.text(5, 7.7, 'Three orthogonal validations of the wave',
          fontsize=10.5, fontweight='bold', ha='center')

# Three boxes vertically
boxes = [
    {'y': 5.5, 'label': 'Cell population\nbulk TSA-seq + Hi-C',
     'detail': 'Phase coordinate from accessibility peaks\n→ predicts radial position $R^2 = 0.90$',
     'color': WONG['gold']},
    {'y': 3.5, 'label': 'Single cells\n14 GM12878 Dip-C 3D structures',
     'detail': '14/14 cells carry the wave\nPooled fit $R^2_{\\rm oct} = 0.97$',
     'color': WONG['vermil']},
    {'y': 1.5, 'label': 'Imaging\n1,041 MERFISH loci × 3,032 cells',
     'detail': 'Locus-by-locus radial validation\n$R^2_{\\rm oct} = 0.93$',
     'color': WONG['blue']},
]
for b in boxes:
    box = FancyBboxPatch((0.5, b['y']-0.6), 4.5, 1.5,
                          boxstyle='round,pad=0.05', fc=b['color'], ec='black',
                          linewidth=0.8, alpha=0.20)
    ax_a.add_patch(box)
    ax_a.text(0.7, b['y']+0.45, b['label'], fontsize=10, fontweight='bold',
              color=b['color'])
    ax_a.text(0.7, b['y']-0.25, b['detail'], fontsize=8.5, color='black',
              style='italic')
# Arrows linking boxes (left side) — vertical arrow going through all three
ax_a.annotate('', xy=(5.3, 1.5), xytext=(5.3, 5.5),
              arrowprops=dict(arrowstyle='->', color=WONG['grey'], lw=2.5))
# Convergence text on the right
ax_a.text(5.5, 4.0,
          'All three converge on\nthe same wave\n\n• Same phase offset\n• Same radial direction\n• Same scaling envelope',
          fontsize=9.5, fontweight='bold', va='center',
          bbox=dict(boxstyle='round,pad=0.4', fc='#fafafa', ec='black', lw=0.6))

# ============================================================
# Panel b — Pooled phase × radial scatter (Test 10)
# ============================================================
ax_b = fig.add_subplot(gs[0, 1])
cmap = plt.cm.tab20
for i, (cell_id, sub) in enumerate(dfm10.groupby('cell_id')):
    sample = sub.sample(min(500, len(sub)), random_state=i)
    ax_b.scatter(sample['cwt_phase'], sample['r_norm'],
                 alpha=0.10, s=3, edgecolor='none',
                 color=cmap(i % 20), rasterized=True)
phi_g = np.linspace(-np.pi, np.pi, 200)
y_g10 = beta10[0] + beta10[1]*np.cos(phi_g) + beta10[2]*np.sin(phi_g)
ax_b.plot(phi_g, y_g10, color='black', lw=2.5, zorder=10)
ax_b.errorbar(oct_phi, oct_y10, fmt='o', color='black', ms=8, mfc='white', mew=2,
              zorder=11, capsize=3)
ax_b.set_xticks(np.linspace(-np.pi, np.pi, 5))
ax_b.set_xticklabels(['−π','−π/2','0','+π/2','+π'])
ax_b.set_xlabel('GM12878 CWT phase $\\varphi$ (bulk-derived)', fontsize=10)
ax_b.set_ylabel(r'$r/\hat r_{\rm med}$ (Dip-C, per-cell normalised)', fontsize=10)
ax_b.set_ylim(np.percentile(dfm10['r_norm'], 0.5), np.percentile(dfm10['r_norm'], 99.5))
ax_b.set_title(f"b   Pooled bulk × single-cell scatter\n"
                f"     $n = {len(dfm10):,}$ beads × 14 cells;  $R^2_{{\\rm oct}}$ = {R2_oct10:.3f}",
                fontsize=10, loc='left')
ax_b.grid(alpha=0.25)

# ============================================================
# Panel c — Per-cell sinusoidal envelopes (Test 10 panel b)
# ============================================================
ax_c = fig.add_subplot(gs[0, 2])
for i, row in fit10.iterrows():
    yc = row['a'] + row['b_cos']*np.cos(phi_g) + row['c_sin']*np.sin(phi_g)
    ax_c.plot(phi_g, yc, color=cmap(i % 20), lw=1.0, alpha=0.7)
ax_c.plot(phi_g, y_g10, color='black', lw=2.5, zorder=10)
ax_c.set_xticks(np.linspace(-np.pi, np.pi, 5))
ax_c.set_xticklabels(['−π','−π/2','0','+π/2','+π'])
ax_c.set_xlabel('GM12878 CWT phase $\\varphi$', fontsize=10)
ax_c.set_ylabel(r'$r/\hat r_{\rm med}$ (per-cell sinusoidal fit)', fontsize=10)
ax_c.set_title('c   Per-cell sinusoidal envelopes\n     (each cell one curve; black = pooled wave)',
                fontsize=10, loc='left')
ax_c.grid(alpha=0.25)

# ============================================================
# Panel d — Per-allele wave (Test 11): maternal vs paternal phase scatter
# ============================================================
ax_d = fig.add_subplot(gs[1, 0])
ax_d.scatter(piv11['phase_offset_a'], piv11['phase_offset_b'],
             alpha=0.6, s=22, color=WONG['gold'], edgecolor='black', linewidth=0.4)
ax_d.plot([-np.pi, np.pi], [-np.pi, np.pi], color='black', lw=1, ls='--', alpha=0.7,
          label='cis-encoded ($y = x$)')
ax_d.set_xlim(-np.pi, np.pi); ax_d.set_ylim(-np.pi, np.pi)
ax_d.set_xticks(np.linspace(-np.pi, np.pi, 5))
ax_d.set_xticklabels(['−π','−π/2','0','+π/2','+π'])
ax_d.set_yticks(np.linspace(-np.pi, np.pi, 5))
ax_d.set_yticklabels(['−π','−π/2','0','+π/2','+π'])
ax_d.set_xlabel('Maternal allele wave phase $\\varphi_0^a$ (rad)', fontsize=10)
ax_d.set_ylabel('Paternal allele wave phase $\\varphi_0^b$ (rad)', fontsize=10)
ax_d.set_title(f"d   Per-allele wave: cis-encoded structure\n"
                f"     median |Δφ| = {median_dphase:.2f} (random null π/2)\n"
                f"     Wilcoxon vs random null: $p = {p_t11:.2g}$",
                fontsize=10, loc='left')
ax_d.legend(fontsize=8.5, loc='upper left')
ax_d.set_aspect('equal'); ax_d.grid(alpha=0.25)

# ============================================================
# Panel e — MERFISH × bulk per-locus scatter (Test 12)
# ============================================================
ax_e = fig.add_subplot(gs[1, 1])
sample = dfm12.sample(min(8000, len(dfm12)), random_state=1)
ax_e.scatter(sample['cos_phi'], sample['median_d_lamina_nm'],
             alpha=0.18, s=4, color=WONG['blue'], edgecolor='none', rasterized=True)
xx = np.linspace(-1, 1, 50)
slope = np.polyfit(dfm12['cos_phi'], dfm12['median_d_lamina_nm'], 1)
ax_e.plot(xx, np.polyval(slope, xx), color='black', lw=2)
ax_e.set_xlabel('Bulk-derived $\\cos(\\varphi)$ at locus', fontsize=10)
ax_e.set_ylabel('Median d_lamina (nm)\n(MERFISH, IMR-90)', fontsize=10)
ax_e.set_title(f"e   MERFISH × bulk locus-by-locus radial validation\n"
                f"     $n = {len(dfm12):,}$ loci × {3032} cells\n"
                f"     Pearson $r = {r12_pearson:+.3f}$, $p < 10^{{-15}}$",
                fontsize=10, loc='left')
ax_e.grid(alpha=0.25)

# ============================================================
# Panel f — Sinusoidal fit on MERFISH median d_lamina (octant means + R²)
# ============================================================
ax_f = fig.add_subplot(gs[1, 2])
sample = dfm12.sample(min(8000, len(dfm12)), random_state=2)
ax_f.scatter(sample['cwt_phase'], sample['d_lamina_z'],
             alpha=0.13, s=3, color=WONG['blue'], edgecolor='none', rasterized=True)
y_g12 = beta12_actual[0] + beta12_actual[1]*np.cos(phi_g) + beta12_actual[2]*np.sin(phi_g)
ax_f.plot(phi_g, y_g12, color='black', lw=2.5, zorder=10)
ax_f.errorbar(oct_phi, oct_y12, fmt='o', color='black', ms=8, mfc='white', mew=2,
              zorder=11, capsize=3)
ax_f.set_xticks(np.linspace(-np.pi, np.pi, 5))
ax_f.set_xticklabels(['−π','−π/2','0','+π/2','+π'])
ax_f.set_xlabel('Bulk CWT phase $\\varphi$', fontsize=10)
ax_f.set_ylabel('Per-locus median $d_{\\rm lamina}$ (z, per chrom)', fontsize=10)
ax_f.set_title(f"f   MERFISH octant-mean sinusoidal fit\n"
                f"     $R^2_{{\\rm oct}}$ = {R2_oct12:.3f};  amplitude = {amp12:.3f}",
                fontsize=10, loc='left')
ax_f.grid(alpha=0.25)

plt.suptitle('Triple validation: cell population × single cells × imaging — every readout shows the same wave',
             fontsize=12.5, y=1.000, fontweight='bold')

plt.savefig(OUT / "main_figure_7_triple_validation.png", dpi=200, bbox_inches='tight')
plt.savefig(OUT / "main_figure_7_triple_validation.svg", bbox_inches='tight')
plt.close(fig)
print(f"\nSaved: main_figure_7_triple_validation.png/svg")
