"""
TEST 11 — Per-allele wave shape (haplotype-resolved Dip-C, GM12878).

Question: do the maternal and paternal alleles of a chromosome IN THE SAME
NUCLEUS make the same wave (cis-encoded) or different waves (trans-factor
dependent)?

Method: for each Tan 2018 GM12878 Dip-C cell, fit two sinusoids separately
(one for chrXa = maternal, one for chrXb = paternal). Compare per-cell:
    Δamplitude = |amp_a - amp_b|
    Δphase = circular distance between phase offsets (a vs b)

If Δphase is small (≪ π/2) on average → cis-encoded wave (same in both alleles)
If Δphase is large (~π/2 average) → independent / trans-factor

For 14 cells × 22 autosomes × 2 alleles per cell, this gives:
   - Per-cell-per-chrom paired comparison
   - Aggregate statistic across cells
"""
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 11: Per-allele wave shape on 14 GM12878 Dip-C cells")
print("="*70)

print("\nLoading GM12878 phase map (hg19, 500-kb bins)...")
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):,}")

def load_3dg(path):
    rows = []
    with gzip.open(path, 'rt') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'): continue
            p = line.split('\t')
            if len(p) < 5: continue
            try:
                rows.append({'chrom_hap': p[0], 'pos': int(p[1]),
                             'x': float(p[2]), 'y': float(p[3]), 'z': float(p[4])})
            except (ValueError, IndexError):
                continue
    return pd.DataFrame(rows)

def fit_sin(phi, y):
    """Fit y = a + b·cos(φ) + c·sin(φ); return amp, phase_offset, R²."""
    if len(phi) < 20: return None
    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)
    return {
        'amp': np.sqrt(beta[1]**2 + beta[2]**2),
        'phase_offset': np.arctan2(beta[2], beta[1]),
        'R2': 1 - SSE/SST if SST > 0 else np.nan,
        'n': len(y),
    }

# Load Dip-C cells, split by haplotype
cells = sorted(DIPC_DIR.glob("gm12878_*.3dg.gz"))
all_pairs = []   # per-cell, per-chrom, per-haplotype fits

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

    # Per-cell normalization
    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)
    df3['r_norm'] = df3['r'] / df3['r'].median()

    df3['chrom_base'] = df3['chrom_hap'].str.replace(r'[ab]$', '', regex=True)
    df3['hap'] = df3['chrom_hap'].str[-1]
    df3['chrom'] = 'chr' + df3['chrom_base']

    # Match to phase map
    df_phi_sorted = df_phi[['chrom','bin_mid','cwt_phase']].sort_values(
        ['chrom','bin_mid']).reset_index(drop=True)

    chunks = []
    for chrom in df3['chrom'].unique():
        if chrom not in df_phi_sorted['chrom'].values: continue
        d3 = df3[df3['chrom']==chrom][['chrom','hap','pos','r_norm']].sort_values('pos')
        dp = df_phi_sorted[df_phi_sorted['chrom']==chrom][['bin_mid','cwt_phase']].sort_values('bin_mid')
        merged = pd.merge_asof(d3, dp, left_on='pos', right_on='bin_mid',
                                direction='nearest', tolerance=200_000)
        chunks.append(merged.dropna(subset=['cwt_phase']))
    if not chunks: continue
    df_match = pd.concat(chunks, ignore_index=True)

    # Per-chrom per-haplotype sinusoidal fit
    for chrom in df_match['chrom'].unique():
        sub_chrom = df_match[df_match['chrom']==chrom]
        for hap in ['a','b']:
            sub = sub_chrom[sub_chrom['hap']==hap]
            if len(sub) < 50: continue
            res = fit_sin(sub['cwt_phase'].values, sub['r_norm'].values)
            if res is None: continue
            all_pairs.append({'cell_id': cell_id, 'chrom': chrom, 'hap': hap, **res})

dfp = pd.DataFrame(all_pairs)
print(f"\nTotal per-cell-per-chrom-per-haplotype fits: {len(dfp)}")

# Pivot: for each (cell, chrom), we should have rows for hap='a' and hap='b'
piv = dfp.pivot_table(index=['cell_id','chrom'], columns='hap',
                      values=['amp','phase_offset','R2','n']).reset_index()
piv.columns = ['_'.join(c).strip('_') for c in piv.columns]
print(f"Per-cell-per-chrom paired rows: {len(piv)}")

# Compute Δphase (circular distance) and Δamp
piv = piv.dropna(subset=['phase_offset_a','phase_offset_b','amp_a','amp_b'])
piv['phase_diff_signed'] = np.angle(np.exp(1j*(piv['phase_offset_a'] - piv['phase_offset_b'])))
piv['phase_diff_abs'] = np.abs(piv['phase_diff_signed'])
piv['amp_diff'] = piv['amp_a'] - piv['amp_b']
piv['amp_diff_abs'] = np.abs(piv['amp_diff'])

print(f"\n=== Per-allele wave-shape comparison ===")
print(f"  n paired allele rows: {len(piv)}")
print(f"  median |Δphase| (radians): {piv['phase_diff_abs'].median():.3f}")
print(f"  mean   |Δphase| (radians): {piv['phase_diff_abs'].mean():.3f}")
print(f"  fraction with |Δphase| < π/4: {(piv['phase_diff_abs'] < np.pi/4).mean():.2f}")
print(f"  fraction with |Δphase| < π/2: {(piv['phase_diff_abs'] < np.pi/2).mean():.2f}")
print(f"  median |Δamp|: {piv['amp_diff_abs'].median():.4f}")
print(f"  median amp_a: {piv['amp_a'].median():.4f}, median amp_b: {piv['amp_b'].median():.4f}")
print()
# Test against null: random circular orientation → mean |Δphase| = π/2 exactly
median_dphase = piv['phase_diff_abs'].median()
expected_random = np.pi / 2
print(f"  Expected |Δphase| under uniform-circular null: π/2 = {expected_random:.3f}")
print(f"  Observed median: {median_dphase:.3f}")
print(f"  Difference from null: {(expected_random - median_dphase):.3f} rad below random")
# Wilcoxon vs π/2
W, pval = stats.wilcoxon(piv['phase_diff_abs'].values - expected_random,
                          alternative='less')
print(f"  Wilcoxon signed-rank vs π/2: W = {W:.0f}, p = {pval:.3g} (one-sided alt: less)")

# Save outputs
dfp.to_csv(OUT / "test11_per_allele_fits.csv", index=False)
piv.to_csv(OUT / "test11_per_allele_paired.csv", index=False)

# ---- Plot ----
fig = plt.figure(figsize=(15, 5.0))
gs = gridspec.GridSpec(1, 3, width_ratios=[1.1, 1.1, 1.0], wspace=0.32)

# Panel a — maternal vs paternal phase scatter
ax_a = fig.add_subplot(gs[0])
ax_a.scatter(piv['phase_offset_a'], piv['phase_offset_b'],
             alpha=0.6, s=24, color=WONG['gold'], edgecolor='black', linewidth=0.4)
ax_a.plot([-np.pi, np.pi], [-np.pi, np.pi], color='black', lw=1, ls='--', alpha=0.7,
          label='y = x (cis-encoded)')
ax_a.set_xlim(-np.pi, np.pi); ax_a.set_ylim(-np.pi, np.pi)
ax_a.set_xticks(np.linspace(-np.pi, np.pi, 5))
ax_a.set_xticklabels(['−π','−π/2','0','+π/2','+π'])
ax_a.set_yticks(np.linspace(-np.pi, np.pi, 5))
ax_a.set_yticklabels(['−π','−π/2','0','+π/2','+π'])
ax_a.set_xlabel('Maternal allele wave phase offset (rad)', fontsize=10)
ax_a.set_ylabel('Paternal allele wave phase offset (rad)', fontsize=10)
ax_a.set_title(f"a   Per-cell, per-chromosome paired phase offsets\n"
                f"     n = {len(piv)} (cell × chrom) pairs;\n"
                f"     median |Δphase| = {median_dphase:.2f} rad (random null = π/2 = 1.57)",
                fontsize=10, loc='left')
ax_a.legend(fontsize=8, loc='upper left')
ax_a.grid(alpha=0.3)
ax_a.set_aspect('equal')

# Panel b — distribution of Δphase
ax_b = fig.add_subplot(gs[1])
bins = np.linspace(0, np.pi, 30)
ax_b.hist(piv['phase_diff_abs'], bins=bins, color=WONG['gold'], edgecolor='black', alpha=0.85)
ax_b.axvline(np.pi/2, color='black', lw=1.2, ls='--', label=f'Random null (π/2)')
ax_b.axvline(piv['phase_diff_abs'].median(), color=WONG['vermil'], lw=2,
             label=f'Median observed = {piv["phase_diff_abs"].median():.2f}')
ax_b.set_xlabel('|Δphase| between maternal & paternal alleles (rad)', fontsize=10)
ax_b.set_ylabel('count', fontsize=10)
ax_b.set_title(f"b   Distribution of allelic phase differences\n"
                f"     (Wilcoxon vs random null: p = {pval:.2g})\n"
                f"     fraction |Δphase| < π/4: {(piv['phase_diff_abs'] < np.pi/4).mean():.0%}",
                fontsize=10, loc='left')
ax_b.legend(fontsize=8, loc='upper right')
ax_b.set_xlim(0, np.pi)
ax_b.grid(alpha=0.3, axis='y')

# Panel c — amplitude scatter (mat vs pat)
ax_c = fig.add_subplot(gs[2])
ax_c.scatter(piv['amp_a'], piv['amp_b'], alpha=0.6, s=24, color=WONG['blue'],
             edgecolor='black', linewidth=0.4)
mx = max(piv['amp_a'].max(), piv['amp_b'].max())
ax_c.plot([0, mx*1.05], [0, mx*1.05], color='black', lw=1, ls='--', alpha=0.7,
          label='y = x')
r_amp, p_amp = stats.pearsonr(piv['amp_a'], piv['amp_b'])
ax_c.set_xlim(0, mx*1.05); ax_c.set_ylim(0, mx*1.05)
ax_c.set_xlabel('Maternal allele wave amplitude', fontsize=10)
ax_c.set_ylabel('Paternal allele wave amplitude', fontsize=10)
ax_c.set_title(f"c   Per-allele wave amplitude correlation\n"
                f"     Pearson r = {r_amp:+.3f}, p = {p_amp:.2g}",
                fontsize=10, loc='left')
ax_c.legend(fontsize=8, loc='upper left')
ax_c.grid(alpha=0.3)
ax_c.set_aspect('equal')

plt.suptitle('Per-allele wave-shape comparison: cis-encoded structure (matched maternal & paternal phase) on 14 GM12878 Dip-C cells × autosomes',
             fontsize=11, y=1.05)
plt.tight_layout()
plt.savefig(OUT / "test11_per_allele_wave.png", dpi=200, bbox_inches='tight')
plt.savefig(OUT / "test11_per_allele_wave.svg", bbox_inches='tight')
plt.close(fig)
print(f"\nSaved: test11_per_allele_wave.png/svg")
print(f"Saved: test11_per_allele_paired.csv ({len(piv)} rows)")
