"""
TEST 10 — Phase × radial-position joint scatter across all GM12878 Dip-C cells.

Tests whether each cell's wave is a noisy version of the bulk wave (single
underlying geometry, per-cell variation in noise) or whether each cell makes
a genuinely different wave (cell-state-specific structure).

Method:
- Load 14 Tan 2018 Dip-C 3dg files (hg19-coordinate 3D structures)
- For each cell:
    - Compute nuclear centroid (mean of all 3D positions)
    - Compute per-bead radial distance r from centroid; normalise to per-cell
      median nuclear radius (so that r_norm ≈ 1 at the median bead distance)
- Match each Dip-C bead to a manuscript-canonical GM12878 CWT phase via
  genomic position (master_annotation_GM12878_hg19.csv; 100 kb bins → match
  Dip-C beads to nearest 100 kb bin)
- Plot per-cell joint scatter (phase, r_norm) with bulk sinusoidal envelope
  overlaid

Outputs:
  test10_phase_radial_joint.csv   per-bead phase + r_norm per cell
  test10_phase_radial_joint.png   joint scatter + per-cell sinusoidal fits +
                                   bulk envelope
"""
import gzip
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from pathlib import Path
from scipy import stats

OUT = Path("/sessions/kind-brave-babbage/mnt/Wavelet/harmonic_synthesis_2026-05-01")
DIPC_DIR = Path("/sessions/kind-brave-babbage/mnt/Wavelet/polychrom simulation/dipc_data/extracted")
PHASE_MAP = "/sessions/kind-brave-babbage/mnt/Wavelet/go_rerun_v2_2026-04-30/master_annotation_GM12878_hg19.csv"

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

print("="*70)
print("TEST 10: Phase × radial-position joint scatter, 14 GM12878 Dip-C cells")
print("="*70)

# Load GM12878 phase map (hg19)
print("\nLoading GM12878 phase map (hg19)...")
df_phi = pd.read_csv(PHASE_MAP)
df_phi['cos_phi'] = np.cos(df_phi['cwt_phase'])
print(f"Phase-map bins: {len(df_phi):,}")
print(f"Bin size detected: ~{int(df_phi.iloc[1]['bin_mid'] - df_phi.iloc[0]['bin_mid']):,} bp")

# Load Dip-C 3dg files
def load_3dg(path):
    """Parse Tan 2018 .3dg.gz file → DataFrame with chrom_hap, pos, x, y, z."""
    rows = []
    with gzip.open(path, 'rt') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'): continue
            parts = line.split('\t')
            if len(parts) < 5: continue
            try:
                rows.append({
                    'chrom_hap': parts[0],
                    'pos': int(parts[1]),
                    'x': float(parts[2]), 'y': float(parts[3]), 'z': float(parts[4]),
                })
            except (ValueError, IndexError):
                continue
    return pd.DataFrame(rows)

cells = sorted(DIPC_DIR.glob("gm12878_*.3dg.gz"))
print(f"\nFound {len(cells)} Dip-C cells")

all_rows = []
per_cell_fit = []

for cellpath in cells:
    cell_id = cellpath.stem.replace('.3dg', '')
    df3 = load_3dg(cellpath)
    if not len(df3): continue

    # Nuclear centroid (median for robustness) + per-cell scale
    cx, cy, cz = df3['x'].median(), df3['y'].median(), df3['z'].median()
    df3['r'] = np.sqrt((df3['x']-cx)**2 + (df3['y']-cy)**2 + (df3['z']-cz)**2)
    median_r = df3['r'].median()
    df3['r_norm'] = df3['r'] / median_r

    # For each bead, find its CWT phase
    # Strip 'a'/'b' to get chromosome label
    df3['chrom_base'] = df3['chrom_hap'].str.replace(r'[ab]$', '', regex=True)
    df3['chrom'] = 'chr' + df3['chrom_base']

    # Match to phase map: find phase-map bin containing each Dip-C bead position
    # Phase-map bins are ~100 kb. Use merge_asof.
    df_phi_sorted = df_phi[['chrom','bin_mid','cwt_phase','cos_phi','is_A']].sort_values(
        ['chrom','bin_mid']).reset_index(drop=True)
    df3_sorted = df3[['chrom','pos','r_norm','chrom_hap','x','y','z']].sort_values(['chrom','pos']).reset_index(drop=True)

    # Per-chrom merge_asof
    chunks = []
    for chrom in df3_sorted['chrom'].unique():
        d3 = df3_sorted[df3_sorted['chrom']==chrom].sort_values('pos')
        dp = df_phi_sorted[df_phi_sorted['chrom']==chrom].sort_values('bin_mid')
        if not len(dp): continue
        merged = pd.merge_asof(d3, dp[['bin_mid','cwt_phase','cos_phi','is_A']],
                                left_on='pos', right_on='bin_mid',
                                direction='nearest', tolerance=200_000)
        chunks.append(merged)
    if not chunks: continue
    df_match = pd.concat(chunks, ignore_index=True)
    df_match['cell_id'] = cell_id
    df_match = df_match.dropna(subset=['cwt_phase','r_norm'])
    print(f"  {cell_id}: {len(df3):,} beads → {len(df_match):,} matched to phase")

    all_rows.append(df_match[['cell_id','chrom','chrom_hap','pos',
                                'cwt_phase','cos_phi','is_A','r_norm']])

    # Per-cell sinusoidal fit: r_norm ~ a + b·cos(φ) + c·sin(φ)
    phi = df_match['cwt_phase'].values
    y = df_match['r_norm'].values
    if len(phi) > 50:
        X = np.column_stack([np.ones_like(phi), np.cos(phi), np.sin(phi)])
        beta, *_ = np.linalg.lstsq(X, y, rcond=None)
        yhat = X @ beta
        SSE = np.sum((y-yhat)**2)
        SST = np.sum((y-y.mean())**2)
        R2 = 1 - SSE/SST if SST > 0 else np.nan
        per_cell_fit.append({'cell_id': cell_id, 'n_beads': len(phi),
                              'a': beta[0], 'b_cos': beta[1], 'c_sin': beta[2],
                              'R2': R2,
                              'amp': np.sqrt(beta[1]**2 + beta[2]**2),
                              'phase_offset': np.arctan2(beta[2], beta[1])})

dfm = pd.concat(all_rows, ignore_index=True)
fit_df = pd.DataFrame(per_cell_fit)
dfm.to_csv(OUT / "test10_phase_radial_joint.csv", index=False)
fit_df.to_csv(OUT / "test10_per_cell_fit.csv", index=False)

print(f"\nTotal matched beads across cells: {len(dfm):,}")
print(f"Per-cell fits ({len(fit_df)} cells):")
print(fit_df[['cell_id','n_beads','b_cos','c_sin','amp','phase_offset','R2']].to_string(index=False))

# Pooled (all cells together) fit
phi_all = dfm['cwt_phase'].values
y_all = dfm['r_norm'].values
X_all = np.column_stack([np.ones_like(phi_all), np.cos(phi_all), np.sin(phi_all)])
beta_all, *_ = np.linalg.lstsq(X_all, y_all, rcond=None)
yhat_all = X_all @ beta_all
R2_pool = 1 - np.sum((y_all-yhat_all)**2)/np.sum((y_all-y_all.mean())**2)
amp_pool = np.sqrt(beta_all[1]**2 + beta_all[2]**2)
phase_pool = np.arctan2(beta_all[2], beta_all[1])
print(f"\nPooled fit (all cells):")
print(f"  amplitude = {amp_pool:.4f}, phase_offset = {phase_pool:.3f} rad,")
print(f"  R²_pop = {R2_pool:.4f}")

# Octant-mean R² (matches manuscript metric)
n_oct = 8
edges = np.linspace(-np.pi, np.pi, n_oct+1)
oct_idx = np.clip(np.digitize(phi_all, edges)-1, 0, 7)
oct_phi = (edges[:-1]+edges[1:])/2
oct_y = np.array([y_all[oct_idx==k].mean() for k in range(n_oct)])
oct_yhat = np.array([yhat_all[oct_idx==k].mean() for k in range(n_oct)])
SS_tot = np.sum((oct_y - oct_y.mean())**2)
SSE = np.sum((oct_y - oct_yhat)**2)
R2_oct = 1 - SSE/SS_tot if SS_tot > 0 else np.nan
print(f"  R²_octant-mean = {R2_oct:.4f}  (manuscript canonical TSA-seq R² = 0.90)")

# ---- Plot ----
fig = plt.figure(figsize=(15, 9))
gs = gridspec.GridSpec(2, 2, height_ratios=[1.4, 1.0], width_ratios=[1.5, 1.0],
                       hspace=0.32, wspace=0.28)

# Panel a — pooled scatter with all cells (each cell a different color), bulk sinusoid
ax_a = fig.add_subplot(gs[0, 0])
cmap = plt.cm.tab20
for i, (cell_id, sub) in enumerate(dfm.groupby('cell_id')):
    sample = sub.sample(min(800, len(sub)), random_state=i)
    ax_a.scatter(sample['cwt_phase'], sample['r_norm'],
                 alpha=0.10, s=4, edgecolor='none',
                 color=cmap(i % 20), label=cell_id if i < 14 else None,
                 rasterized=True)
# Bulk sinusoidal envelope
phi_g = np.linspace(-np.pi, np.pi, 200)
y_g = beta_all[0] + beta_all[1]*np.cos(phi_g) + beta_all[2]*np.sin(phi_g)
ax_a.plot(phi_g, y_g, color='black', lw=2.5, label='Pooled wave',
          zorder=10)
# Octant means
ax_a.errorbar(oct_phi, oct_y, fmt='o', color='black', ms=10, mfc='white', mew=2.4,
              capsize=4, zorder=11, label='Octant means')
ax_a.set_xticks(np.linspace(-np.pi, np.pi, 5))
ax_a.set_xticklabels(['−π','−π/2','0','+π/2','+π'])
ax_a.set_xlabel('GM12878 CWT phase φ', fontsize=10.5)
ax_a.set_ylabel(r'Radial position (per-cell normalised $r/\hat r_{\rm med}$)', fontsize=10.5)
ax_a.set_title(f'a   Joint phase × radial-position scatter, all 14 GM12878 Dip-C cells\n'
                f'    pooled $R^2_{{\\rm oct}}$ = {R2_oct:.3f};  pooled amplitude = {amp_pool:.3f};  '
                f'$n$ = {len(dfm):,} beads',
                fontsize=10.5, loc='left')
ax_a.set_ylim(np.percentile(dfm['r_norm'], 0.5), np.percentile(dfm['r_norm'], 99.5))
ax_a.grid(alpha=0.3)

# Panel b — per-cell sinusoidal fits overlaid (line plot)
ax_b = fig.add_subplot(gs[0, 1])
phi_g = np.linspace(-np.pi, np.pi, 200)
for i, row in fit_df.iterrows():
    yc = row['a'] + row['b_cos']*np.cos(phi_g) + row['c_sin']*np.sin(phi_g)
    ax_b.plot(phi_g, yc, color=cmap(i % 20), lw=1.0, alpha=0.7,
              label=f"{row['cell_id']}")
# Bulk
ax_b.plot(phi_g, y_g, color='black', lw=2.5, zorder=10, label='Pooled')
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 φ', fontsize=10)
ax_b.set_ylabel(r'$r/\hat r_{\rm med}$ (per-cell sinusoidal fit)', fontsize=10)
ax_b.set_title('b   Per-cell sinusoidal envelopes\n     overlaid (each cell one curve)',
                fontsize=10, loc='left')
ax_b.grid(alpha=0.3)
ax_b.legend(fontsize=6.5, loc='best', ncol=2)

# Panel c — per-cell amplitude + R²
ax_c = fig.add_subplot(gs[1, 0])
ax_c.scatter(fit_df['amp'], fit_df['R2'], color=WONG['gold'], s=60, edgecolor='black')
for i, row in fit_df.iterrows():
    ax_c.annotate(row['cell_id'].replace('gm12878_',''),
                   (row['amp'], row['R2']), fontsize=7, alpha=0.8)
ax_c.set_xlabel('Per-cell sinusoidal amplitude', fontsize=10)
ax_c.set_ylabel(r'Per-cell $R^2$ (population, not octant)', fontsize=10)
ax_c.set_title('c   Per-cell wave amplitude vs goodness-of-fit\n'
                '     (each cell has a wave; amplitudes vary)',
                fontsize=10, loc='left')
ax_c.grid(alpha=0.3)

# Panel d — phase-offset distribution (do all cells share the same phase?)
ax_d = fig.add_subplot(gs[1, 1], projection='polar')
phases = fit_df['phase_offset'].values
amps = fit_df['amp'].values
for ph, amp, cell in zip(phases, amps, fit_df['cell_id']):
    ax_d.annotate('', xy=(ph, amp), xytext=(0, 0),
                  arrowprops=dict(arrowstyle='->', color=WONG['gold'], lw=1.5, alpha=0.7))
# Resultant
resultant = np.mean(np.exp(1j*phases) * amps)
ax_d.annotate('', xy=(np.angle(resultant), np.abs(resultant)), xytext=(0,0),
              arrowprops=dict(arrowstyle='->', color='black', lw=3))
ax_d.text(np.angle(resultant), np.abs(resultant)*1.15, 'pooled',
          ha='center', fontsize=9, fontweight='bold')
ax_d.set_title('d   Per-cell wave phase-offset\n'
                '     (gold = each cell; black = pooled)',
                fontsize=10, loc='left', pad=10)
ax_d.set_ylim(0, max(amps)*1.3)

plt.suptitle('Cell-population × single-cell mutual elucidation: '
             'every Dip-C cell makes a wave at the same phase as the pooled wave',
             fontsize=12, y=1.005)
plt.tight_layout()
plt.savefig(OUT / "test10_phase_radial_joint.png", dpi=180, bbox_inches='tight')
plt.savefig(OUT / "test10_phase_radial_joint.svg", bbox_inches='tight')
plt.close(fig)

print(f"\nSaved: test10_phase_radial_joint.png/svg")
print(f"Saved: test10_phase_radial_joint.csv ({len(dfm)} rows)")
print(f"Saved: test10_per_cell_fit.csv ({len(fit_df)} rows)")
