"""
Build the hierarchy-completing schematic figure for v81 manuscript.

Six levels in a horizontal sequence:
  1. Nucleosomes (~10 nm; 147 bp)
  2. Loops (kilobase-scale; CTCF-anchored)
  3. TADs (~megabase; topologically-associating domain)
  4. A/B compartments (5-10 Mbp; speckle vs lamina)
  5. **WAVE** (highlighted; supra-compartment periodic radial fold)
  6. Chromosome territory (whole-chromosome volume)

Each level shown as a stylized icon + name + characteristic length scale.
The wave is highlighted in gold (the "new" level) with the radial-accordion
geometry illustrated.

Output: hierarchy_schematic.png (suitable for Fig 1 panel A)
"""
import matplotlib.pyplot as plt
import matplotlib.patches as mp
import matplotlib.patches as patches
from matplotlib.patches import FancyBboxPatch, Circle, Rectangle, Arc, FancyArrowPatch
import numpy as np

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

fig, ax = plt.subplots(figsize=(15, 5.5))
ax.set_xlim(0, 15)
ax.set_ylim(0, 5.5)
ax.axis('off')

# Title
ax.text(7.5, 5.1, 'A new level in the hierarchy of chromosome organisation',
        fontsize=13, fontweight='bold', ha='center', family='DejaVu Sans')

# Level 1: nucleosomes
xc = 1.0
ax.text(xc, 4.55, '1. Nucleosomes', fontsize=10, ha='center', fontweight='bold')
# Draw 4 nucleosomes as small circles connected by line
for i, x in enumerate([0.5, 0.75, 1.0, 1.25, 1.5]):
    ax.add_patch(Circle((x, 3.0), 0.10, fc=WONG['grey'], ec='black', linewidth=0.6))
ax.plot([0.4, 1.6], [3.0, 3.0], color='black', lw=0.5, zorder=0)
ax.text(xc, 2.4, '~ 147 bp / nuclesome\n~ 10 nm', fontsize=8, ha='center', style='italic')

# Arrow to level 2
ax.annotate('', xy=(2.0, 3.0), xytext=(1.7, 3.0),
            arrowprops=dict(arrowstyle='->', color='black', lw=1))

# Level 2: loops
xc = 3.0
ax.text(xc, 4.55, '2. Loops', fontsize=10, ha='center', fontweight='bold')
# Draw loop with CTCF anchors
ax.add_patch(Arc((xc, 3.0), 1.4, 1.0, theta1=0, theta2=180,
                  edgecolor=WONG['blue'], lw=1.6))
ax.plot([xc-0.7, xc+0.7], [3.0, 3.0], color=WONG['blue'], lw=0.5, ls=':')
ax.add_patch(Circle((xc-0.7, 3.0), 0.08, fc=WONG['vermil'], ec='black', linewidth=0.6))
ax.add_patch(Circle((xc+0.7, 3.0), 0.08, fc=WONG['vermil'], ec='black', linewidth=0.6))
ax.text(xc-0.85, 3.20, 'CTCF', fontsize=6.5, ha='center', color=WONG['vermil'])
ax.text(xc, 2.4, '~ 100 kb\nCTCF-anchored', fontsize=8, ha='center', style='italic')

# Arrow
ax.annotate('', xy=(3.95, 3.0), xytext=(3.7, 3.0),
            arrowprops=dict(arrowstyle='->', color='black', lw=1))

# Level 3: TADs
xc = 5.0
ax.text(xc, 4.55, '3. TADs', fontsize=10, ha='center', fontweight='bold')
# Draw nested loop ensemble
for offset, w, h in [(-0.4, 0.7, 0.4), (0.3, 0.5, 0.3)]:
    ax.add_patch(Arc((xc + offset, 3.0), w, h, theta1=0, theta2=180,
                      edgecolor=WONG['blue'], lw=1.0))
# TAD boundary box
ax.add_patch(FancyBboxPatch((xc-0.85, 2.7), 1.7, 0.6,
                              boxstyle='round,pad=0.05',
                              fc='none', ec=WONG['green'], lw=1.5,
                              linestyle='--'))
ax.text(xc, 2.4, '~ 1 Mb\nself-interacting domain', fontsize=8, ha='center', style='italic')

# Arrow
ax.annotate('', xy=(5.95, 3.0), xytext=(5.7, 3.0),
            arrowprops=dict(arrowstyle='->', color='black', lw=1))

# Level 4: A/B compartments
xc = 7.0
ax.text(xc, 4.55, '4. A/B compartments', fontsize=10, ha='center', fontweight='bold')
# Draw alternating colored stripes
colors_alt = [WONG['vermil'], WONG['blue'], WONG['vermil'], WONG['blue'], WONG['vermil']]
labels_alt = ['A','B','A','B','A']
for i, (col, lab) in enumerate(zip(colors_alt, labels_alt)):
    x0 = xc - 0.85 + i * 0.34
    ax.add_patch(Rectangle((x0, 2.85), 0.30, 0.30, fc=col, ec='black', linewidth=0.4, alpha=0.7))
    ax.text(x0 + 0.15, 3.0, lab, fontsize=8, ha='center', va='center', fontweight='bold', color='white')
ax.text(xc, 2.4, '~ 5–10 Mb\nspeckle ↔ lamina', fontsize=8, ha='center', style='italic')

# Arrow
ax.annotate('', xy=(7.95, 3.0), xytext=(7.7, 3.0),
            arrowprops=dict(arrowstyle='->', color='black', lw=1))

# Level 5: WAVE (highlighted)
xc = 9.5
# Highlight box
ax.add_patch(FancyBboxPatch((xc-1.25, 2.0), 2.5, 2.6,
                              boxstyle='round,pad=0.05',
                              fc=WONG['lightgold'], ec=WONG['darkgold'], lw=2, alpha=0.30))
ax.text(xc, 4.55, 'NEW: 5. Wave', fontsize=10.5, ha='center', fontweight='bold',
        color=WONG['darkgold'])
# Sinusoidal wave
xs = np.linspace(xc-1.0, xc+1.0, 200)
ys = 3.0 + 0.40*np.sin((xs - xc + 1.0) * np.pi * 2.5)
ax.plot(xs, ys, color=WONG['darkgold'], lw=2.0)
# Show A-crests and B-troughs
for x_p in np.linspace(xc-1.0, xc+1.0, 6):
    yp = 3.0 + 0.40*np.sin((x_p - xc + 1.0) * np.pi * 2.5)
    if yp > 3.0:
        ax.add_patch(Circle((x_p, yp), 0.06, fc=WONG['vermil'], ec='black', linewidth=0.5, zorder=5))
    else:
        ax.add_patch(Circle((x_p, yp), 0.06, fc=WONG['blue'], ec='black', linewidth=0.5, zorder=5))
# A-crest / B-trough labels
ax.text(xc - 0.85, 3.6, 'A-crest', fontsize=6.5, color=WONG['vermil'], fontweight='bold', ha='center')
ax.text(xc - 0.40, 2.5, 'B-trough', fontsize=6.5, color=WONG['blue'], fontweight='bold', ha='center')
ax.text(xc, 2.20, '~ 25–50 Mb / cycle\n$T \\propto L^{0.83}$, $N \\approx 5$ cycles/chromosome',
        fontsize=8, ha='center', style='italic',
        bbox=dict(boxstyle='round,pad=0.2', fc='white', ec=WONG['darkgold'], lw=0.8))

# Arrow
ax.annotate('', xy=(11.5, 3.0), xytext=(11.0, 3.0),
            arrowprops=dict(arrowstyle='->', color='black', lw=1))

# Level 6: Chromosome territory
xc = 12.7
ax.text(xc, 4.55, '6. Chr territory', fontsize=10, ha='center', fontweight='bold')
# Draw 3D blob
from matplotlib.patches import Ellipse
ax.add_patch(Ellipse((xc, 3.0), 1.6, 1.2, angle=20,
                      fc=WONG['sky'], ec='black', linewidth=0.8, alpha=0.6))
ax.add_patch(Ellipse((xc-0.1, 3.0), 1.2, 0.9, angle=-15,
                      fc=WONG['pink'], ec='black', linewidth=0.6, alpha=0.5))
ax.text(xc, 2.4, '~ 250 Mbp\nwhole-chromosome volume', fontsize=8, ha='center', style='italic')

# Bottom panel — radial accordion in 3D nuclear context
yc_bottom = 0.9
ax.text(7.5, 1.7, 'The wave: a radial accordion within the nucleus',
        fontsize=11, fontweight='bold', ha='center', style='italic')
# Nuclear envelope (large oval)
nuclear_x, nuclear_y = 7.5, yc_bottom
ax.add_patch(Ellipse((nuclear_x, nuclear_y), 8.5, 1.4,
                      fc='#fafafa', ec='black', linewidth=1.5))
# Lamina label
ax.text(nuclear_x - 4.5, nuclear_y + 0.8, 'Nuclear lamina', fontsize=9, color=WONG['blue'],
        fontweight='bold', ha='left', style='italic')
ax.text(nuclear_x - 4.5, nuclear_y - 0.85, 'Lamina (B-trough)', fontsize=8, color=WONG['blue'],
        ha='left')
ax.text(nuclear_x, nuclear_y + 0.85, 'Nuclear speckle (A-crest)', fontsize=8, color=WONG['vermil'],
        ha='center', fontweight='bold')

# Speckle dots at center
for x_off, y_off in [(-0.5, 0.05), (0.0, 0.10), (0.5, 0.0), (-0.2, -0.10)]:
    ax.add_patch(Circle((nuclear_x + x_off, nuclear_y + y_off), 0.12, fc=WONG['vermil'], ec='black', linewidth=0.5))

# Wave-trace through nucleus
t = np.linspace(0, 1, 200)
xs = nuclear_x - 4.0 + 8.0 * t
ys_radial = nuclear_y + 0.55 * np.cos(t * np.pi * 5)  # 5 cycles per chromosome
ax.plot(xs, ys_radial, color=WONG['darkgold'], lw=2.5, alpha=0.8)
# Mark 5 cycles
for i in range(5):
    cx = nuclear_x - 4.0 + 8.0 * (i + 0.5) / 5.0
    if (i % 2) == 0:
        ax.add_patch(Circle((cx, nuclear_y + 0.55), 0.10, fc=WONG['vermil'], ec='black', zorder=5))
    else:
        ax.add_patch(Circle((cx, nuclear_y - 0.55), 0.10, fc=WONG['blue'], ec='black', zorder=5))

ax.text(nuclear_x + 4.5, nuclear_y, 'Chromosome\nterritory boundary', fontsize=8, ha='left',
        style='italic', color='black')

# Save
plt.tight_layout()
import os
os.makedirs('/sessions/kind-brave-babbage/mnt/Wavelet/submission_FINAL_v24_April16/02_figures/main', exist_ok=True)
os.makedirs('/sessions/kind-brave-babbage/mnt/Wavelet/submission_FINAL_v24_April16/01_manuscript/figures', exist_ok=True)
plt.savefig('/sessions/kind-brave-babbage/mnt/Wavelet/go_rerun_v2_2026-04-30/hierarchy_schematic.png', dpi=300, bbox_inches='tight')
plt.savefig('/sessions/kind-brave-babbage/mnt/Wavelet/go_rerun_v2_2026-04-30/hierarchy_schematic.svg', bbox_inches='tight')
plt.savefig('/sessions/kind-brave-babbage/mnt/Wavelet/submission_FINAL_v24_April16/02_figures/main/hierarchy_schematic.png', dpi=300, bbox_inches='tight')
plt.savefig('/sessions/kind-brave-babbage/mnt/Wavelet/submission_FINAL_v24_April16/01_manuscript/figures/hierarchy_schematic.png', dpi=300, bbox_inches='tight')
plt.close(fig)
print("Wrote hierarchy_schematic.png to all 3 locations")
