tests/test_phase2_wigner.py

"""Phase 2: Wigner-D HOA rotation accuracy and speed."""

from __future__ import annotations

import sys
import time
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, doa_from_intensity
from hoa64.encode import encode_points
from hoa64.rotate import rotate_matrix_order1, rotate_source_directions, rotate_yaw_pitch_roll
from hoa64.wigner import hoa_rotation_matrix, rotation_matrix_zyx, apply_hoa_rotation


def test_order1_wigner_matches_cartesian():
    a = encode_points([25.0], [-12.0], [1.0])
    for yaw, pitch, roll in [(90, 0, 0), (0, 45, 0), (30, -20, 15), (180, 0, 0)]:
        a_w = rotate_yaw_pitch_roll(
            a, yaw=yaw, pitch=pitch, roll=roll, max_order=1, method="wigner"
        )
        a_c = rotate_matrix_order1(a, yaw=yaw, pitch=pitch, roll=roll)
        np.testing.assert_allclose(a_w[:4], a_c[:4], atol=1e-9, rtol=1e-9)


def test_wigner_matches_plane_wave_reencode():
    """Ground truth: rotate source direction, re-encode."""
    cases = [
        (0.0, 0.0, 90.0, 0.0, 0.0),
        (40.0, -15.0, 35.0, 0.0, 0.0),
        (10.0, 20.0, 0.0, 40.0, 0.0),
        (-70.0, 5.0, 20.0, -25.0, 30.0),
    ]
    for az0, el0, yaw, pitch, roll in cases:
        a0 = encode_points([az0], [el0], [1.0], max_order=7)
        az1, el1 = rotate_source_directions(az0, el0, yaw=yaw, pitch=pitch, roll=roll)
        a_gt = encode_points([float(az1)], [float(el1)], [1.0], max_order=7)
        a_w = rotate_yaw_pitch_roll(
            a0, yaw=yaw, pitch=pitch, roll=roll, max_order=7, method="wigner"
        )
        # Relative error on full 64-vector
        denom = np.linalg.norm(a_gt) + 1e-15
        rel = np.linalg.norm(a_w[:64] - a_gt[:64]) / denom
        assert rel < 1e-6, (
            f"rel={rel:.3e} for src=({az0},{el0}) rot=({yaw},{pitch},{roll})"
        )


def test_wigner_orthogonal_blocks():
    R = rotation_matrix_zyx(33.0, -17.0, 8.0, degrees=True)
    M = hoa_rotation_matrix(R, max_order=7)
    # Each order block should be orthogonal (rotation)
    for n in range(0, 8):
        i0 = n * n
        dim = 2 * n + 1
        B = M[i0 : i0 + dim, i0 : i0 + dim]
        I = B.T @ B
        np.testing.assert_allclose(I, np.eye(dim), atol=1e-8)


def test_wigner_faster_than_dense():
    a = encode_points([15.0], [10.0], [1.0], max_order=7)
    # warm-up
    rotate_yaw_pitch_roll(a, yaw=20.0, pitch=10.0, roll=5.0, method="wigner")
    rotate_yaw_pitch_roll(
        a, yaw=20.0, pitch=10.0, roll=5.0, method="dense", n_azi=48, n_el=24
    )

    t0 = time.perf_counter()
    for _ in range(50):
        rotate_yaw_pitch_roll(a, yaw=20.0, pitch=10.0, roll=5.0, method="wigner")
    t_w = time.perf_counter() - t0

    t0 = time.perf_counter()
    for _ in range(5):
        rotate_yaw_pitch_roll(
            a, yaw=20.0, pitch=10.0, roll=5.0, method="dense", n_azi=48, n_el=24
        )
    t_d = time.perf_counter() - t0
    # per-call times
    tw = t_w / 50
    td = t_d / 5
    # Wigner should be substantially faster (typically 50–1000×)
    assert tw < td, f"wigner {tw:.4f}s not faster than dense {td:.4f}s"
    print(f"  timing: wigner={tw*1e3:.3f} ms/call  dense={td*1e3:.3f} ms/call  speedup={td/tw:.0f}x")


def test_stream_rotation_CT():
    # (C,T) path
    from hoa64.encode import encode_plane_waves

    t = np.linspace(0, 1, 32, endpoint=False)
    sig = np.sin(2 * np.pi * 3 * t)[None, :]
    hoa = encode_plane_waves([0.0], [0.0], sig, max_order=3)
    out = rotate_yaw_pitch_roll(hoa, yaw=90.0, max_order=3, method="wigner")
    assert out.shape[0] >= 16
    # DOA of first frame energy via products
    W, Y, Z, X = out[0], out[1], out[2], out[3]
    I = np.array([np.mean(W * X), np.mean(W * Y), np.mean(W * Z)])
    n = np.linalg.norm(I)
    assert n > 1e-9


if __name__ == "__main__":
    test_order1_wigner_matches_cartesian()
    print("OK test_order1_wigner_matches_cartesian")
    test_wigner_matches_plane_wave_reencode()
    print("OK test_wigner_matches_plane_wave_reencode")
    test_wigner_orthogonal_blocks()
    print("OK test_wigner_orthogonal_blocks")
    test_wigner_faster_than_dense()
    print("OK test_wigner_faster_than_dense")
    test_stream_rotation_CT()
    print("OK test_stream_rotation_CT")