"""
TEST 3 — Within-compartment scaling.

Question: Inside each A or B compartment segment, is the dominant sub-domain
   wavelength a function of the parent segment length (T_sub ∝ L_segment^b)?

Interpretation:
- b ≈ 0.83 → fractal/self-similar in shape: same scaling law operates at TAD scale.
- b ≈ 0     → nested-but-not-fractal: sub-domain has fixed characteristic size,
              independent of parent segment length.

Method:
- Use the K562 v38 hg38 phase map (per-bin, 500-kb resolution).
- Identify compartment segments as maximal runs of same-sign cwt_phase
  (cos(phase) > 0 = A, < 0 = B).
- Filter to segments with ≥ 8 bins (≥ 4 Mbp) so a sub-domain CWT is well-posed.
- For each segment, measure within-segment dominant sub-period using:
    (a) the autocorrelation function (zero-crossing) of the n_peaks_in_bin track,
    (b) the FFT peak of the n_peaks_in_bin track,
    (c) a within-segment CWT on the cwt_amplitude track.
- Regress log(T_sub) ~ log(L_segment) for each method; report b and 95% CI.

Outputs:
- test3_segment_table.csv (one row per segment with L, T_sub, method)
- test3_within_compartment_scaling.csv (regression summary per method)
- test3_within_compartment_scaling.png/svg (3-panel scaling-law plot)
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
from scipy.optimize import curve_fit
from pathlib import Path
import importlib.util, sys

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"

# Reuse chromperiod's Paul wavelet machinery
SRC = "/sessions/kind-brave-babbage/mnt/Wavelet/k562_regen_option_A_2026-04-18/regen_k562_phase_maps.py"
spec = importlib.util.spec_from_file_location('chromperiod', SRC)
cp = importlib.util.module_from_spec(spec)
sys.modules['chromperiod'] = cp
spec.loader.exec_module(cp)

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

print("="*70)
print("TEST 3: Within-compartment scaling (fractal vs nested-but-not-fractal)")
print("="*70)

df = pd.read_csv(ANNO)
df['cos_phi'] = np.cos(df['cwt_phase'])
df['compartment'] = np.where(df['cos_phi'] > 0, 'A', 'B')

# ---- Identify segments per chromosome ----

def find_runs(arr):
    """Return list of (start_idx, end_idx_exclusive, value) for runs of identical values."""
    runs = []
    if len(arr) == 0: return runs
    start = 0
    for i in range(1, len(arr)):
        if arr[i] != arr[i-1]:
            runs.append((start, i, arr[i-1]))
            start = i
    runs.append((start, len(arr), arr[-1]))
    return runs

segments = []
for chrom, cdf in df.groupby('chrom'):
    cdf = cdf.sort_values('bin_mid_hg38').reset_index(drop=True)
    cdf['compartment'] = np.where(cdf['cos_phi'] > 0, 'A', 'B')
    runs = find_runs(cdf['compartment'].values)
    for (s, e, comp) in runs:
        if comp not in ('A', 'B'): continue
        L_bins = e - s
        L_mbp = L_bins * 0.5
        if L_bins >= 8:   # ≥ 4 Mbp — minimum for sub-domain CWT
            segments.append({
                'chrom': chrom,
                'start_bin': s, 'end_bin': e, 'compartment': comp,
                'L_bins': L_bins, 'L_Mbp': L_mbp,
                'start_bp': cdf.loc[s, 'bin_mid_hg38'],
                'end_bp': cdf.loc[e-1, 'bin_mid_hg38'],
            })
print(f"\nFound {len(segments)} compartment segments with ≥8 bins (≥4 Mbp)")
seg_df = pd.DataFrame(segments)
print(f"  A segments: {(seg_df['compartment']=='A').sum()}")
print(f"  B segments: {(seg_df['compartment']=='B').sum()}")
print(f"  L range: {seg_df['L_Mbp'].min():.1f} – {seg_df['L_Mbp'].max():.1f} Mbp")
print()

# ---- Within-segment sub-period measurement ----

def acf_zero_crossing(x):
    """Find the lag of the first zero-crossing of the autocorrelation function.
       Returns 2 × lag (one full cycle) in same units as x."""
    x = x - np.nanmean(x)
    if np.nanstd(x) == 0:
        return np.nan
    n = len(x)
    acf = np.correlate(x, x, mode='full')[n-1:] / (np.var(x) * n)
    # find first zero crossing (sign change)
    for i in range(1, len(acf)):
        if acf[i-1] > 0 and acf[i] < 0:
            return 2 * i   # full cycle = 2 × half cycle
    return np.nan

def fft_peak_period(x, period_min=4, period_max=None):
    """Find the period (in samples) of the FFT peak."""
    x = x - np.nanmean(x)
    if np.nanstd(x) == 0:
        return np.nan
    n = len(x)
    if period_max is None:
        period_max = n / 2
    # FFT
    X = np.abs(np.fft.rfft(x))
    freqs = np.fft.rfftfreq(n, d=1.0)
    # period = 1/freq
    with np.errstate(divide='ignore', invalid='ignore'):
        periods = np.where(freqs > 0, 1.0/freqs, np.inf)
    valid = (periods >= period_min) & (periods <= period_max)
    if valid.sum() < 2:
        return np.nan
    idx = np.argmax(X[valid])
    return periods[valid][idx]

def cwt_dominant_period(x, period_min=4, period_max=None, wavelet_order=2):
    """CWT global-wavelet-spectrum max period."""
    n = len(x)
    if n < 16: return np.nan
    if period_max is None:
        period_max = n * 0.5
    sigma = np.nanstd(x, ddof=1)
    if sigma == 0 or not np.isfinite(sigma):
        return np.nan
    sig = (x - np.nanmean(x)) / sigma
    sig = np.nan_to_num(sig)
    fourier_factor, coi_factor = cp._get_paul_params(wavelet_order)
    s0 = period_min / fourier_factor
    smax = period_max / fourier_factor
    if smax <= s0:
        return np.nan
    n_scales = 60
    scales = np.geomspace(s0, smax, n_scales)
    periods = scales * fourier_factor
    N_pad = int(2 ** np.ceil(np.log2(n)))
    x_hat = np.fft.fft(sig, n=N_pad)
    omega_pos = 2.0 * np.pi * np.arange(1, N_pad // 2 + 1) / N_pad
    omega_neg = -2.0 * np.pi * np.arange(N_pad // 2 - 1, 0, -1) / N_pad
    omega = np.concatenate([[0.0], omega_pos, omega_neg])
    coi = cp.compute_coi(n, 1.0, coi_factor)
    gws = np.zeros(n_scales)
    for j, s in enumerate(scales):
        psi_vals = np.conj(cp._paul_psi_hat(s, omega, m=wavelet_order))
        norm = np.sqrt(2.0 * np.pi * s)
        W_hat = norm * x_hat * psi_vals
        W_n = np.fft.ifft(W_hat)[:n]
        mask = periods[j] < coi
        if mask.sum() > 0:
            gws[j] = np.mean(np.abs(W_n[mask])**2)
    if not np.any(gws > 0):
        return np.nan
    return periods[np.argmax(gws)]


# ---- Within-segment measurement, three methods ----

print("Measuring sub-period within each segment (3 methods)...")
seg_records = []
for s in segments:
    chrom = s['chrom']
    s0, e0 = s['start_bin'], s['end_bin']
    cdf = df[df['chrom']==chrom].sort_values('bin_mid_hg38').reset_index(drop=True)
    # Track 1: peak count per bin (the strongest signal at fine scale)
    n_peaks = cdf['n_peaks_in_bin'].iloc[s0:e0].fillna(0).values
    # Track 2: cwt_amplitude
    cwt_amp = cdf['cwt_amplitude'].iloc[s0:e0].fillna(0).values
    # Track 3: cos_phi (modulation of compartment sign within segment is rare;
    #           but cwt_phase advances even within a segment)
    cwt_ph = cdf['cwt_phase'].iloc[s0:e0].fillna(0).values

    L_bins = s['L_bins']
    pmax_eff = max(L_bins // 2, 4)

    # ACF zero-crossing on n_peaks
    T_acf_npeaks = acf_zero_crossing(n_peaks)
    # FFT peak on n_peaks
    T_fft_npeaks = fft_peak_period(n_peaks, period_min=2, period_max=pmax_eff)
    # CWT dominant on cwt_amplitude
    T_cwt_amp = cwt_dominant_period(cwt_amp, period_min=2, period_max=pmax_eff)

    rec = {**s,
           'T_acf_npeaks_bins': T_acf_npeaks,
           'T_acf_npeaks_Mbp': T_acf_npeaks * 0.5 if not np.isnan(T_acf_npeaks) else np.nan,
           'T_fft_npeaks_bins': T_fft_npeaks,
           'T_fft_npeaks_Mbp': T_fft_npeaks * 0.5 if not np.isnan(T_fft_npeaks) else np.nan,
           'T_cwt_amp_bins': T_cwt_amp,
           'T_cwt_amp_Mbp': T_cwt_amp * 0.5 if not np.isnan(T_cwt_amp) else np.nan,
           }
    seg_records.append(rec)

seg_full = pd.DataFrame(seg_records)
seg_full.to_csv(OUT / "test3_segment_table.csv", index=False)
print(f"Saved {len(seg_full)} segments → test3_segment_table.csv")
print()

# ---- Regression: log(T_sub) ~ log(L_segment) ----

def fit_loglog(L, T):
    """OLS on log-log; return slope, intercept, slope CI (parametric), R^2."""
    valid = (np.isfinite(L) & np.isfinite(T) & (L > 0) & (T > 0))
    Lv = np.log10(L[valid])
    Tv = np.log10(T[valid])
    n = len(Lv)
    if n < 5:
        return None
    slope, intercept = np.polyfit(Lv, Tv, 1)
    Tpred = slope * Lv + intercept
    SSE = np.sum((Tv - Tpred)**2)
    SS_tot = np.sum((Tv - np.mean(Tv))**2)
    R2 = 1 - SSE / SS_tot if SS_tot > 0 else np.nan
    SE_slope = np.sqrt((SSE / (n - 2)) / np.sum((Lv - np.mean(Lv))**2))
    ci_lo = slope - 1.96 * SE_slope
    ci_hi = slope + 1.96 * SE_slope
    return {'slope': slope, 'intercept': intercept,
            'ci_lo': ci_lo, 'ci_hi': ci_hi, 'R2': R2, 'n': n}

regs = {}
for col_L, col_T, label in [
        ('L_Mbp', 'T_acf_npeaks_Mbp', 'ACF zero-crossing on n_peaks_in_bin'),
        ('L_Mbp', 'T_fft_npeaks_Mbp', 'FFT peak on n_peaks_in_bin'),
        ('L_Mbp', 'T_cwt_amp_Mbp',    'CWT max on cwt_amplitude')]:
    res = fit_loglog(seg_full[col_L].values, seg_full[col_T].values)
    if res is None:
        print(f"{label}: fit failed (too few valid points)")
        continue
    regs[label] = res
    print(f"{label}:")
    print(f"  n = {res['n']}, slope b = {res['slope']:+.3f} "
          f"[{res['ci_lo']:+.3f}, {res['ci_hi']:+.3f}], R² = {res['R2']:.3f}")
    if res['ci_lo'] < 0.83 < res['ci_hi']:
        print(f"  → 95% CI INCLUDES manuscript b = 0.83 (consistent with self-similarity)")
    elif abs(res['slope']) < 0.20 and res['ci_lo'] > -0.20 and res['ci_hi'] < 0.20:
        print(f"  → b ≈ 0 (sub-period independent of parent length; "
              f"NESTED-but-not-fractal)")
    else:
        print(f"  → b is between 0 and manuscript value; partial scaling")
    print()

reg_summary = pd.DataFrame([{'method': k, **v} for k,v in regs.items()])
reg_summary.to_csv(OUT / "test3_within_compartment_scaling.csv", index=False)

# ---- Plot ----

fig, axes = plt.subplots(1, 3, figsize=(15, 4.6))
methods = [
    ('T_acf_npeaks_Mbp', 'ACF zero-crossing on n_peaks_in_bin', WONG['blue']),
    ('T_fft_npeaks_Mbp', 'FFT peak on n_peaks_in_bin',          WONG['vermil']),
    ('T_cwt_amp_Mbp',    'CWT max on cwt_amplitude',             WONG['green']),
]
for ax, (col_T, label, col) in zip(axes, methods):
    L = seg_full['L_Mbp'].values
    T = seg_full[col_T].values
    valid = np.isfinite(L) & np.isfinite(T) & (T > 0)
    ax.scatter(L[valid], T[valid], color=col, alpha=0.5, s=12,
               edgecolor='none', label=f"n = {valid.sum()} segments")
    if label in regs:
        r = regs[label]
        Lx = np.linspace(np.log10(L[valid].min()), np.log10(L[valid].max()), 50)
        Ly = r['slope'] * Lx + r['intercept']
        ax.plot(10**Lx, 10**Ly, color='black', lw=1.4,
                label=f"b = {r['slope']:+.2f} [{r['ci_lo']:+.2f}, {r['ci_hi']:+.2f}], R² = {r['R2']:.2f}")
        # Reference lines
        Lx2 = np.linspace(np.log10(L[valid].min()), np.log10(L[valid].max()), 50)
        ax.plot(10**Lx2, 10**(0.83*Lx2 + r['intercept'] + np.mean(Ly) - np.mean(0.83*Lx2 + r['intercept'])),
                color=WONG['gold'], lw=1.0, ls='--',
                label='Reference b = 0.83 (manuscript scaling)')
    ax.set_xscale('log'); ax.set_yscale('log')
    ax.set_xlabel('Parent compartment length L (Mbp)', fontsize=10)
    ax.set_ylabel('Sub-period $T_\\mathrm{sub}$ (Mbp)', fontsize=10)
    ax.set_title(label, fontsize=10)
    ax.legend(fontsize=8, loc='best')
    ax.grid(alpha=0.3, which='both')
plt.suptitle('TEST 3 — Within-compartment scaling: T_sub ∝ L^b ?  '
             '(K562 v38; segments with ≥8 bins / ≥4 Mbp)',
             fontsize=12, y=1.02)
plt.tight_layout()
plt.savefig(OUT / "test3_within_compartment_scaling.png", dpi=200, bbox_inches='tight')
plt.savefig(OUT / "test3_within_compartment_scaling.svg", bbox_inches='tight')
plt.close(fig)

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