"""
Per-chromosome global wavelet spectrum (GWS) visualization — every
autosome + chrX in K562 v38, with all local maxima (the "ups and downs")
marked. The 2017 AutoSignal-style view of individual periodicities.

Method:
- Load K562 v38 ENCFF361VGY narrowPeak (171,118 autosomal peaks).
- For each chromosome, run consecutive-peak CWT (Paul m=2; canonical
  manuscript pipeline; n_scales=160, period band 4 to min(14,000, 0.4*N_peaks)).
- Compute GWS, AR(1) 95% threshold per chromosome.
- Find all local maxima with prominence ≥ 0.05 on the log-GWS curve.
- Plot as 6×4 grid: 23 panels for autosomes + chrX.
- Each panel: log10(GWS) vs log10(period in Mbp), with all peaks marked
  (rank-1 = gold; rank-2/3/4 = vermillion; lower-rank = grey), AR(1) 95%
  as black dashed line, dominant period and chr name + n_peaks in title.

Outputs:
- per_chrom_GWS_K562_v38.csv  (one row per (chrom, peak) with period_Mbp + gws + above_ar1)
- supplementary_figure_41_per_chrom_spectra.png/svg  (24-panel figure)
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
from scipy.stats import chi2
from pathlib import Path
import importlib.util, sys

OUT  = Path("/sessions/kind-brave-babbage/mnt/Wavelet/harmonic_synthesis_2026-05-01")
PEAKS = "/sessions/kind-brave-babbage/mnt/Wavelet/k562_regen_option_A_2026-04-18/ENCFF361VGY.bed"

# Reuse chromperiod pipeline
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'}

# ---- Run CWT per chromosome with full GWS + multi-peak detection ----

def cwt_with_harmonics(signal_raw, wavelet_order=2, n_scales=240,
                       period_min=4, period_max=14000, prominence=0.015):
    """Returns periods (peak-index units), gws, ar1_threshold, peak ranking."""
    N = len(signal_raw)
    if N < 32:
        return None
    sigma = np.std(signal_raw, ddof=1)
    if sigma == 0:
        return None
    signal = (signal_raw - np.mean(signal_raw)) / sigma
    ar1_alpha = cp.estimate_ar1(signal)

    fourier_factor, coi_factor = cp._get_paul_params(wavelet_order)
    s0 = period_min / fourier_factor
    smax = period_max / fourier_factor
    if smax <= s0:
        return None
    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(signal, 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)

    power = np.zeros((n_scales, N))
    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]
        power[j, :] = np.abs(W_n)**2

    chi2_crit = chi2.ppf(0.95, df=2)
    ar1_threshold = np.zeros(n_scales)
    for j in range(n_scales):
        freq_j = 1.0 / periods[j]
        P_k = (1.0 - ar1_alpha**2) / (
              1.0 - 2.0 * ar1_alpha * np.cos(2.0 * np.pi * freq_j) + ar1_alpha**2)
        ar1_threshold[j] = P_k * chi2_crit / 2.0

    gws = np.zeros(n_scales)
    for j in range(n_scales):
        mask = periods[j] < coi
        gws[j] = np.mean(power[j, mask]) if mask.sum() > 0 else np.nan

    valid = np.isfinite(gws) & (gws > 0)
    if valid.sum() < 5:
        return None
    log_gws = np.log(gws + 1e-30)
    valid_idx = np.where(valid)[0]
    log_valid = log_gws[valid]
    peaks_local, _ = find_peaks(log_valid, prominence=prominence)
    peak_idx = valid_idx[peaks_local]
    peak_periods = periods[peak_idx]
    peak_gws = gws[peak_idx]
    # include global max
    global_idx = valid_idx[np.argmax(log_valid)]
    if global_idx not in peak_idx:
        peak_idx = np.append(peak_idx, global_idx)
        peak_periods = np.append(peak_periods, periods[global_idx])
        peak_gws = np.append(peak_gws, gws[global_idx])
    order = np.argsort(-peak_gws)
    peak_idx = peak_idx[order]
    peak_periods = peak_periods[order]
    peak_gws = peak_gws[order]

    return {'periods': periods, 'gws': gws, 'ar1_threshold': ar1_threshold,
            'peak_idx': peak_idx, 'peak_periods': peak_periods, 'peak_gws': peak_gws,
            'ar1_alpha': ar1_alpha, 'N_peaks': N}

# ---- Run on every chromosome ----

print("="*70)
print("Per-chromosome GWS spectra: K562 v38 ENCFF361VGY")
print("="*70)

print("Loading peaks...")
df_peaks = cp.load_narrowpeak(PEAKS)
print(f"Total peaks: {len(df_peaks):,}")

CHROMS = [f'chr{i}' for i in range(1, 23)] + ['chrX']
results_rows = []
spectra = {}

for chrom in CHROMS:
    sub = df_peaks[df_peaks['chrom']==chrom].sort_values('start').reset_index(drop=True)
    if len(sub) < 100:
        print(f"  {chrom}: skipping (only {len(sub)} peaks)")
        continue
    signal = sub['signalValue'].values.astype(float)
    positions = ((sub['start'].values + sub['end'].values)/2.0)
    spacings = np.diff(positions)
    median_spacing_bp = float(np.median(spacings[spacings > 0]))
    chrom_len_bp = int(sub['end'].max())
    res = cwt_with_harmonics(signal, wavelet_order=2, n_scales=240,
                              period_min=4, period_max=14000, prominence=0.015)
    if res is None:
        print(f"  {chrom}: CWT failed")
        continue
    spectra[chrom] = {**res, 'median_spacing_bp': median_spacing_bp,
                      'chrom_len_bp': chrom_len_bp, 'N_peaks': len(sub)}

    # Convert peak_periods from peak-index to Mbp
    peak_Mbp = res['peak_periods'] * median_spacing_bp / 1e6
    n_peaks_to_record = min(8, len(res['peak_periods']))
    for r in range(n_peaks_to_record):
        T_mbp = peak_Mbp[r]
        T_idx = res['peak_periods'][r]
        gws_val = res['peak_gws'][r]
        ar1_at = res['ar1_threshold'][res['peak_idx'][r]]
        above = gws_val > ar1_at
        n_cycles = chrom_len_bp / 1e6 / T_mbp if T_mbp > 0 else np.nan
        results_rows.append({
            'chrom': chrom, 'rank': r,
            'period_index': T_idx, 'period_Mbp': T_mbp,
            'gws': gws_val, 'ar1_threshold': ar1_at, 'above_ar1': above,
            'n_cycles_in_chrom': n_cycles,
            'N_peaks': len(sub),
            'median_spacing_bp': median_spacing_bp,
            'ar1_alpha': res['ar1_alpha'],
        })
    print(f"  {chrom}: N={len(sub):6,}, dom T = {peak_Mbp[0]:5.1f} Mbp, "
          f"n_peaks_in_GWS = {len(res['peak_periods'])}, "
          f"AR1_α = {res['ar1_alpha']:.3f}")

res_df = pd.DataFrame(results_rows)
res_df.to_csv(OUT / "per_chrom_GWS_K562_v38.csv", index=False)
print(f"\nSaved: per_chrom_GWS_K562_v38.csv ({len(res_df)} rows)")

# ---- Plot grid ----

n_chrom = len(spectra)
ncols = 4
nrows = int(np.ceil(n_chrom / ncols))
fig, axes = plt.subplots(nrows, ncols, figsize=(16, 3.0 * nrows),
                         sharey=False, sharex=False)
axes_flat = axes.flatten()

for ax, chrom in zip(axes_flat, sorted(spectra.keys(),
                                        key=lambda c: (c=='chrX', int(c.replace('chr','').replace('X','99'))))):
    s = spectra[chrom]
    msp = s['median_spacing_bp'] / 1e6   # peak-index to Mbp factor
    T_mbp = s['periods'] * msp
    log_gws = np.log10(s['gws'] + 1e-30)
    log_ar1 = np.log10(s['ar1_threshold'] + 1e-30)

    # Shade the region above AR(1) threshold (significant)
    sig_mask = s['gws'] > s['ar1_threshold']
    ax.fill_between(T_mbp, log_ar1, log_gws,
                    where=sig_mask, alpha=0.15, color=WONG['gold'],
                    interpolate=True)

    ax.plot(T_mbp, log_gws, color='black', lw=1.3, label='GWS')
    ax.plot(T_mbp, log_ar1, color=WONG['grey'], lw=0.8, ls='--', label='AR(1) 95%')

    # Mark peaks: rank-1 gold thick, rank-2/3/4 vermillion thin, rest grey dotted
    n_peak_to_mark = min(6, len(s['peak_idx']))
    annotated_xs = []   # to avoid label overlaps
    for r in range(n_peak_to_mark):
        p_idx = s['peak_idx'][r]
        T_pk = s['peak_periods'][r] * msp
        gws_pk = s['peak_gws'][r]
        col = WONG['gold'] if r == 0 else (WONG['vermil'] if r < 4 else WONG['grey'])
        ls = '-' if r == 0 else ('--' if r < 4 else ':')
        lw = 1.6 if r == 0 else (0.9 if r < 4 else 0.5)
        ax.axvline(T_pk, color=col, lw=lw, alpha=0.85, ls=ls)
        # Annotate top-4 peaks; offset y to avoid overlap
        if r < 4:
            # Skip if too close to already-annotated peak
            too_close = any(abs(np.log10(T_pk) - np.log10(x)) < 0.08 for x in annotated_xs)
            if not too_close:
                y_text = np.log10(gws_pk + 1e-30) + (0.025 + 0.04*(r))
                ax.text(T_pk, y_text, f"{T_pk:.1f}", color=col, fontsize=7.5,
                        ha='center', va='bottom', fontweight='bold')
                annotated_xs.append(T_pk)

    ax.set_xscale('log')
    dom_T = s['peak_periods'][0] * msp
    n_cyc = s['chrom_len_bp']/1e6 / dom_T
    ax.set_title(f"{chrom}    T₁ = {dom_T:.1f} Mbp;  N = {s['N_peaks']:,};  "
                 f"≈ {n_cyc:.1f} cycles",
                 fontsize=9.5, loc='left')
    ax.tick_params(axis='both', labelsize=7.5)
    ax.set_xlim(0.05, 200)
    ax.grid(alpha=0.25, which='both', lw=0.4)

# Hide unused panels
for k in range(n_chrom, len(axes_flat)):
    axes_flat[k].set_visible(False)

# Bottom row x-label, leftmost y-label
for ax in axes_flat[-ncols:]:
    if ax.get_visible():
        ax.set_xlabel('period (Mbp)', fontsize=9)
for r in range(nrows):
    axes_flat[r*ncols].set_ylabel(r'$\log_{10}$ GWS', fontsize=9)

plt.suptitle('Global wavelet spectra per chromosome (K562 v38; ENCFF361VGY; Paul $m=2$)\n'
             'Black: GWS;  grey dashed: AR(1) 95% threshold;  '
             'gold solid: rank-1 dominant period (Mbp);  '
             'vermillion dashed: rank-2--4 harmonics;  grey dotted: lower-rank peaks',
             fontsize=12, y=1.005)
plt.tight_layout()
plt.savefig(OUT / "supplementary_figure_41_per_chrom_spectra.png", dpi=180,
            bbox_inches='tight')
plt.savefig(OUT / "supplementary_figure_41_per_chrom_spectra.svg",
            bbox_inches='tight')
plt.close(fig)
print(f"Saved: supplementary_figure_41_per_chrom_spectra.{{png,svg}}")
print()
print("DONE")
