"""
02_run_chromperiod_per_sample.py
=================================

Run the canonical consecutive-peak CWT (Paul m=2, substrate (a)) on a single
ATAC-seq narrowPeak BED and emit the per-chromosome summary + 500-kb phase map
that the manuscript pipeline uses.

This is a direct adaptation of `chromperiod_phase_extraction_standalone.py`
(the audited canonical reference). All wavelet parameters, normalisation,
significance test, and phase-binning rules are byte-for-byte preserved. The
ONLY changes are:
  - Generalised to take an arbitrary input BED via CLI argument.
  - Writes per-sample output rows tagged with `line` and `disease` columns
    so downstream HGPS-vs-control aggregation can stitch results across
    samples without re-running CWT.
  - No K562 hilbert-eig comparison (was a one-off diagnostic in the standalone).
  - Uses chr1..chr22 only (autosomes — manuscript discipline; chrX excluded
    from per-line averages because progerin perturbations have known X-chromosome
    asymmetries that would confound the autosome-pooled mixed model).

Substrate-(a) verification, locked in PRE_COMMITTED_DISPOSITION.md:
  * Signal           : narrowPeak `signalValue` column at each peak
  * Index            : peak number (NOT base pairs)
  * Wavelet          : Paul m=2; fourier_factor = 4*pi/(2*m+1) = 4*pi/5 = 2.5133
  * Period axis      : "peaks"; converted to Mbp only via median inter-peak spacing
  * Significance     : AR(1) red-noise null at 95% (chi^2 df=2), outside COI
  * Phase binning    : amplitude-weighted circular mean over 500-kb genomic bins

Author: Wolf H. Gebhardt / chromperiod v0.1.0 recipe
Adapted: 2026-05-02 for HGPS GSE150136 pre-submission addition
"""

import argparse
import math
import os
import sys

import numpy as np
import pandas as pd
from scipy import stats as scipy_stats

# ===================================================================
# CHROMPERIOD CORE FUNCTIONS (audited; do not modify without
# updating the substrate (a) lock in PRE_COMMITTED_DISPOSITION.md)
# ===================================================================


def _paul_psi_hat(s, omega, m=2):
    """Paul wavelet in Fourier space (T&C eq. 6)."""
    norm = (2.0 ** m) / np.sqrt(m * float(math.factorial(2 * m - 1)))
    exponent = np.clip(-s * omega, -500.0, 0.0)
    return np.where(omega > 0, norm * (s * omega) ** m * np.exp(exponent), 0.0)


def _get_paul_params(m=2):
    """Return (fourier_factor, coi_factor) for Paul wavelet order m.

    Paul m=2 -> fourier_factor = 4*pi/5 = 2.5133, coi_factor = ff/sqrt(2).
    """
    fourier_factor = 4.0 * np.pi / (2.0 * m + 1.0)
    coi_factor = fourier_factor / np.sqrt(2.0)
    return fourier_factor, coi_factor


def estimate_ar1(signal):
    N = len(signal)
    if N < 3:
        return 0.0
    x = signal - np.mean(signal)
    r0 = np.dot(x, x) / N
    r1 = np.dot(x[:-1], x[1:]) / N
    if r0 == 0:
        return 0.0
    return float(np.clip(r1 / r0, 0.0, 0.999))


def compute_coi(N, delta_t, coi_factor):
    n = np.arange(N)
    return coi_factor * delta_t * np.maximum(1, np.minimum(n + 1, N - n))


def consecutive_peaks_cwt_inline(signal_raw, wavelet_order=2, n_scales=80,
                                 period_min=10, period_max=7000,
                                 significance_level=0.95):
    """Substrate (a) recipe: consecutive-peak CWT on narrowPeak signalValue."""
    from scipy.stats import chi2

    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 = estimate_ar1(signal)
    fourier_factor, coi_factor = _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(max(N, 2))))
    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(_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 = compute_coi(N, 1.0, coi_factor)

    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

    significance = power > ar1_threshold[:, np.newaxis]

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

    valid_gws = np.where(np.isfinite(gws), gws, -np.inf)
    dom_idx = int(np.argmax(valid_gws))
    dominant_period = float(periods[dom_idx])

    n_outside = outside_coi.sum()
    n_sig_outside = (significance & outside_coi).sum()
    sig95 = float(n_sig_outside) / float(n_outside) if n_outside > 0 else 0.0

    return {
        'scales': scales,
        'periods': periods,
        'dominant_period': dominant_period,
        'sig95': sig95,
        'power': power,
        'gws': gws,
        'gws_max': float(np.nanmax(gws)),
        'ar1_alpha': ar1_alpha,
        'signal': signal,
        'n_peaks': N,
    }


def extract_cwt_phase_inline(signal, scales, periods, wavelet_order=2,
                             dominant_period=None, bandwidth_log10=0.15):
    """Carrier phase at the dominant period band (substrate (a))."""
    N = len(signal)
    N_pad = int(2 ** np.ceil(np.log2(max(N, 2))))
    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])

    if dominant_period is None:
        dominant_period = periods[len(periods) // 2]

    log_periods = np.log10(periods)
    log_dom = np.log10(dominant_period)
    in_band = np.abs(log_periods - log_dom) <= bandwidth_log10
    band_indices = np.where(in_band)[0]
    if len(band_indices) == 0:
        band_indices = np.array([np.argmin(np.abs(log_periods - log_dom))])
        in_band = np.zeros(len(periods), dtype=bool)
        in_band[band_indices[0]] = True

    reconstructed_complex = np.zeros(N, dtype=complex)
    delta_j = np.log(scales[1] / scales[0])

    for j, (s, in_b) in enumerate(zip(scales, in_band)):
        if not in_b:
            continue
        psi_vals = np.conj(_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]
        reconstructed_complex += W_n / np.sqrt(s)

    C_delta = 1.132
    psi_0 = 1.0
    reconstructed_complex = (delta_j * np.sqrt(1.0) / (C_delta * psi_0)) * reconstructed_complex

    return {
        'phase': np.angle(reconstructed_complex),
        'amplitude': np.abs(reconstructed_complex),
    }


# ===================================================================
# I/O + per-sample wrapper
# ===================================================================


def load_narrowpeak(filepath):
    """Load BED6+4 narrowPeak. Keep autosomes only (manuscript discipline)."""
    cols = ['chrom', 'start', 'end', 'name', 'score', 'strand',
            'signalValue', 'pValue', 'qValue', 'peak']
    df = pd.read_csv(filepath, sep='\t', header=None, names=cols,
                     comment='#', low_memory=False)
    valid_chroms = [f'chr{i}' for i in range(1, 23)]  # autosomes only
    df = df[df['chrom'].isin(valid_chroms)].copy()
    return df


def process_chromosome(peaks_df, chrom, line, disease,
                       wavelet_order=2, n_scales=80, bin_size=500_000):
    cpk = peaks_df[peaks_df['chrom'] == chrom].sort_values('start').copy()
    if len(cpk) < 100:
        return None, None  # too few peaks for CWT

    positions = ((cpk['start'].values + cpk['end'].values) / 2.0)
    signal_raw = cpk['signalValue'].values.astype(float)

    spacings = np.diff(positions)
    spacings = spacings[spacings > 0]
    median_spacing_bp = float(np.median(spacings))

    # Substrate (a) period bounds: period_min=10 peaks; period_max extends to
    # the chromosome length in peak space so the long-wavelength dominants
    # detected in v38 canonical are not clipped (e.g. chr2 dom_period =
    # 11,406 peaks on N_peaks = 12,881 chrom). The CWT does not have a hard
    # Nyquist cap at N/2 the way FFT does; long scales are penalised by the
    # cone-of-influence (COI) outside-COI mask in the GWS.
    period_max_chrom = max(7000, len(cpk))

    cwt = consecutive_peaks_cwt_inline(
        signal_raw, wavelet_order=wavelet_order, n_scales=n_scales,
        period_min=10, period_max=period_max_chrom)

    dom_period_peaks = cwt['dominant_period']
    dom_period_mbp = dom_period_peaks * median_spacing_bp / 1e6
    # Manuscript-canonical N_cycles: genomic coordinate L_bp / T_bp (the
    # ~5 cycles-per-chromosome figure in Fig 2d). chrom_len_bp uses peak
    # positions max as a lower bound on chromosome length.
    chrom_len_bp = float(positions.max()) - float(positions.min())
    n_cycles_genomic = chrom_len_bp / (dom_period_mbp * 1e6) if dom_period_mbp > 0 else float('nan')
    # Peak-coord cycle count retained as diagnostic only.
    n_cycles_peakcoord = float(len(cpk)) / dom_period_peaks
    n_cycles = n_cycles_genomic  # primary metric

    # Phase
    phase_result = extract_cwt_phase_inline(
        cwt['signal'], cwt['scales'], cwt['periods'],
        wavelet_order=wavelet_order, dominant_period=dom_period_peaks,
        bandwidth_log10=0.15)
    peak_phase = phase_result['phase']
    peak_amp = phase_result['amplitude']

    # Substrate-(a) sanity: median phase advance per peak should equal
    # 2*pi / dominant_period_in_peaks. Save the ratio for the audit.
    unwrapped = np.unwrap(peak_phase)
    adv_per_peak = float(np.median(np.abs(np.diff(unwrapped))))
    expected_adv = 2.0 * np.pi / dom_period_peaks
    advance_ratio = adv_per_peak / expected_adv if expected_adv else float('nan')

    # 500-kb amplitude-weighted circular-mean phase map
    chrom_len = int(positions.max()) + bin_size
    bin_starts = np.arange(0, chrom_len, bin_size)
    rows = []
    for bs in bin_starts:
        be = bs + bin_size
        bm = bs + bin_size // 2
        mask = (positions >= bs) & (positions < be)
        n_peaks_in_bin = int(mask.sum())
        if n_peaks_in_bin > 0:
            amps = peak_amp[mask]
            phases = peak_phase[mask]
            cx = float(np.sum(amps * np.cos(phases)))
            cy = float(np.sum(amps * np.sin(phases)))
            bin_phase = float(np.arctan2(cy, cx))
            bin_amp = float(np.sqrt(cx ** 2 + cy ** 2) / (np.sum(amps) + 1e-30))
        else:
            bin_phase = float('nan')
            bin_amp = float('nan')
        rows.append({
            'line': line, 'disease': disease, 'chrom': chrom,
            'start': int(bs), 'end': int(be), 'bin_mid': int(bm),
            'cwt_phase': bin_phase, 'cwt_amplitude': bin_amp,
            'n_peaks_in_bin': n_peaks_in_bin,
            'dom_period_Mbp': dom_period_mbp,
        })
    phase_map_chrom = pd.DataFrame(rows)

    summary_row = {
        'line': line, 'disease': disease, 'chrom': chrom,
        'n_peaks': int(len(cpk)),
        'median_inter_peak_bp': median_spacing_bp,
        'chrom_len_bp_observed': chrom_len_bp,
        'ar1_alpha': cwt['ar1_alpha'],
        'sig95': cwt['sig95'],
        'gws_max': cwt['gws_max'],
        'dom_period_peaks': dom_period_peaks,
        'dom_period_Mbp': dom_period_mbp,
        'N_cycles': n_cycles,                  # genomic L_bp / T_bp
        'N_cycles_peakcoord': n_cycles_peakcoord,
        'phase_advance_ratio': advance_ratio,
    }
    return summary_row, phase_map_chrom


def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument('--bed', required=True, help='Input narrowPeak BED file')
    p.add_argument('--line', required=True, help='Sample line ID (e.g. HGADFN143)')
    p.add_argument('--disease', required=True, choices=['HGPS', 'control'],
                   help='Disease label for downstream aggregation')
    p.add_argument('--out-summary', required=True,
                   help='Output CSV path: per-chromosome summary')
    p.add_argument('--out-phase-map', required=True,
                   help='Output CSV path: 500-kb phase map (long format)')
    p.add_argument('--n-scales', type=int, default=80,
                   help='CWT scales (default 80, manuscript canonical)')
    p.add_argument('--period-min', type=int, default=10,
                   help='Min period in peaks (default 10)')
    p.add_argument('--period-max', type=int, default=7000,
                   help='Max period in peaks (default 7000)')
    args = p.parse_args()

    print(f"[chromperiod] Loading narrowPeak: {args.bed}", flush=True)
    peaks = load_narrowpeak(args.bed)
    print(f"[chromperiod] Total autosomal peaks: {len(peaks):,}", flush=True)
    print(f"[chromperiod] Per-chromosome peak counts:", flush=True)
    for chrom, n in peaks.groupby('chrom').size().items():
        print(f"           {chrom}: {n:,}", flush=True)

    chroms = [f'chr{i}' for i in range(1, 23)]
    summary_rows = []
    phase_chunks = []
    for chrom in chroms:
        s, m = process_chromosome(
            peaks, chrom, line=args.line, disease=args.disease,
            n_scales=args.n_scales)
        if s is None:
            print(f"  {chrom}: SKIPPED (too few peaks)", flush=True)
            continue
        print(f"  {chrom}: n_peaks={s['n_peaks']:,} sig95={s['sig95']:.3f} "
              f"dom_period={s['dom_period_peaks']:.0f} peaks "
              f"({s['dom_period_Mbp']:.1f} Mbp) "
              f"N={s['N_cycles']:.2f} alpha={s['ar1_alpha']:.3f} "
              f"phase_adv_ratio={s['phase_advance_ratio']:.3f}", flush=True)
        summary_rows.append(s)
        phase_chunks.append(m)

    summary = pd.DataFrame(summary_rows)
    phase_map = pd.concat(phase_chunks, ignore_index=True)

    os.makedirs(os.path.dirname(os.path.abspath(args.out_summary)) or '.',
                exist_ok=True)
    summary.to_csv(args.out_summary, index=False)
    phase_map.to_csv(args.out_phase_map, index=False)

    print(f"\n[chromperiod] Saved per-chrom summary: {args.out_summary}",
          flush=True)
    print(f"[chromperiod] Saved phase map         : {args.out_phase_map}",
          flush=True)

    # Top-level QC numbers for the audit
    print("\n=== Per-sample QC ===", flush=True)
    print(f"  Sample           : {args.line} ({args.disease})", flush=True)
    print(f"  Total peaks      : {len(peaks):,}", flush=True)
    print(f"  Chroms processed : {len(summary)}/22", flush=True)
    if len(summary):
        print(f"  Mean sig95       : {summary['sig95'].mean():.3f}", flush=True)
        print(f"  Mean dom_period  : {summary['dom_period_Mbp'].mean():.1f} Mbp",
              flush=True)
        print(f"  Median N_cycles  : {summary['N_cycles'].median():.2f}",
              flush=True)
        print(f"  Mean phase_advance_ratio (should be ~1.0 by construction): "
              f"{summary['phase_advance_ratio'].mean():.3f}", flush=True)


if __name__ == '__main__':
    main()
