"""
TEST 1 — Rao 2014 sub-compartment phase positions.

Question: Do A1, A2, B1, B2, B3 occupy fixed phase positions within the supra-compartment cycle?
If yes (evenly spaced + ordered), the 5th harmonic IS the sub-compartment ordering
   → direct mechanistic interpretation of the harmonics.
If no (random / overlapping), the harmonics are NOT the Rao sub-compartments
   → some other origin (TAD-scale modulation, duty-cycle asymmetry, etc.)

Method:
- Use master_annotation_K562_v38_hg38.csv (6094 500-kb bins, with cwt_phase in radians + subcomp_mode).
- For each sub-compartment class: circular mean phase, mean cos(phi), bootstrap 95% CI (10,000 reps).
- Test ordering: do the 5 mean phases trace an ordered sequence around the cycle?
- Permutation null: shuffle subcomp_mode labels 10,000 times, recompute Rayleigh R for each class,
  check whether the observed concentration exceeds the null at p < 0.05 per class.
- Plot: polar diagram with each sub-compartment as a vector + 95% CI arc.
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path

OUT = Path("/sessions/kind-brave-babbage/mnt/Wavelet/harmonic_synthesis_2026-05-01")
ANNO = "/sessions/kind-brave-babbage/mnt/Wavelet/go_rerun_v2_2026-04-30/master_annotation_K562_v38_hg38.csv"

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

# Color per sub-compartment (active = warm, inactive = cool)
COL = {'A1': WONG['vermil'], 'A2': WONG['gold'],
       'B1': WONG['green'],  'B2': WONG['sky'],  'B3': WONG['blue']}

print("="*70)
print("TEST 1: Rao sub-compartment phase positions on K562 v38")
print("="*70)

df = pd.read_csv(ANNO)
print(f"Loaded {len(df):,} 500-kb bins")

# Filter to autosomes + X with valid phase + valid subcompartment
sub = df[df['cwt_phase'].notna() & df['subcomp_mode'].isin(['A1','A2','B1','B2','B3'])].copy()
print(f"Bins with valid phase + Rao subcompartment: {len(sub):,}")
print()

# ---- Per-class circular statistics ----

def circ_mean(angles):
    """Circular mean of angles in radians."""
    return np.arctan2(np.sin(angles).sum(), np.cos(angles).sum())

def rayleigh_R(angles):
    """Mean resultant length (concentration); 0 = uniform, 1 = perfectly aligned."""
    n = len(angles)
    return np.sqrt(np.sin(angles).sum()**2 + np.cos(angles).sum()**2) / n

def bootstrap_circ_mean(angles, n_boot=10000, ci=0.95, seed=42):
    """Bootstrap CI for circular mean."""
    rng = np.random.default_rng(seed)
    boot_means = np.empty(n_boot)
    n = len(angles)
    for b in range(n_boot):
        idx = rng.integers(0, n, n)
        boot_means[b] = circ_mean(angles[idx])
    # Use unwrapping around the central estimate
    center = circ_mean(angles)
    diffs = np.angle(np.exp(1j*(boot_means - center)))   # in [-pi, pi]
    lo = np.quantile(diffs, (1-ci)/2)
    hi = np.quantile(diffs, 1 - (1-ci)/2)
    return center, center + lo, center + hi, boot_means

results = []
boot_distros = {}
for cls in ['A1','A2','B1','B2','B3']:
    angles = sub.loc[sub['subcomp_mode']==cls, 'cwt_phase'].values
    n = len(angles)
    cmean, cmean_lo, cmean_hi, boot = bootstrap_circ_mean(angles, n_boot=10000, seed=42)
    R = rayleigh_R(angles)
    cos_mean = np.cos(angles).mean()
    cos_sem = np.cos(angles).std(ddof=1) / np.sqrt(n)
    results.append({'subcomp': cls, 'n': n,
                    'circ_mean_rad': cmean, 'circ_mean_deg': np.degrees(cmean),
                    'ci_lo_rad': cmean_lo, 'ci_hi_rad': cmean_hi,
                    'R': R, 'cos_mean': cos_mean, 'cos_sem': cos_sem})
    boot_distros[cls] = boot
    print(f"{cls}: n={n:5d}  mean phase = {np.degrees(cmean):+7.1f}° "
          f"[{np.degrees(cmean_lo):+7.1f}°, {np.degrees(cmean_hi):+7.1f}°]  "
          f"R = {R:.4f}  <cos> = {cos_mean:+.4f} ± {cos_sem:.4f}")

res = pd.DataFrame(results)
res.to_csv(OUT / "test1_subcomp_phase_summary.csv", index=False)
print()

# ---- Permutation null: are individual class concentrations significant? ----

print("Permutation null (shuffle subcomp_mode labels 10,000 times) ...")
all_phases = sub['cwt_phase'].values
labels = sub['subcomp_mode'].values
n_perm = 10000
rng = np.random.default_rng(123)

null_R = {cls: np.empty(n_perm) for cls in ['A1','A2','B1','B2','B3']}
for p in range(n_perm):
    perm_labels = rng.permutation(labels)
    for cls in ['A1','A2','B1','B2','B3']:
        null_R[cls][p] = rayleigh_R(all_phases[perm_labels == cls])

print()
print("Rayleigh R observed vs permutation null (p = fraction of nulls >= observed):")
for cls in ['A1','A2','B1','B2','B3']:
    obs_R = res.loc[res['subcomp']==cls, 'R'].iloc[0]
    p_perm = (null_R[cls] >= obs_R).mean()
    print(f"  {cls}: R_obs = {obs_R:.4f}, null mean = {null_R[cls].mean():.4f}, "
          f"null 95th = {np.quantile(null_R[cls], 0.95):.4f}, p_perm = {p_perm:.4g}")
    res.loc[res['subcomp']==cls, 'p_perm'] = p_perm

res.to_csv(OUT / "test1_subcomp_phase_summary.csv", index=False)
print()

# ---- Are sub-compartments ORDERED on the cycle? ----
print("Ordering test: do the 5 mean phases trace an ordered sequence?")
order_by_phase = res.sort_values('circ_mean_rad')
print("Order around cycle (most negative → most positive phase):")
print(order_by_phase[['subcomp','circ_mean_deg','n','R','p_perm']].to_string(index=False))
print()

# Check if A1<A2<...<B3 monotone or not
biological_order = ['A1','A2','B1','B2','B3']
phase_at_biorder = [res.loc[res['subcomp']==c,'circ_mean_rad'].iloc[0] for c in biological_order]
# Also check the sequence around the cycle
diffs_consec = np.diff([np.angle(np.exp(1j*p)) for p in phase_at_biorder])
print(f"Sequential phase differences A1→A2→B1→B2→B3 (deg): "
      f"{[f'{np.degrees(d):+.1f}' for d in diffs_consec]}")
print(f"Mean spacing if even (5 classes around cycle): 360/5 = 72.0°")
print()

# ---- Plot polar diagram ----

fig, axes = plt.subplots(1, 2, figsize=(13, 6.0),
                         subplot_kw={'projection':'polar'})

# Panel a: mean direction with 95% CI arc
ax = axes[0]
ax.set_title('Mean phase per Rao sub-compartment\n(K562 v38; n bins per class shown)',
             fontsize=11, pad=18)
for _, row in res.iterrows():
    cls = row['subcomp']
    theta = row['circ_mean_rad']
    r = row['R']
    ci_lo, ci_hi = row['ci_lo_rad'], row['ci_hi_rad']
    # Draw arrow from center
    ax.annotate("", xy=(theta, r), xytext=(0, 0),
                arrowprops=dict(arrowstyle='->', color=COL[cls], lw=2.5))
    # CI arc
    arc_theta = np.linspace(ci_lo, ci_hi, 50)
    ax.plot(arc_theta, [r]*len(arc_theta), color=COL[cls], lw=3, alpha=0.7)
    # Label
    label_r = max(r * 1.18, 0.05)
    ax.text(theta, label_r, f"{cls}\n(n={int(row['n'])})",
            color=COL[cls], fontweight='bold', ha='center', va='center', fontsize=10)
ax.set_ylim(0, max(0.18, res['R'].max()*1.30))
ax.set_theta_zero_location('E')
ax.set_theta_direction(1)  # counter-clockwise

# Panel b: cos-mean ± SEM bar (linear axis for readability)
axes[1].remove()
ax2 = fig.add_subplot(1, 2, 2)
classes = ['A1','A2','B1','B2','B3']
xs = np.arange(len(classes))
cos_means = [res.loc[res['subcomp']==c, 'cos_mean'].iloc[0] for c in classes]
cos_sems  = [res.loc[res['subcomp']==c, 'cos_sem'].iloc[0]  for c in classes]
colors = [COL[c] for c in classes]
ax2.bar(xs, cos_means, yerr=cos_sems, capsize=4, color=colors, edgecolor='black', linewidth=0.6)
ax2.axhline(0, color='black', lw=0.5)
ax2.set_xticks(xs)
ax2.set_xticklabels(classes, fontsize=11, fontweight='bold')
ax2.set_ylabel(r'$\langle \cos(\varphi) \rangle$  (A-side $> 0$, B-side $< 0$)', fontsize=10)
ax2.set_title('Mean cos(phase) per Rao sub-compartment\n(error bars = SEM)', fontsize=11, pad=10)
ax2.grid(axis='y', alpha=0.3)
for i, (m, s) in enumerate(zip(cos_means, cos_sems)):
    ax2.text(i, m + (0.005 if m >= 0 else -0.015) + (np.sign(m) * s),
             f"{m:+.3f}", ha='center', fontsize=9,
             va='bottom' if m >= 0 else 'top')

plt.tight_layout()
plt.savefig(OUT / "test1_subcomp_phase_polar.png", dpi=200, bbox_inches='tight')
plt.savefig(OUT / "test1_subcomp_phase_polar.svg", bbox_inches='tight')
plt.close(fig)
print(f"Saved: {OUT}/test1_subcomp_phase_polar.{{png,svg}}")
print(f"Saved: {OUT}/test1_subcomp_phase_summary.csv")
print()
print("DONE TEST 1")
