#!/usr/bin/env python3
# Part of JoeyLib, Copyright (c) 2026 Scott Duensing, under the MIT License.
# See LICENSE.md for license information.
# SPDX-License-Identifier: MIT
"""Grade a recording of AUDPAR (examples/audpar) against the audio contract.

AUDPAR plays one timeline and logs when each step started
("AUDPAR <ms> <event> <values>"). This finds the timeline in the recording
by its two sync tones (1 kHz, at the start and the end, so an emulator whose
clock runs a little fast or slow is mapped straight), cuts each step out of
it, and measures it:

    tone, voice   the fundamental within PITCH_TOLERANCE of the asked Hz
    voice levels  voice 0's attenuation steps 2 dB apart, LEVEL_TOLERANCE
    voices 1, 2   their pitch, and as loud as voice 0 (VOICE_MATCH_DB)
    noise         the generator's rate, 125000 / pitch Hz: read off the
                  first spectral null where every machine's output passes
                  it (NOISE_GRADED_HZ, within NOISE_TOLERANCE), darker as
                  the pitch rises everywhere, and its steps 2 dB apart
    sfx           a sine of 8 samples a cycle sounds rate / 8 Hz, starting
                  within SFX_LATENCY_MAX of the call; panned left or right
                  it is in one channel (--stereo ports)
    mod           a module that does not loop sounds, then ends, and
                  jlAudioIsPlayingMod goes false within MOD_END_SLACK of
                  its sound stopping
    tick          the callbacks that came, against hz * time
    silence       every gap is quiet: nothing left sounding

Usage:
    gradeAudioConformance.py PORT LOG WAV [--stereo] [--skip NAME ...]

    --stereo  the port has stereo output: grade the pans
    --skip    a check this port cannot pass, by name prefix (a documented
              hardware limit), reported as SKIP
    --level-tolerance DB
              how far an attenuation step may land from 2 dB a step: a
              chip whose own levels are coarser (the YM2149's) gets its
              own figure

Prints one line per check, then "AUDPAR PORT passed/graded". Exit 0 when
every graded check passes, 1 when one fails, 2 on a usage or input error.
"""

import argparse
import math
import re
import sys
import wave

import numpy as np
from scipy import signal

SYNC_HZ = 1000.0
SYNC_BAND_HZ = 60.0
FRAME_S = 0.010
WINDOW_S = 0.040
SYNC_SHARE = 0.5               # band power share that marks a sync frame
SYNC_MIN_DBFS = -45.0
SYNC_RUN_FRAMES = 10           # consecutive frames: 100 ms of sync
MARGIN_S = 0.060               # trimmed from each end of a step
PITCH_LOW_HZ = 40.0
PITCH_HIGH_HZ = 8000.0
PITCH_TOLERANCE = 0.01
LEVEL_STEP_DB = 2.0
LEVEL_TOLERANCE = 1.0
NOISE_LEVEL_TOLERANCE = 1.5
VOICE_MATCH_DB = 3.0
TONE_MATCH_DB = 4.0
NOISE_RATE_HZ = 125000.0
NOISE_TOLERANCE = 0.15
NOISE_FIT_LIMIT = 0.45         # of the sample rate: the band a null is sought in
NULL_SEARCH_LOW = 0.6          # of the rate asked: where its first null is sought
NULL_SEARCH_HIGH = 1.6
NULL_LEVEL_LOW = 0.2           # the level a null is measured down from
NULL_DEPTH_DB = 6.0
NULL_SMOOTH = 5                # bins of the spectrum averaged
NOISE_GRADED_HZ = 4100.0       # rates up to this are read off the spectrum:
                               # above it the machines' own output filters
                               # (the Amiga's, the ST's) bury the null
SILENCE_DBFS = -55.0
SOUND_DBFS = -45.0
PAN_APART_DB = 12.0
PAN_CENTRE_DB = 3.0
CENTRE_CHECK_HZ = 440.0        # the tone whose stereo position is graded
NOISE_LOUDNESS_PITCH = 31      # the noise compared with voice 0 for loudness
MOD_END_SLACK_S = 0.6
SFX_LATENCY_MAX = 0.050        # a sound effect starts this soon after the call
ONSET_SEARCH_S = 1.0
ONSET_BAND = 0.08              # of the pitch: the band an onset is heard in
ONSET_SHARE = 0.5
MOD_LEAD_S = 0.1
TICK_SHARE = 0.03
FULL_SCALE = 32768.0
LEVEL_FRAME_S = 0.020
SFX_SAMPLES = 4000             # AUDPAR's sine: its length at each rate


class Grader:
    def __init__(self, port, skips):
        self.port = port
        self.skips = skips
        self.passed = 0
        self.graded = 0
        self.failed = 0

    def check(self, name, ok, detail):
        if any(name.startswith(s) for s in self.skips):
            print(f"SKIP {name}: {detail}")
            return
        self.graded += 1
        if ok:
            self.passed += 1
        else:
            self.failed += 1
        print(f"{'PASS' if ok else 'FAIL'} {name}: {detail}")


def dbfs(x):
    """Level without DC: a DAC resting off centre (the ST's YM) is silence."""
    rms = float(np.sqrt(np.mean(np.square(x - x.mean())))) if len(x) else 0.0
    return 20.0 * math.log10(max(rms, 1e-9) / FULL_SCALE)


def steadyDbfs(x, rate):
    """The median level of LEVEL_FRAME_S frames: a step's own level, which
    a glitch of the recording path (a real-time emulator's sound server
    dropping or repeating a block) cannot move the way it moves a mean."""
    frame = int(LEVEL_FRAME_S * rate)
    if len(x) < frame:
        return dbfs(x)
    levels = [dbfs(x[i:i + frame]) for i in range(0, len(x) - frame + 1, frame)]
    return float(np.median(levels))


def loadWav(path):
    with wave.open(path, 'rb') as w:
        channels = w.getnchannels()
        width = w.getsampwidth()
        rate = w.getframerate()
        raw = w.readframes(w.getnframes())
    if width != 2:
        raise ValueError(f"{path}: expected 16-bit samples, got {width * 8}-bit")
    data = np.frombuffer(raw, dtype='<i2').astype(np.float64)
    data = data[:(len(data) // channels) * channels].reshape(-1, channels)
    return data, rate


def parseLog(path):
    events = []
    with open(path, errors='replace') as f:
        for line in f:
            m = re.search(r'AUDPAR (\d+) (\w+)(.*)', line)
            if m:
                events.append((int(m.group(1)) / 1000.0, m.group(2), m.group(3).split()))
    return events


def syncFrames(mono, rate):
    """Per FRAME_S hop: True where the 1 kHz sync tone dominates."""
    window = int(WINDOW_S * rate)
    hop = int(FRAME_S * rate)
    freqs = np.fft.rfftfreq(window, 1.0 / rate)
    band = (freqs > SYNC_HZ - SYNC_BAND_HZ) & (freqs < SYNC_HZ + SYNC_BAND_HZ)
    taper = np.hanning(window)
    flags = []
    for start in range(0, len(mono) - window, hop):
        frame = mono[start:start + window]
        power = np.abs(np.fft.rfft(frame * taper)) ** 2
        total = power.sum()
        loud = dbfs(frame) > SYNC_MIN_DBFS
        flags.append(bool(loud and total > 0 and power[band].sum() / total > SYNC_SHARE))
    return flags


def syncOnset(flags, fromFrame):
    """First frame at or after fromFrame that starts SYNC_RUN_FRAMES of sync."""
    run = 0
    for i in range(fromFrame, len(flags)):
        run = run + 1 if flags[i] else 0
        if run == SYNC_RUN_FRAMES:
            return i - SYNC_RUN_FRAMES + 1
    return None


def pitchOf(x, rate):
    """Strongest frequency in PITCH_LOW_HZ..PITCH_HIGH_HZ, refined between bins."""
    if len(x) < 256:
        return 0.0
    size = 1 << (int(math.log2(len(x))) + 2)
    spectrum = np.abs(np.fft.rfft((x - x.mean()) * np.hanning(len(x)), size))
    freqs = np.fft.rfftfreq(size, 1.0 / rate)
    allowed = np.where((freqs >= PITCH_LOW_HZ) & (freqs <= PITCH_HIGH_HZ))[0]
    peak = allowed[np.argmax(spectrum[allowed])]
    if 0 < peak < len(spectrum) - 1:
        a, b, c = np.log(spectrum[peak - 1:peak + 2] + 1e-12)
        offset = 0.5 * (a - c) / (a - 2.0 * b + c) if (a - 2.0 * b + c) != 0 else 0.0
    else:
        offset = 0.0
    return (peak + offset) * rate / size


def noiseRate(x, rate, expected):
    """The noise's sample-and-hold rate, read off its first spectral null:
    noise held at R Hz has a sinc^2 spectrum that first falls to nothing at
    R. Looked for between NULL_SEARCH_LOW and NULL_SEARCH_HIGH of the rate
    asked, and only a real dip counts -- NULL_DEPTH_DB below the spectrum's
    level under it -- so a wrong rate shows as a null in the wrong place or
    as none. (Fitting the whole spectrum was thrown by the images a chip's
    own output rate leaves above it: the IIgs's DOC runs at 26.3 kHz.)
    Returns None when there is no such null."""
    nper = min(2048, len(x))
    freqs, psd = signal.welch(x - x.mean(), rate, nperseg=nper)
    level = 10.0 * np.log10(np.convolve(psd, np.ones(NULL_SMOOTH) / NULL_SMOOTH, mode='same') + 1e-12)
    search = np.where((freqs >= NULL_SEARCH_LOW * expected) & (freqs <= min(NULL_SEARCH_HIGH * expected, NOISE_FIT_LIMIT * rate)))[0]
    below = np.where((freqs >= NULL_LEVEL_LOW * expected) & (freqs <= NULL_SEARCH_LOW * expected))[0]
    if len(search) < 3 or len(below) == 0:
        return None
    dip = search[np.argmin(level[search])]
    if dip in (search[0], search[-1]) or level[below].max() - level[dip] < NULL_DEPTH_DB:
        return None
    return float(freqs[dip])


def centroid(x, rate):
    freqs, psd = signal.welch(x - x.mean(), rate, nperseg=min(2048, len(x)))
    return float((freqs * psd).sum() / max(psd.sum(), 1e-12))


def main():
    parser = argparse.ArgumentParser(add_help=True)
    parser.add_argument('port')
    parser.add_argument('log')
    parser.add_argument('wav')
    parser.add_argument('--stereo', action='store_true')
    parser.add_argument('--skip', nargs='*', default=[])
    parser.add_argument('--level-tolerance', type=float, default=LEVEL_TOLERANCE)
    args = parser.parse_args()

    try:
        data, rate = loadWav(args.wav)
        events = parseLog(args.log)
    except (OSError, ValueError, wave.Error) as error:
        print(f"AUDPAR {args.port}: cannot read the recording or the log: {error}")
        return 2
    syncs = [e for e in events if e[1] == 'sync']
    done = [e for e in events if e[1] == 'done']
    if len(syncs) != 2 or not done:
        print(f"AUDPAR {args.port}: the log does not hold a whole run (syncs {len(syncs)}, done {len(done)})")
        return 2
    mono = data.mean(axis=1)
    grader = Grader(args.port, args.skip)

    # Align: the two sync onsets map log seconds onto the recording.
    flags = syncFrames(mono, rate)
    first = syncOnset(flags, 0)
    span = syncs[1][0] - syncs[0][0]
    last = None if first is None else syncOnset(flags, first + int(0.8 * span / FRAME_S))
    if first is None or last is None:
        grader.check('sync', False, f"sync tones found: {'start' if first is not None else 'none'}{', end' if last is not None else ''}")
        print(f"AUDPAR {args.port} {grader.passed}/{grader.graded}")
        return 1
    t0 = first * FRAME_S
    scale = (last - first) * FRAME_S / span
    grader.check('sync', abs(scale - 1.0) < 0.05, f"start {t0:.3f} s, the timeline runs at {scale:.4f} of the program's clock")

    def cut(start, end, channel=None, shift=0.0):
        a = int((t0 + start * scale + shift + MARGIN_S) * rate)
        b = int((t0 + end * scale + shift - MARGIN_S) * rate)
        if b <= a:
            return np.zeros(0)
        return mono[a:b] if channel is None else data[a:b, channel]

    def onset(start, hz):
        """Seconds from a step's log time to a sound at hz, or None: the tail
        of whatever came before is a different pitch and does not count."""
        window = int(WINDOW_S * rate)
        hop = int(FRAME_S * rate)
        freqs = np.fft.rfftfreq(window, 1.0 / rate)
        band = (freqs > hz * (1.0 - ONSET_BAND)) & (freqs < hz * (1.0 + ONSET_BAND))
        taper = np.hanning(window)
        a = int((t0 + start * scale) * rate)
        for i in range(a, min(len(mono), a + int(ONSET_SEARCH_S * rate)) - window, hop):
            frame = mono[i:i + window]
            power = np.abs(np.fft.rfft((frame - frame.mean()) * taper)) ** 2
            if dbfs(frame) > SOUND_DBFS and power[band].sum() > ONSET_SHARE * power.sum():
                return (i - a) / rate
        return None

    steps = []
    for i, event in enumerate(events):
        nextStart = events[i + 1][0] if i + 1 < len(events) else event[0]
        steps.append((event[0], nextStart, event[1], event[2]))

    def centred(start, end, name):
        left = dbfs(cut(start, end, 0))
        right = dbfs(cut(start, end, 1))
        grader.check(f"{name} centred", abs(left - right) <= PAN_CENTRE_DB, f"left {left:.1f} dBFS, right {right:.1f} dBFS")

    lag = 0.0                  # the sample path's latency, from the last sfx
    toneLevels = []
    voiceLevels = {}
    noiseLevels = {}
    noiseCentroids = []
    for start, end, kind, values in steps:
        if kind in ('tone', 'sync'):
            hz = float(values[0])
            x = cut(start, end)
            got = pitchOf(x, rate)
            level = dbfs(x)
            grader.check(f"{kind} {int(hz)} Hz", abs(got / hz - 1.0) <= PITCH_TOLERANCE and level > SOUND_DBFS, f"{got:.1f} Hz at {level:.1f} dBFS")
            if kind == 'tone':
                toneLevels.append((int(hz), level))
                if args.stereo and hz == CENTRE_CHECK_HZ:
                    centred(start, end, 'tone')
        elif kind == 'voice':
            voice, hz, atten = int(values[0]), float(values[1]), int(values[2])
            x = cut(start, end)
            level = steadyDbfs(x, rate)
            if args.stereo and atten == 0:
                centred(start, end, f"voice {voice}")
            voiceLevels[(voice, atten)] = level
            if atten <= 8:
                got = pitchOf(x, rate)
                grader.check(f"voice {voice} {int(hz)} Hz atten {atten}", abs(got / hz - 1.0) <= PITCH_TOLERANCE, f"{got:.1f} Hz at {level:.1f} dBFS")
        elif kind == 'noise':
            pitch, atten = int(values[0]), int(values[1])
            x = cut(start, end)
            level = steadyDbfs(x, rate)
            if atten == 0 and (pitch, 0) not in noiseLevels:
                expected = NOISE_RATE_HZ / max(pitch, 1)
                noiseCentroids.append((pitch, centroid(x, rate)))
                if level <= SOUND_DBFS:
                    grader.check(f"noise {pitch} sounds", False, f"{level:.1f} dBFS: it does not hold for the step")
                elif expected <= NOISE_GRADED_HZ:
                    got = noiseRate(x, rate, expected)
                    grader.check(f"noise {pitch} rate", got is not None and abs(got / expected - 1.0) <= NOISE_TOLERANCE, f"{got:.0f} Hz, {expected:.0f} asked, at {level:.1f} dBFS" if got is not None else f"no null near {expected:.0f} Hz, at {level:.1f} dBFS")
                else:
                    grader.check(f"noise {pitch} sounds", True, f"{level:.1f} dBFS (its rate, {expected:.0f} Hz, is graded by the darkening below)")
            noiseLevels[(pitch, atten)] = level
        elif kind == 'sfx':
            hz, pan = float(values[2]), int(values[3])
            length = SFX_SAMPLES / float(values[1])
            panName = {0: 'left', 1: 'centre', 2: 'right'}.get(pan, str(pan))
            found = onset(start, hz)
            grader.check(f"sfx {values[1]} Hz {panName} starts", found is not None and found <= SFX_LATENCY_MAX, f"{found * 1000:.0f} ms after the call" if found is not None else "never")
            lag = found if found is not None else 0.0
            x = cut(start, start + length, shift=lag)
            got = pitchOf(x, rate)
            level = dbfs(x)
            grader.check(f"sfx {values[1]} Hz {panName}", abs(got / hz - 1.0) <= PITCH_TOLERANCE and level > SOUND_DBFS, f"{got:.1f} Hz ({hz:.0f} asked) at {level:.1f} dBFS")
            if data.shape[1] >= 2:
                left = dbfs(cut(start, start + length, 0, lag))
                right = dbfs(cut(start, start + length, 1, lag))
                apart = left - right
                if pan == 0:
                    ok = apart >= PAN_APART_DB
                elif pan == 2:
                    ok = -apart >= PAN_APART_DB
                else:
                    ok = abs(apart) <= PAN_CENTRE_DB
                if args.stereo:
                    grader.check(f"pan {panName}", ok, f"left {left:.1f} dBFS, right {right:.1f} dBFS")
        elif kind == 'silence':
            # After a sample-path sound, its tail is late by the latency
            # the sfx check graded: the silence is judged after it.
            x = cut(start, end, shift=lag)
            level = steadyDbfs(x, rate)
            grader.check(f"silence at {start:.2f} s", level < SILENCE_DBFS, f"{level:.1f} dBFS")
        elif kind == 'mod' and values and values[0] == 'start':
            endEvents = [e for e in events if e[1] == 'mod' and e[2] and e[2][0] == 'end']
            if not endEvents:
                grader.check('mod', False, "no mod end in the log")
                continue
            endAt, endValues = endEvents[0][0], endEvents[0][2]
            stillPlaying = len(endValues) > 1 and endValues[1] != '0'
            x = cut(start, endAt + 1.0)
            frames = int(FRAME_S * rate)
            loud = [dbfs(x[i:i + frames]) > SOUND_DBFS for i in range(0, len(x) - frames, frames)]
            sounded = any(loud)
            lastLoud = (max(i for i, v in enumerate(loud) if v) + 1) * FRAME_S if sounded else 0.0
            soundEnd = start + (MARGIN_S + lastLoud) / scale
            grader.check('mod sounds', sounded, f"{sum(loud) * FRAME_S:.1f} s of it above {SOUND_DBFS:.0f} dBFS")
            grader.check('mod ends', not stillPlaying and endAt - start < 11.5, f"jlAudioIsPlayingMod false {endAt - start:.2f} s after the start")
            grader.check('mod end reported', -MOD_LEAD_S <= endAt - soundEnd <= MOD_END_SLACK_S, f"its sound stopped at {soundEnd - start:.2f} s, reported at {endAt - start:.2f} s")
        elif kind == 'tick':
            hz, ok, calls, ms = int(values[0]), values[1] == '1', int(values[2]), int(values[3])
            expected = hz * ms / 1000.0
            good = ok and abs(calls - expected) <= max(1.0, TICK_SHARE * expected)
            grader.check(f"tick {hz} Hz", good, f"registered {'yes' if ok else 'no'}, {calls} calls in {ms} ms, {expected:.0f} expected")

    if toneLevels:
        base = toneLevels[0][1]
        spread = max(abs(level - base) for _, level in toneLevels)
        grader.check('tone levels even', spread <= TONE_MATCH_DB, ", ".join(f"{hz} Hz {level:.1f}" for hz, level in toneLevels))
    if (0, 0) in voiceLevels:
        base = voiceLevels[(0, 0)]
        for atten in range(1, 15):
            if (0, atten) in voiceLevels:
                drop = base - voiceLevels[(0, atten)]
                grader.check(f"voice level {atten}", abs(drop - LEVEL_STEP_DB * atten) <= args.level_tolerance, f"{drop:.1f} dB below attenuation 0, {LEVEL_STEP_DB * atten:.0f} asked")
        for voice in (1, 2):
            if (voice, 0) in voiceLevels:
                grader.check(f"voice {voice} as loud as voice 0", abs(voiceLevels[(voice, 0)] - base) <= VOICE_MATCH_DB, f"{voiceLevels[(voice, 0)]:.1f} against {base:.1f} dBFS")
        # Full-swing noise and a full-swing square carry the same power, as
        # on the PSG every port's voices are modelled on.
        if (NOISE_LOUDNESS_PITCH, 0) in noiseLevels:
            grader.check("noise as loud as a voice", abs(noiseLevels[(NOISE_LOUDNESS_PITCH, 0)] - base) <= VOICE_MATCH_DB, f"{noiseLevels[(NOISE_LOUDNESS_PITCH, 0)]:.1f} against {base:.1f} dBFS")
    stepPitch = None
    for (pitch, atten) in noiseLevels:
        if atten > 0:
            stepPitch = pitch
    if stepPitch is not None and (stepPitch, 0) in noiseLevels:
        base = noiseLevels[(stepPitch, 0)]
        for atten in sorted(a for p, a in noiseLevels if p == stepPitch and a > 0):
            drop = base - noiseLevels[(stepPitch, atten)]
            grader.check(f"noise level {atten}", abs(drop - LEVEL_STEP_DB * atten) <= max(NOISE_LEVEL_TOLERANCE, args.level_tolerance), f"{drop:.1f} dB below attenuation 0, {LEVEL_STEP_DB * atten:.0f} asked")
    # Each pitch darker than the one before, where the later one's rate is
    # inside the recording's band: rates past half the sample rate all
    # record as near-white noise and cannot be told apart.
    if len(noiseCentroids) > 1:
        darker = all(b[1] < a[1] for a, b in zip(noiseCentroids, noiseCentroids[1:]) if NOISE_RATE_HZ / max(b[0], 1) < rate / 2.0)
        grader.check('noise darkens', darker, ", ".join(f"{p}: {c:.0f} Hz" for p, c in noiseCentroids))

    print(f"AUDPAR {args.port} {grader.passed}/{grader.graded}")
    return 0 if grader.failed == 0 else 1


if __name__ == '__main__':
    sys.exit(main())