#!/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

"""pspAudioGrade.py - grade what AUDIOPRB played on the PSP.

AUDIOPRB is examples/audio built to press its own keys (make/examples.mk):
the tracker song from the start, then S every PROBE_STREAM_FRAMES frames
from PROBE_STREAM_FIRST (and N once, whose noise sweep the streams cover).
Each S streams a falling square wave, each take STREAM_TAKE_HZ higher than
the last. scripts/verify-psp-audio.sh records what reached the PSP's audio
output (PPSSPP's own recording, 44.1 kHz stereo on the emulated clock), and
this grades it from the moment the music starts:

    music   the tracker song alone, before the first take: sound, and
            quieter than a take
    take N  each of the first three takes, a little way in: loud, and at
            its pitch, from the zero crossings -- so the PSP plays the mix
            at the rate it was made for
    end     the recording stops within END_SLACK seconds of the program
            being told to quit: it ran until then, and ended on it

    pspAudioGrade.py LABEL WAV QUIT_SECONDS

QUIT_SECONDS is when (since the recording began) the gate pressed Select,
which the PSP port gives as Escape. Prints one line per check and PASS or
FAIL; exit 1 on a failure.
"""

import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "tools"))
from gradeAudioPhases import readMono, rms  # noqa: E402

# examples/audio/audio.c: the probe's timetable and the streamed tone.
FRAME_HZ = 60
PROBE_STREAM_FIRST = 30
PROBE_STREAM_FRAMES = 90
STREAM_RATE_HZ = 8000
STREAM_BLOCK = 256
STREAM_BLOCKS = 48
STREAM_START_HZ = 330
STREAM_DROP_HZ = 220
STREAM_TAKE_HZ = 55
TAKES = 3

# The music starts when its first sample rises above ONSET_LEVEL.
ONSET_LEVEL = 200
MUSIC_WINDOW = (0.05, 0.40)
MUSIC_RMS = (1000, 8000)
# A take is measured from TAKE_WINDOW seconds after it is due to start: far
# enough in to miss the audio queue's latency, early in its fall.
TAKE_WINDOW = (0.10, 0.20)
TAKE_MIN_RMS = 15000
PITCH_TOLERANCE = 0.10
END_SLACK = 0.5


def crossingsPerSecond(samples, rate, start, end):
    a = int(start * rate)
    b = min(int(end * rate), len(samples))
    count = 0
    for i in range(a + 1, b):
        if (samples[i - 1] < 0) != (samples[i] < 0):
            count += 1
    return count / ((b - a) / rate)


def takeHz(take, seconds):
    """The pitch take is at, seconds after it starts."""
    block = int(seconds * STREAM_RATE_HZ) // STREAM_BLOCK
    return STREAM_START_HZ + STREAM_TAKE_HZ * take - STREAM_DROP_HZ * block / STREAM_BLOCKS


def main(argv):
    if len(argv) != 4:
        print(__doc__, file=sys.stderr)
        return 2
    label, path, quitSeconds = argv[1], argv[2], float(argv[3])
    samples, rate = readMono(path)
    duration = len(samples) / rate
    onset = next((i / rate for i, v in enumerate(samples) if abs(v) > ONSET_LEVEL), None)
    if onset is None:
        print(f"{label}: FAIL - {duration:.1f}s recorded and all of it silent", file=sys.stderr)
        return 1
    fails = []

    level = rms(samples, rate, onset + MUSIC_WINDOW[0], onset + MUSIC_WINDOW[1])
    print(f"{label}: music from {onset:.2f}s, RMS {level:.0f}")
    if level is None or not MUSIC_RMS[0] <= level <= MUSIC_RMS[1]:
        fails.append(f"music RMS {level}, want {MUSIC_RMS[0]}-{MUSIC_RMS[1]}")

    for take in range(TAKES):
        due = onset + (PROBE_STREAM_FIRST + PROBE_STREAM_FRAMES * take) / FRAME_HZ
        start = due + TAKE_WINDOW[0]
        end = due + TAKE_WINDOW[1]
        level = rms(samples, rate, start, end)
        if level is None:
            fails.append(f"take {take + 1} is past the end of the recording")
            continue
        hz = crossingsPerSecond(samples, rate, start, end) / 2
        want = takeHz(take, (TAKE_WINDOW[0] + TAKE_WINDOW[1]) / 2)
        print(f"{label}: take {take + 1} at {start:.2f}s: RMS {level:.0f}, {hz:.0f} Hz (want {want:.0f})")
        if level < TAKE_MIN_RMS:
            fails.append(f"take {take + 1} RMS {level:.0f}, want at least {TAKE_MIN_RMS}")
        if abs(hz - want) > want * PITCH_TOLERANCE:
            fails.append(f"take {take + 1} at {hz:.0f} Hz, want {want:.0f}")

    print(f"{label}: the recording ends at {duration:.2f}s, Select at {quitSeconds:.2f}s")
    if abs(duration - quitSeconds) > END_SLACK:
        fails.append(f"the recording ends at {duration:.2f}s, not within {END_SLACK}s of Select at {quitSeconds:.2f}s")

    for f in fails:
        print(f"{label}: FAIL - {f}", file=sys.stderr)
    if fails:
        return 1
    print(f"{label}: PASS")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))