"""
Analyse voice onset timing from WAV recordings at multiple RMS thresholds.

Mirrors the detection logic in decision.js / record.js:
  - 4096-sample chunks (ScriptProcessor buffer size)
  - RMS per chunk > threshold → onset at that chunk's start time
  - Reports onset as seconds from start of file

Thresholds tested: 0.005, 0.01 (JS default), 0.015, 0.02, 0.025, 0.03, 0.05

Output columns:
  session_code, participant_code, phase, filename,
  sample_rate, duration_s, rms_max,
  onset_<threshold>  (seconds from file start, blank if never detected)

Usage:
  python analyse_voice_onset.py
"""

import csv
import re
import wave
import numpy as np
from pathlib import Path

CHUNK_SIZE  = 4096   # must match ScriptProcessor bufferSize in decision.js / record.js
THRESHOLDS  = [0.005, 0.01, 0.015, 0.02, 0.025, 0.03, 0.05]

SCRIPTS_DIR    = Path(__file__).parent
DATA_DIR       = SCRIPTS_DIR.parent.parent / 'data'
RECORDINGS_DIR = DATA_DIR / 'recordings'
OUT_PATH       = DATA_DIR / 'processed' / 'timestamps.csv'

# {session}_{participant}_{phase}.{wav|webm}  — phase in baseline/prelim/final
FILENAME_RE = re.compile(r'^([a-z0-9]+)_([a-z0-9]+)_(baseline|prelim|final)\.(wav|webm)$')


def read_wav_mono_f32(path: Path):
    """Return (sample_rate, float32 mono array in [-1, 1])."""
    with wave.open(str(path), 'rb') as wf:
        n_ch      = wf.getnchannels()
        sampwidth = wf.getsampwidth()
        rate      = wf.getframerate()
        raw       = wf.readframes(wf.getnframes())

    if sampwidth == 2:
        samples = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0
    elif sampwidth == 4:
        samples = np.frombuffer(raw, dtype=np.int32).astype(np.float32) / 2_147_483_648.0
    elif sampwidth == 1:
        samples = (np.frombuffer(raw, dtype=np.uint8).astype(np.float32) - 128) / 128.0
    else:
        raise ValueError(f'Unsupported sample width: {sampwidth}')

    if n_ch > 1:
        samples = samples.reshape(-1, n_ch).mean(axis=1)

    return rate, samples


def chunk_rms(samples: np.ndarray, chunk_size: int = CHUNK_SIZE) -> np.ndarray:
    """RMS for each non-overlapping chunk; trailing samples are dropped."""
    n = (len(samples) // chunk_size) * chunk_size
    chunks = samples[:n].reshape(-1, chunk_size)
    return np.sqrt(np.mean(chunks ** 2, axis=1))


def find_onset(rms_per_chunk: np.ndarray, rate: int, threshold: float) -> float | None:
    """Return seconds from file start of first chunk exceeding threshold, or None."""
    idx = np.argmax(rms_per_chunk > threshold)
    if not (rms_per_chunk > threshold).any():
        return None
    return float(idx * CHUNK_SIZE / rate)


def threshold_col(t: float) -> str:
    return f'onset_{t:.3f}'.replace('.', 'p')


def main():
    OUT_PATH.parent.mkdir(parents=True, exist_ok=True)

    wav_files = sorted(RECORDINGS_DIR.glob('*.wav'))
    if not wav_files:
        print(f'No WAV files found in {RECORDINGS_DIR}')
        return

    t_cols     = [threshold_col(t) for t in THRESHOLDS]
    fieldnames = [
        'session_code', 'participant_code', 'phase', 'filename',
        'sample_rate', 'duration_s', 'rms_max',
        *t_cols,
    ]

    rows = []
    for path in wav_files:
        m = FILENAME_RE.match(path.name)
        if not m:
            print(f'Skipping {path.name} — unexpected filename')
            continue

        session_code, participant_code, phase, _ = m.groups()

        try:
            rate, samples = read_wav_mono_f32(path)
        except Exception as e:
            print(f'Error reading {path.name}: {e}')
            continue

        rms = chunk_rms(samples)
        row = {
            'session_code':     session_code,
            'participant_code': participant_code,
            'phase':            phase,
            'filename':         path.name,
            'sample_rate':      rate,
            'duration_s':       round(len(samples) / rate, 3),
            'rms_max':          round(float(rms.max()) if len(rms) else 0.0, 5),
        }
        for t, col in zip(THRESHOLDS, t_cols):
            onset = find_onset(rms, rate, t)
            row[col] = '' if onset is None else round(onset, 4)

        rows.append(row)

    _phase_order = {'baseline': 0, 'prelim': 1, 'final': 2}
    rows.sort(key=lambda r: (r['session_code'], r['participant_code'],
                             _phase_order.get(r['phase'], 9)))

    with open(OUT_PATH, 'w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    print(f'Wrote {len(rows)} rows → {OUT_PATH}')
    print()
    print(f'{"Threshold":>12}  {"Detected":>10}  {"Coverage":>10}')
    print('-' * 38)
    for t, col in zip(THRESHOLDS, t_cols):
        n = sum(1 for r in rows if r[col] != '')
        marker = '  ← JS default' if t == 0.01 else ''
        print(f'{t:>12.3f}  {n:>7}/{len(rows)}  {n/len(rows)*100:>8.1f}%{marker}')


if __name__ == '__main__':
    main()