#!/usr/bin/env python3
"""Offline residual analysis for scripts/gsharp_ab_probe.mjs.

Input is the probe's trace plus its two-channel WebM.  The right channel is
only the known accompaniment reference; the added solo is intentionally
identified by its logged schedule rather than an unavailable third channel.
"""
import argparse
import json
import math
import pathlib
import subprocess

import numpy as np

HARMONICS = (103.826174, 207.652349, 311.478523, 415.304698)


def decode(path):
    raw = subprocess.check_output([
        'ffmpeg', '-v', 'error', '-i', str(path), '-f', 'f32le', '-acodec',
        'pcm_f32le', '-ac', '2', '-ar', '48000', 'pipe:1',
    ])
    return np.frombuffer(raw, dtype='<f4').reshape(-1, 2), 48000


def idx(ms, capture_ms, sr):
    return round((ms - capture_ms) * sr / 1000)


def rms(x):
    return float(np.sqrt(np.mean(np.square(x)))) if len(x) else 0.0


def db(x):
    return 20 * math.log10(max(x, 1e-12))


def best_delay(ref, mic, trials, capture_ms, sr):
    # Correlate actual known-output intervals only. Direct waveform correlation
    # can invert under the room response, so maximize absolute normalized value
    # and retain its signed value in the report.
    lags = range(0, round(.30 * sr) + 1)
    scores = np.zeros(len(lags), dtype=np.float64)
    counts = np.zeros(len(lags), dtype=np.int32)
    for tr in trials:
        if tr['condition'] != 'A':
            continue
        a = idx(tr['atMs'], capture_ms, sr)
        n = round(1.65 * sr)
        x = ref[a:a+n]
        if len(x) != n:
            continue
        x = x - x.mean()
        for j, lag in enumerate(lags):
            y = mic[a+lag:a+lag+n]
            if len(y) != n:
                continue
            y = y - y.mean()
            den = np.linalg.norm(x) * np.linalg.norm(y)
            if den:
                scores[j] += float(np.dot(x, y) / den)
                counts[j] += 1
    valid = counts > 0
    mean = np.divide(scores, counts, out=np.zeros_like(scores), where=valid)
    pick = int(np.argmax(np.abs(mean)))
    return {'lagMs': lags[pick] * 1000 / sr, 'corr': float(mean[pick]),
            'searchMs': [0, 300], 'controlTrials': int(counts[pick])}


def fit_transfer(ref, mic, train, capture_ms, sr, seconds=2.8):
    n = round(seconds * sr)
    fft_n = 1 << (n - 1).bit_length()
    numerator = np.zeros(fft_n // 2 + 1, dtype=np.complex128)
    denominator = np.zeros(fft_n // 2 + 1, dtype=np.float64)
    used = 0
    for tr in train:
        a = idx(tr['atMs'], capture_ms, sr)
        x, y = ref[a:a+n], mic[a:a+n]
        if len(x) != n or len(y) != n:
            continue
        # A shared, fixed Hann makes the estimate a local complex STFT transfer
        # rather than a circular whole-file subtraction.
        w = np.hanning(n)
        X = np.fft.rfft(x * w, fft_n)
        Y = np.fft.rfft(y * w, fft_n)
        numerator += Y * np.conj(X)
        denominator += np.real(X * np.conj(X))
        used += 1
    if used < 1:
        raise ValueError('No complete control segments available for transfer fit')
    ridge = max(float(np.median(denominator[denominator > 0])) * 1e-4, 1e-12)
    return numerator / (denominator + ridge), fft_n, n, used, ridge


def residual_for(ref, mic, tr, h, fft_n, n, capture_ms, sr):
    a = idx(tr['atMs'], capture_ms, sr)
    x, y = ref[a:a+n], mic[a:a+n]
    if len(x) != n or len(y) != n:
        return None
    w = np.hanning(n)
    # Predict only in the well-supported interior. Dividing by the window would
    # amplify endpoints, so metrics below never inspect the first/last 250 ms.
    pred = np.fft.irfft(np.fft.rfft(x * w, fft_n) * h, fft_n)[:n]
    return x, y, y * w - pred


def band_rms(x, sr, center, half_width=5.0):
    if not len(x):
        return 0.0
    spec = np.fft.rfft(x * np.hanning(len(x)))
    fr = np.fft.rfftfreq(len(x), 1 / sr)
    keep = (fr >= center - half_width) & (fr <= center + half_width)
    if not np.any(keep):
        return 0.0
    filtered = np.fft.irfft(spec * keep, len(x))
    return rms(filtered)


def summary(values):
    if not values:
        return {'n': 0, 'meanDbfs': None, 'sdDb': None}
    powers = np.array([v * v for v in values])
    return {'n': len(values), 'meanDbfs': db(math.sqrt(float(np.mean(powers)))),
            'sdDb': float(np.std([db(v) for v in values], ddof=1)) if len(values) > 1 else 0.0}


def salience_for(events, tr, arrival_ms):
    # salience rows are MIDI 33..96, so G#2/3/4/5 map to 11/23/35/47.
    rows = [11, 23, 35, 47]
    frames = [e for e in events if e.get('type') == 'follow.frame'
              and arrival_ms + 30 <= e.get('t', -1) <= arrival_ms + 360
              and e.get('pos') == 1 and isinstance(e.get('sal'), list)]
    if not frames:
        return {'frames': 0, 'rows': {}, 'evMean': None, 'pcomMean': None, 'pcurMean': None}
    out_rows = {}
    for row in rows:
        vals = [f['sal'][row] for f in frames if len(f['sal']) > row]
        out_rows[str(row + 33)] = {'mean': float(np.mean(vals)), 'max': int(max(vals))} if vals else None
    nums = lambda key: [f[key] for f in frames if isinstance(f.get(key), (int, float))]
    return {'frames': len(frames), 'rows': out_rows,
            'evMean': float(np.mean(nums('ev'))) if nums('ev') else None,
            # pcom is the posterior of the currently expected G#2 event.
            # pcur is its successor, and is retained only as a separate value.
            'pcomMean': float(np.mean(nums('pcom'))) if nums('pcom') else None,
            'pcurMean': float(np.mean(nums('pcur'))) if nums('pcur') else None,
            'pMean': float(np.mean(nums('p'))) if nums('p') else None}


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('trace', type=pathlib.Path)
    ap.add_argument('--paired', type=pathlib.Path)
    ap.add_argument('--out', type=pathlib.Path)
    ap.add_argument('--export-residual', type=pathlib.Path,
                    help='write frozen-A-transfer residual chunks as f32le, in trial order')
    ns = ap.parse_args()
    trace = json.loads(ns.trace.read_text())
    paired = ns.paired or pathlib.Path(str(ns.trace) + '.paired.webm')
    audio, sr = decode(paired)
    events = trace['events']
    capture = next(e for e in events if e.get('type') == 'test.captureStart')
    trials = [e for e in events if e.get('type') == 'test.abTrial']
    if not trials or set(t['condition'] for t in trials) != {'A', 'B'}:
        raise SystemExit('This is not a complete A/B probe trace.')
    mic, ref = audio[:, 0].astype(np.float64), audio[:, 1].astype(np.float64)
    controls = [t for t in trials if t['condition'] == 'A']
    # Alternating controls train and hold out the transfer estimate. Treatment
    # never participates in fitting the speaker path.
    train, held = controls[::2], controls[1::2]
    if not held:
        train, held = controls[:-1], controls[-1:]
    delay = best_delay(ref, mic, trials, capture['t'], sr)
    h, fft_n, n, used, ridge = fit_transfer(ref, mic, train, capture['t'], sr)
    per_trial = []
    residual_chunks = []
    groups = {'A': {str(f): [] for f in HARMONICS}, 'B': {str(f): [] for f in HARMONICS}}
    for tr in trials:
        r = residual_for(ref, mic, tr, h, fft_n, n, capture['t'], sr)
        if r is None:
            continue
        x, y, residual = r
        if ns.export_residual:
            residual_chunks.append(residual.astype('<f4', copy=False))
        # The solo arrives at the physical mic after the measured loop delay.
        onset = delay['lagMs'] + 250
        lo, hi = round((onset + 35) * sr / 1000), round((onset + 360) * sr / 1000)
        section = residual[lo:hi]
        bands = {str(f): band_rms(section, sr, f) for f in HARMONICS}
        for f, val in bands.items():
            groups[tr['condition']][f].append(val)
        arrival_abs = tr['soloAtMs'] + delay['lagMs']
        per_trial.append({'trial': tr['trial'], 'condition': tr['condition'], 'atMs': tr['atMs'],
                          'soloAtMs': tr['soloAtMs'], 'arrivalEstimateMs': arrival_abs,
                          'micRmsDbfs': db(rms(y)), 'referenceRmsDbfs': db(rms(x)),
                          'residualRmsDbfs': db(rms(residual)),
                          'residualBandsDbfs': {f: db(v) for f, v in bands.items()},
                          'salience': salience_for(events, tr, arrival_abs)})
    harmonic = {}
    for f in map(str, HARMONICS):
        a, b = summary(groups['A'][f]), summary(groups['B'][f])
        # Difference of mean power, expressed as an RMS-equivalent dB increase.
        delta = b['meanDbfs'] - a['meanDbfs'] if a['meanDbfs'] is not None and b['meanDbfs'] is not None else None
        harmonic[f] = {'control': a, 'treatment': b, 'treatmentMinusControlDb': delta}
    result = {
        'trace': str(ns.trace), 'paired': str(paired), 'sampleRate': sr,
        'channels': {'left': 'physical microphone', 'right': 'digital accompaniment reference'},
        'transfer': {'method': 'complex frequency-domain least-squares fit on alternating accompaniment-only trials',
                     'trainControlTrials': [t['trial'] for t in train], 'heldControlTrials': [t['trial'] for t in held],
                     'segmentSeconds': n / sr, 'fftSize': fft_n, 'ridge': ridge, 'used': used},
        'alignment': delay,
        'residualWindow': {'relativeToSoloMicArrivalMs': [35, 360], 'bandHalfWidthHz': 5.0},
        'harmonics': harmonic, 'trials': per_trial,
    }
    if ns.export_residual:
        if len(residual_chunks) != len(per_trial):
            raise RuntimeError('Residual export lost a trial')
        ns.export_residual.write_bytes(b''.join(x.tobytes() for x in residual_chunks))
        result['residualExport'] = {'path': str(ns.export_residual), 'format': 'f32le',
                                    'samplesPerTrial': n,
                                    'trialOrder': [x['trial'] for x in per_trial]}
    destination = ns.out or pathlib.Path(str(ns.trace) + '.analysis.json')
    destination.write_text(json.dumps(result, indent=2) + '\n')
    print(json.dumps({'analysis': str(destination), 'alignment': delay, 'harmonics': harmonic}, indent=2))


if __name__ == '__main__':
    main()