"""
TEST 2 — Per-annotation harmonic spectra.

Question: Do different functional annotations share the same harmonic structure
   (universal scaffold), or do they have feature-specific harmonics
   (layered functional readouts)?

Method:
- Use the master annotation matrix (500-kb bins) on chr1, chr7, chr19.
- For each track, compute the CWT in 500-kb-bin coordinates with Paul m=2.
- Report the rank-1 (dominant) period and the next 3 local maxima of GWS.
- Compare harmonic ratios across tracks.

Tracks tested:
- cwt_amplitude (the wave itself; baseline)
- son_tsa, lmnb_tsa (radial geometry)
- repliseq_K562 (replication timing)
- log_gene_density (gene density)
- cpg_island_bp (regulatory)
- L1, Alu (repeat classes)
- total_repeat_bp (compartment proxy)

Honest disclosure of caveats:
- The wave was originally computed in PEAK-INDEX coordinates; here we test the
  500-kb bin-index version, so the dominant period in bin-index might differ.
- CWT on a track that's nearly constant or has zero variance will fail; we skip those.
- Local-max detection uses prominence ≥ 0.05 on the log-GWS curve, same as the harmonic_sweep.
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
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'}

# ---- CWT engine that returns full GWS + multi-peak detection ----

def cwt_track_with_harmonics(signal_raw, n_scales=120,
                              period_min=4, period_max=2000,
                              wavelet_order=2, prominence=0.05):
    """CWT on a 1D track (in bin-index coords). Returns full GWS + multi-peak metadata."""
    N = len(signal_raw)
    if N < 32:
        return None
    sigma = np.std(signal_raw, ddof=1)
    if sigma == 0 or not np.isfinite(sigma):
        return None
    signal = (signal_raw - np.nanmean(signal_raw)) / sigma
    signal = np.nan_to_num(signal, nan=0.0, posinf=0.0, neginf=0.0)

    fourier_factor, coi_factor = cp._get_paul_params(wavelet_order)
    period_max_eff = min(period_max, N * 0.8)
    s0 = period_min / fourier_factor
    smax = period_max_eff / 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

    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, props = find_peaks(log_valid, prominence=prominence)
    peak_idx = valid_idx[peaks_local]
    peak_periods = periods[peak_idx]
    peak_gws = gws[peak_idx]
    # also 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,
            'peak_periods': peak_periods, 'peak_gws': peak_gws}


# ---- Run ----

print("="*70)
print("TEST 2: Per-annotation harmonic spectra")
print("="*70)

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

CHROMS = ['chr1', 'chr7', 'chr19']
# Tracks to scan; column name → human label
TRACKS = {
    'cwt_amplitude':   'CWT amplitude (wave proxy)',
    'son_tsa':         'SON TSA-seq (speckle prox.)',
    'lmnb_tsa':        'LMNB1 TSA-seq (lamina prox.)',
    'repliseq_K562':   'Repli-seq (RT)',
    'log_gene_density':'log(gene density)',
    'cpg_island_bp':   'CpG-island bp',
    'L1':              'LINE-1 bp',
    'Alu':             'Alu bp',
    'total_repeat_bp': 'Total repeat bp',
}

results = []
spectra = {}   # (chrom, track) → dict
for chrom in CHROMS:
    chrom_df = df[df['chrom']==chrom].sort_values('bin_mid_hg38').reset_index(drop=True)
    print(f"\n{chrom}: {len(chrom_df)} bins (each 500 kb → {len(chrom_df)*0.5:.1f} Mbp)")
    for track, label in TRACKS.items():
        if track not in chrom_df.columns:
            print(f"  {track}: column missing")
            continue
        sig = chrom_df[track].values.astype(float)
        # Replace NaNs with chromosome mean (avoid distorting; track that's all NaN is skipped by std=0 check)
        if np.all(np.isnan(sig)):
            print(f"  {track}: all NaN")
            continue
        sig = np.where(np.isnan(sig), np.nanmean(sig), sig)
        # Constrain period range to the supra-compartment band:
        #   period_min = 4 bins (2 Mbp; TAD scale)
        #   period_max = 100 bins (50 Mbp; covers the dominant wave for chr1/7/19)
        # This isolates the wave + its harmonics from chromosome-arm-scale envelopes.
        res = cwt_track_with_harmonics(sig, n_scales=120,
                                         period_min=4, period_max=100,
                                         wavelet_order=2, prominence=0.05)
        if res is None:
            print(f"  {track}: skipped (zero variance or too short)")
            continue
        spectra[(chrom, track)] = res
        # Top 4 peaks
        n_peaks = min(4, len(res['peak_periods']))
        rank1 = res['peak_periods'][0]
        for r in range(n_peaks):
            T_bins = res['peak_periods'][r]
            T_mbp = T_bins * 0.5
            ratio = rank1 / T_bins
            results.append({'chrom': chrom, 'track': track, 'label': label,
                            'rank': r,
                            'period_bins': T_bins, 'period_Mbp': T_mbp,
                            'gws': res['peak_gws'][r],
                            'ratio_to_rank1': ratio if r > 0 else 1.0})
        print(f"  {track:18s}  rank1 = {rank1*0.5:.1f} Mbp; "
              f"peaks (Mbp) = {[f'{T*0.5:.1f}' for T in res['peak_periods'][:n_peaks]]}")

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

# ---- Cross-track comparison: are rank-1 periods consistent? ----

print("\n--- Rank-1 period (Mbp) by track and chromosome ---")
piv = res_df[res_df['rank']==0].pivot_table(
    index='track', columns='chrom', values='period_Mbp', aggfunc='first')
piv = piv.reindex([k for k in TRACKS if k in piv.index])
print(piv.to_string(float_format=lambda x: f"{x:.1f}"))

print("\n--- Top harmonic ratios (rank2/rank1, rank3/rank1) by track and chromosome ---")
piv2 = res_df[(res_df['rank']==1) & (res_df['ratio_to_rank1']>=2) & (res_df['ratio_to_rank1']<=8)]
piv2_t = piv2.pivot_table(index='track', columns='chrom', values='ratio_to_rank1', aggfunc='first')
piv2_t = piv2_t.reindex([k for k in TRACKS if k in piv2_t.index])
print(piv2_t.to_string(float_format=lambda x: f"{x:.2f}x"))

# ---- Plot 3x9 grid: spectra ----
fig, axes = plt.subplots(len(CHROMS), len(TRACKS), figsize=(22, 7.5),
                         sharex=False, sharey=False)
COLOR = {'cwt_amplitude': WONG['gold'], 'son_tsa': WONG['vermil'],
         'lmnb_tsa': WONG['blue'], 'repliseq_K562': WONG['green'],
         'log_gene_density': WONG['pink'], 'cpg_island_bp': WONG['sky'],
         'L1': WONG['grey'], 'Alu': WONG['grey'], 'total_repeat_bp': WONG['grey']}
for i, chrom in enumerate(CHROMS):
    for j, track in enumerate(TRACKS):
        ax = axes[i, j] if len(CHROMS) > 1 else axes[j]
        s = spectra.get((chrom, track))
        if s is None:
            ax.set_visible(False)
            continue
        T_mbp = s['periods'] * 0.5
        ax.semilogx(T_mbp, np.log10(s['gws']), color=COLOR.get(track, 'k'), lw=1.2)
        # mark top-4 peaks
        for r in range(min(4, len(s['peak_periods']))):
            T = s['peak_periods'][r] * 0.5
            ax.axvline(T, color='black' if r==0 else 'tab:red',
                       lw=1.0 if r==0 else 0.5,
                       ls='-' if r==0 else '--', alpha=0.6 if r==0 else 0.5)
        if i == 0:
            ax.set_title(TRACKS[track], fontsize=8.5)
        if j == 0:
            ax.set_ylabel(f'{chrom}\nlog$_{{10}}$ GWS', fontsize=9)
        if i == len(CHROMS)-1:
            ax.set_xlabel('period (Mbp)', fontsize=8)
        ax.tick_params(axis='both', labelsize=7)
        ax.grid(alpha=0.25)

plt.suptitle('Per-annotation harmonic spectra (K562 v38, 500-kb bins, Paul m=2)\n'
             'Solid black line = rank-1 (dominant); red dashed = rank-2..4 harmonics',
             fontsize=11, y=1.005)
plt.tight_layout()
plt.savefig(OUT / "test2_annotation_harmonics_spectra.png", dpi=200, bbox_inches='tight')
plt.savefig(OUT / "test2_annotation_harmonics_spectra.svg", bbox_inches='tight')
plt.close(fig)
print(f"\nSaved: test2_annotation_harmonics_spectra.{{png,svg}}")
print("\nDONE TEST 2")
