tests/test_phase1_audio.py

"""Phase 1: WAV I/O, streams, JSON spatial reports."""

from __future__ import annotations

import json
import sys
import tempfile
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from hoa64.analysis import angular_error_deg
from hoa64.audio_io import read_wav, write_wav
from hoa64.report import (
    REPORT_SCHEMA_VERSION,
    report_from_ambix_wav,
    report_from_mono_wav,
    report_from_scene,
)
from hoa64.stream import SourceSpec, encode_mono_plane_wave, encode_scene
from hoa64.synth import envelope_adsr, tone


def test_wav_roundtrip(tmp_path: Path | None = None):
    base = Path(tmp_path) if tmp_path else Path(tempfile.mkdtemp())
    sr = 16000
    x = tone(440, 0.2, sr, amplitude=0.3)
    path = base / "mono.wav"
    write_wav(path, x, sr)
    audio, sr2 = read_wav(path)
    assert sr2 == sr
    assert audio.shape[0] == 1
    assert audio.shape[1] == x.shape[0]
    assert np.corrcoef(audio[0], x)[0, 1] > 0.99


def test_mono_encode_report_doa():
    sr = 24000
    sig = tone(1000, 0.4, sr, amplitude=0.5) * envelope_adsr(
        int(0.4 * sr), sr
    )
    hoa = encode_mono_plane_wave(sig, 45.0, 0.0, max_order=3)
    rep = report_from_hoa_local(hoa, sr, az_true=45.0)
    err = angular_error_deg(45.0, 0.0, rep.doa_az_deg, rep.doa_el_deg)
    assert err < 5.0, f"DOA err={err} got ({rep.doa_az_deg},{rep.doa_el_deg})"
    assert rep.schema == REPORT_SCHEMA_VERSION
    assert rep.duration_sec > 0.3
    d = rep.to_dict()
    assert "bands" in d and "frames" in d
    # JSON serializable
    json.dumps(d)


def report_from_hoa_local(hoa, sr, az_true=None):
    from hoa64.report import report_from_hoa

    return report_from_hoa(hoa, sr, max_order=3, include_peak_map=True)


def test_scene_two_sources_json(tmp_path: Path | None = None):
    base = Path(tmp_path) if tmp_path else Path(tempfile.mkdtemp())
    sr = 16000
    n = int(0.35 * sr)
    env = envelope_adsr(n, sr)
    sources = [
        SourceSpec(-30.0, 0.0, tone(500, 0.35, sr, amplitude=0.5) * env, "A"),
        SourceSpec(120.0, 10.0, tone(900, 0.35, sr, amplitude=0.35) * env, "B"),
    ]
    rep = report_from_scene(sources, sr, max_order=3)
    assert len(rep.sources_hint) == 2
    assert rep.energy > 0
    out = base / "report.json"
    rep.save(out)
    loaded = json.loads(out.read_text())
    assert loaded["schema"] == REPORT_SCHEMA_VERSION
    assert loaded["n_channels"] == 16  # order 3
    assert "one_liner" not in loaded
    assert rep.one_liner().startswith("spatial:")


def test_ambix_wav_analyze(tmp_path: Path | None = None):
    base = Path(tmp_path) if tmp_path else Path(tempfile.mkdtemp())
    sr = 16000
    sig = tone(700, 0.25, sr, amplitude=0.4)
    hoa = encode_mono_plane_wave(sig, -60.0, 5.0, max_order=1)
    path = base / "bformat.wav"
    write_wav(path, hoa[:4], sr)
    rep = report_from_ambix_wav(path, max_order=1, include_bands=True)
    err = angular_error_deg(-60.0, 5.0, rep.doa_az_deg, rep.doa_el_deg)
    assert err < 8.0, f"ambix DOA err={err}"


def test_mono_wav_cli_path(tmp_path: Path | None = None):
    base = Path(tmp_path) if tmp_path else Path(tempfile.mkdtemp())
    sr = 16000
    sig = tone(800, 0.2, sr, amplitude=0.4)
    wav = base / "m.wav"
    write_wav(wav, sig, sr)
    rep = report_from_mono_wav(wav, 0.0, 0.0, max_order=1)
    err = angular_error_deg(0.0, 0.0, rep.doa_az_deg, rep.doa_el_deg)
    assert err < 5.0


if __name__ == "__main__":
    test_wav_roundtrip()
    print("OK test_wav_roundtrip")
    test_mono_encode_report_doa()
    print("OK test_mono_encode_report_doa")
    test_scene_two_sources_json()
    print("OK test_scene_two_sources_json")
    test_ambix_wav_analyze()
    print("OK test_ambix_wav_analyze")
    test_mono_wav_cli_path()
    print("OK test_mono_wav_cli_path")