"""
TEST 4 — Radial-geometry refinement: harmonic-augmented sinusoid fit on TSA-seq.

Question: Does adding rank-2 + rank-3 harmonics of cwt_phase to the SON–LMNB1
   sinusoidal fit boost R² above the manuscript-reported 0.90?

If yes:  the harmonics carry real radial-geometric information; the radial
         accordion has substructure (sub-radial layers).
If no:   the radial position is a clean rank-1 sinusoid of the wave coordinate;
         the harmonics in the consecutive-peak GWS are NOT radial sub-structure.

Method:
- Use master_annotation_K562_v38_hg38.csv (cwt_phase, son_tsa, lmnb_tsa columns).
- Compute speckle_axis = SON_z - LMNB1_z (manuscript's axis).
  Or use the existing 'speckle_axis' column if present.
- Fit three nested OLS models:
    M1: y = a0 + a1·cos(φ) + b1·sin(φ)                 [rank-1 only]
    M2: M1 + a2·cos(2φ) + b2·sin(2φ)
    M3: M2 + a3·cos(3φ) + b3·sin(3φ)
- For comparison, also fit the per-octant means form (matches manuscript).
- Report R² (population), R² (octant-mean), and partial F-tests for nested
  improvement.

Population vs octant-mean R²: the manuscript's R² = 0.90 is the OCTANT-MEAN R²
(8-octant binning) which is much higher than the population R² (~0.20) because
the population scatter includes within-octant noise. We report both.

Outputs:
- test4_tsa_fit_summary.csv
- test4_tsa_fit.png/svg (population scatter + octant-mean overlay + 3 model fits)
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import f as f_dist
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'}

print("="*70)
print("TEST 4: Radial-geometry refinement — harmonic-augmented TSA-seq fit")
print("="*70)

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

# Use the SON_z − LMNB1_z axis. Z-normalize within chromosome to remove
# chromosome-mean offsets.
def zscore_per_chrom(df, col):
    return df.groupby('chrom')[col].transform(lambda x: (x - x.mean()) / x.std(ddof=1))

if 'son_tsa' in df.columns and 'lmnb_tsa' in df.columns:
    df['son_z']  = zscore_per_chrom(df, 'son_tsa')
    df['lmnb_z'] = zscore_per_chrom(df, 'lmnb_tsa')
    df['axis']   = df['son_z'] - df['lmnb_z']
elif 'speckle_axis' in df.columns:
    df['axis']   = df['speckle_axis']
else:
    raise ValueError("No SON/LMNB1 columns found")

# Filter to bins with valid phase + axis
sub = df[df['cwt_phase'].notna() & df['axis'].notna()].copy()
print(f"Bins with valid phase + TSA-seq axis: {len(sub):,}")

phi = sub['cwt_phase'].values
y   = sub['axis'].values
n   = len(y)

# --- Build three nested design matrices ---
def design(phi, K):
    """Design matrix with intercept + harmonics 1..K."""
    cols = [np.ones_like(phi)]
    for k in range(1, K+1):
        cols.append(np.cos(k*phi))
        cols.append(np.sin(k*phi))
    return np.column_stack(cols)

def fit_ols(X, y):
    """OLS fit; return beta, R², SSE, df."""
    beta, _, _, _ = np.linalg.lstsq(X, y, rcond=None)
    yhat = X @ beta
    SSE = np.sum((y - yhat)**2)
    SS_tot = np.sum((y - y.mean())**2)
    R2 = 1 - SSE / SS_tot
    df_residual = X.shape[0] - X.shape[1]
    return {'beta': beta, 'yhat': yhat, 'SSE': SSE, 'R2': R2,
            'df_resid': df_residual, 'p': X.shape[1]}

models = {}
for K in [1, 2, 3]:
    X = design(phi, K)
    res = fit_ols(X, y)
    models[K] = res
    print(f"\nModel M{K} (harmonics 1..{K}, p={res['p']} params):")
    print(f"  R² (population) = {res['R2']:.4f}")

# --- Octant-mean R² (matches manuscript's 0.90 metric) ---
n_oct = 8
edges = np.linspace(-np.pi, np.pi, n_oct+1)
oct_idx = np.digitize(phi, edges) - 1
oct_idx = np.clip(oct_idx, 0, n_oct-1)
oct_phi  = (edges[:-1] + edges[1:]) / 2
oct_y    = np.array([y[oct_idx==k].mean() for k in range(n_oct)])

# For octant-mean R², predict the octant means from each model
oct_R2 = {}
for K, res in models.items():
    yhat_oct = np.array([res['yhat'][oct_idx==k].mean() for k in range(n_oct)])
    SS_tot_oct = np.sum((oct_y - oct_y.mean())**2)
    SSE_oct    = np.sum((oct_y - yhat_oct)**2)
    R2_oct = 1 - SSE_oct/SS_tot_oct if SS_tot_oct > 0 else np.nan
    oct_R2[K] = R2_oct
    print(f"  R² (octant-mean, M{K}) = {R2_oct:.4f}")

# --- Partial F-test: M2 vs M1, M3 vs M2 ---
print("\n--- Nested F-tests (testing whether higher harmonics improve fit) ---")
for K_full, K_red in [(2, 1), (3, 2)]:
    SSE_full = models[K_full]['SSE']
    SSE_red  = models[K_red]['SSE']
    df_full  = models[K_full]['df_resid']
    df_red   = models[K_red]['df_resid']
    df_diff  = df_red - df_full   # number of added params
    F = ((SSE_red - SSE_full)/df_diff) / (SSE_full/df_full)
    p_F = 1 - f_dist.cdf(F, df_diff, df_full)
    print(f"  M{K_full} vs M{K_red}: F = {F:.2f}, df = ({df_diff}, {df_full}), p = {p_F:.4g}")
    # Partial R² gain
    SS_tot = SSE_red + (models[K_red]['R2'] / (1 - models[K_red]['R2'])) * SSE_red \
             if models[K_red]['R2'] != 1 else 0
    dR2 = models[K_full]['R2'] - models[K_red]['R2']
    dR2_oct = oct_R2[K_full] - oct_R2[K_red]
    print(f"    ΔR² (population) = {dR2:+.4f}")
    print(f"    ΔR² (octant)     = {dR2_oct:+.4f}")

# --- Save summary ---
summary = []
for K in [1,2,3]:
    summary.append({'model': f'M{K}',
                    'n_params': models[K]['p'],
                    'R2_population': models[K]['R2'],
                    'R2_octant_mean': oct_R2[K],
                    'SSE': models[K]['SSE']})
pd.DataFrame(summary).to_csv(OUT / "test4_tsa_fit_summary.csv", index=False)

# --- Plot ---
fig, axes = plt.subplots(1, 2, figsize=(13, 5.0))

# Panel a: population scatter + 3 model fits
ax = axes[0]
ax.scatter(phi, y, alpha=0.05, s=4, color=WONG['grey'], rasterized=True,
           edgecolor='none', label=f"all bins (n={n:,})")
phi_grid = np.linspace(-np.pi, np.pi, 400)
for K, color, ls in [(1, WONG['blue'], '-'),
                     (2, WONG['vermil'], '--'),
                     (3, WONG['gold'], ':')]:
    Xg = design(phi_grid, K)
    yg = Xg @ models[K]['beta']
    ax.plot(phi_grid, yg, color=color, ls=ls, lw=2.0,
            label=f"M{K} (R²_pop={models[K]['R2']:.3f}, R²_oct={oct_R2[K]:.3f})")
# Octant-mean overlay
ax.errorbar(oct_phi, oct_y, fmt='o', color='black', ms=8, capsize=4, mfc='white', mew=2,
            label='Octant means')
ax.axhline(0, color='black', lw=0.5)
ax.set_xticks(np.linspace(-np.pi, np.pi, 5))
ax.set_xticklabels(['−π', '−π/2', '0', '+π/2', '+π'])
ax.set_xlabel('CWT phase φ', fontsize=11)
ax.set_ylabel(r'SON−LMNB1 axis (z, A-side $> 0$)', fontsize=11)
ax.set_title('TSA-seq vs CWT phase: nested harmonic models', fontsize=11)
ax.legend(fontsize=8.5, loc='upper right')
ax.grid(alpha=0.3)

# Panel b: bar chart of R² by model
ax2 = axes[1]
xs = [1, 2, 3]
labels = ['M1\n(rank-1)', 'M2\n(+rank-2)', 'M3\n(+rank-3)']
R2_pop = [models[K]['R2'] for K in xs]
R2_oct = [oct_R2[K] for K in xs]
w = 0.35
ax2.bar([x-w/2 for x in xs], R2_pop, w, color=WONG['blue'], label='R² population')
ax2.bar([x+w/2 for x in xs], R2_oct, w, color=WONG['gold'], label='R² octant-mean')
for i, x in enumerate(xs):
    ax2.text(x-w/2, R2_pop[i]+0.01, f"{R2_pop[i]:.3f}", ha='center', fontsize=9)
    ax2.text(x+w/2, R2_oct[i]+0.01, f"{R2_oct[i]:.3f}", ha='center', fontsize=9)
ax2.axhline(0.90, color='black', ls='--', lw=1, alpha=0.5,
            label='Manuscript R²_oct = 0.90')
ax2.set_xticks(xs)
ax2.set_xticklabels(labels, fontsize=10)
ax2.set_ylabel('R² (model fit)', fontsize=11)
ax2.set_ylim(0, 1.0)
ax2.set_title('Improvement from adding harmonics', fontsize=11)
ax2.legend(fontsize=9, loc='upper left')
ax2.grid(alpha=0.3, axis='y')

plt.suptitle('TEST 4 — Does adding rank-2/3 harmonics improve the radial-geometry fit?',
             fontsize=12, y=1.02)
plt.tight_layout()
plt.savefig(OUT / "test4_tsa_fit.png", dpi=200, bbox_inches='tight')
plt.savefig(OUT / "test4_tsa_fit.svg", bbox_inches='tight')
plt.close(fig)

print(f"\nSaved: test4_tsa_fit.{{png,svg}}")
print("\nDONE TEST 4")
