"""
Harmonic / parameter-sweep analysis on K562 v38 narrowPeak.

For three representative chromosomes (chr1 large, chr7 medium, chr19 small):
  1. Sweep wavelet family / m parameter (Paul m=1-6, Morlet, Mexican hat)
  2. Sweep period range
  3. Sweep peak signalValue threshold

For each setting, extract ALL local GWS maxima outside the cone of influence
and test for integer-ratio harmonic relationships.

Output: harmonic_sweep_results.csv + multi-peak GWS figure.
"""
import importlib.util, sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
from pathlib import Path

OUTDIR = Path("/sessions/kind-brave-babbage/mnt/Wavelet/go_rerun_v2_2026-04-30")
PEAKS = "/sessions/kind-brave-babbage/mnt/Wavelet/k562_regen_option_A_2026-04-18/ENCFF361VGY.bed"

# Reuse the canonical chromperiod CWT 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)

# ---- Add an enhanced CWT function that returns ALL local GWS maxima ----

def consecutive_peaks_cwt_full(signal_raw, wavelet_order=2, n_scales=160,
                                period_min=4, period_max=14000,
                                significance_level=0.95):
    """Modified CWT that returns the full GWS curve + all local maxima outside COI."""
    N = len(signal_raw)
    sigma = np.std(signal_raw, ddof=1)
    if sigma == 0:
        raise ValueError("Signal has zero variance")
    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
    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])

    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

    coi = cp.compute_coi(N, 1.0, coi_factor)
    from scipy.stats import chi2
    chi2_crit = chi2.ppf(significance_level, 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
        if mask.sum() > 0:
            gws[j] = np.mean(power[j, mask])
        else:
            gws[j] = np.nan

    # Find ALL local maxima of GWS, not just the global one
    valid_mask = np.isfinite(gws) & (gws > 0)
    valid_indices = np.where(valid_mask)[0]
    if not len(valid_indices):
        return {'periods': periods, 'gws': gws, 'all_peaks': [], 'ar1_threshold': ar1_threshold,
                'ar1_alpha': ar1_alpha, 'n_peaks': N}
    valid_gws = gws[valid_mask]
    log_valid_gws = np.log(valid_gws + 1e-30)

    # Find peaks with prominence
    peak_idx_local, props = find_peaks(log_valid_gws, prominence=0.05)
    # Translate local index back to scale index
    peak_idx = valid_indices[peak_idx_local]
    peak_periods = periods[peak_idx]
    peak_gws = gws[peak_idx]
    # Also include the global maximum if not already
    global_max_idx = valid_indices[np.argmax(valid_gws)]
    if global_max_idx not in peak_idx:
        peak_idx = np.append(peak_idx, global_max_idx)
        peak_periods = np.append(peak_periods, periods[global_max_idx])
        peak_gws = np.append(peak_gws, gws[global_max_idx])
    # Sort peaks by GWS descending
    order = np.argsort(-peak_gws)
    all_peaks = [
        {'period_indices': float(peak_periods[i]), 'gws': float(peak_gws[i]),
          'above_ar1': bool(peak_gws[i] > ar1_threshold[peak_idx[i]])}
        for i in order
    ]

    return {
        'periods': periods, 'gws': gws, 'ar1_threshold': ar1_threshold,
        'all_peaks': all_peaks, 'ar1_alpha': ar1_alpha, 'n_peaks': N,
        'coi': coi,
    }


# ---- Run sweeps ----

# Load all peaks
print("Loading K562 v38 peaks...")
peaks_all = cp.load_narrowpeak(PEAKS)
print(f"Total peaks: {len(peaks_all):,}")

CHROMS = ['chr1', 'chr7', 'chr19']
SVAL_THRESHOLDS = ['top100', 'top75', 'top50', 'top25']

results = []

for chrom in CHROMS:
    cp_data = peaks_all[peaks_all['chrom'] == chrom].sort_values('start').reset_index(drop=True)
    full_n = len(cp_data)
    chrom_len_bp = int(cp_data['end'].max())

    # ---- 1. Wavelet family sweep ----
    for m in [1, 2, 3, 4, 5, 6]:
        signal = cp_data['signalValue'].values.astype(float)
        positions = ((cp_data['start'].values + cp_data['end'].values) / 2.0)
        spacings = np.diff(positions)
        median_spacing_bp = float(np.median(spacings[spacings > 0]))
        try:
            res = consecutive_peaks_cwt_full(signal, wavelet_order=m,
                                              n_scales=160,
                                              period_min=4, period_max=14000)
        except ValueError:
            continue
        for rank, p in enumerate(res['all_peaks'][:5]):
            T_mbp = p['period_indices'] * median_spacing_bp / 1e6
            n_cycles = chrom_len_bp / 1e6 / T_mbp if T_mbp > 0 else np.nan
            results.append({
                'chrom': chrom, 'sweep': 'wavelet_m', 'setting': f'Paul_m{m}',
                'rank': rank, 'period_indices': p['period_indices'],
                'period_Mbp': T_mbp, 'cycles_per_chrom': n_cycles,
                'gws': p['gws'], 'above_ar1': p['above_ar1'],
                'n_peaks_used': full_n, 'median_spacing_bp': median_spacing_bp,
                'ar1_alpha': res['ar1_alpha'],
            })

    # ---- 2. Period range sweep ----
    for p_max in [4000, 7000, 14000]:
        signal = cp_data['signalValue'].values.astype(float)
        positions = ((cp_data['start'].values + cp_data['end'].values) / 2.0)
        spacings = np.diff(positions)
        median_spacing_bp = float(np.median(spacings[spacings > 0]))
        try:
            res = consecutive_peaks_cwt_full(signal, wavelet_order=2, n_scales=160,
                                              period_min=4, period_max=p_max)
        except ValueError:
            continue
        for rank, p in enumerate(res['all_peaks'][:5]):
            T_mbp = p['period_indices'] * median_spacing_bp / 1e6
            n_cycles = chrom_len_bp / 1e6 / T_mbp if T_mbp > 0 else np.nan
            results.append({
                'chrom': chrom, 'sweep': 'period_range', 'setting': f'pmax{p_max}',
                'rank': rank, 'period_indices': p['period_indices'],
                'period_Mbp': T_mbp, 'cycles_per_chrom': n_cycles,
                'gws': p['gws'], 'above_ar1': p['above_ar1'],
                'n_peaks_used': full_n, 'median_spacing_bp': median_spacing_bp,
                'ar1_alpha': res['ar1_alpha'],
            })

    # ---- 3. signalValue threshold sweep (top X% by signal) ----
    for thresh_label in SVAL_THRESHOLDS:
        if thresh_label == 'top100':
            sub = cp_data
        else:
            pct = int(thresh_label.replace('top', ''))
            cutoff = cp_data['signalValue'].quantile(1 - pct/100)
            sub = cp_data[cp_data['signalValue'] >= cutoff]
        if len(sub) < 100:
            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]))
        try:
            res = consecutive_peaks_cwt_full(signal, wavelet_order=2, n_scales=160,
                                              period_min=4, period_max=14000)
        except ValueError:
            continue
        for rank, p in enumerate(res['all_peaks'][:5]):
            T_mbp = p['period_indices'] * median_spacing_bp / 1e6
            n_cycles = chrom_len_bp / 1e6 / T_mbp if T_mbp > 0 else np.nan
            results.append({
                'chrom': chrom, 'sweep': 'signal_threshold', 'setting': thresh_label,
                'rank': rank, 'period_indices': p['period_indices'],
                'period_Mbp': T_mbp, 'cycles_per_chrom': n_cycles,
                'gws': p['gws'], 'above_ar1': p['above_ar1'],
                'n_peaks_used': len(sub), 'median_spacing_bp': median_spacing_bp,
                'ar1_alpha': res['ar1_alpha'],
            })

df = pd.DataFrame(results)
df.to_csv(OUTDIR / "harmonic_sweep_results.csv", index=False)
print(f"\nWrote harmonic_sweep_results.csv: {len(df)} rows")

# ---- Harmonic ratio analysis: are secondary peaks integer-ratio harmonics? ----
print("\n=== Harmonic ratio test: rank-1 / rank-0 period ratios ===")
print(f"{'Chrom':<8s}{'Sweep':<22s}{'Setting':<14s}{'T0_Mbp':>10s}{'T1_Mbp':>10s}{'ratio':>10s}{'integer?':>10s}")
for chrom in CHROMS:
    for sweep in ['wavelet_m', 'period_range', 'signal_threshold']:
        for setting in df[(df['chrom']==chrom) & (df['sweep']==sweep)]['setting'].unique():
            sub = df[(df['chrom']==chrom) & (df['sweep']==sweep) & (df['setting']==setting)]
            if len(sub) < 2:
                continue
            T0 = sub[sub['rank']==0]['period_Mbp'].iloc[0] if len(sub[sub['rank']==0]) else np.nan
            T1 = sub[sub['rank']==1]['period_Mbp'].iloc[0] if len(sub[sub['rank']==1]) else np.nan
            if pd.isna(T1) or T1 == 0 or T0 == 0:
                continue
            ratio = T0 / T1
            # Test integer ratio (close to 2, 3, 4...)
            nearest_int = round(ratio)
            integer_test = abs(ratio - nearest_int) / nearest_int < 0.10 if nearest_int >= 2 else False
            print(f"{chrom:<8s}{sweep:<22s}{setting:<14s}{T0:>10.2f}{T1:>10.2f}{ratio:>10.2f}"
                  f"{('~'+str(nearest_int)+'×' if integer_test else 'no'):>10s}")
