tests/test_rotate_rnn.py

"""Rotation + iterative RNN-stub motion tests."""

from __future__ import annotations

import sys
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, peak_direction
from hoa64.encode import encode_points
from hoa64.rotate import (
    rotate_matrix_order1,
    rotate_source_directions,
    rotate_yaw_pitch_roll,
)
from hoa64.rnn_stub import step_rotate, world_from_sources


def test_order1_yaw_matches_reencode():
    a = encode_points([0.0], [0.0], [1.0])
    # rotate field +90° yaw: front → left
    a_fast = rotate_matrix_order1(a, yaw=90.0)
    az, el = doa_from_intensity(a_fast)
    err = angular_error_deg(90.0, 0.0, az, el)
    assert err < 1.0, f"intensity DOA after yaw90: ({az},{el}) err={err}"


def test_source_dir_rotation_exact_reencode():
    """Ground truth: rotate source coordinates, re-encode."""
    az0, el0 = 20.0, -10.0
    a0 = encode_points([az0], [el0], [1.0])
    az1, el1 = rotate_source_directions(az0, el0, yaw=45.0)
    a1 = encode_points([float(az1)], [float(el1)], [1.0])
    # dense rotate of a0 by same yaw should ≈ a1 for low orders
    a_rot = rotate_yaw_pitch_roll(a0, yaw=45.0, max_order=3, n_azi=120, n_el=60)
    # compare order ≤ 3 channels
    nch = 16
    rel = np.linalg.norm(a_rot[:nch] - a1[:nch]) / (np.linalg.norm(a1[:nch]) + 1e-12)
    assert rel < 0.15, f"relative L2 order≤3 after rotate: {rel}"


def test_rnn_loop_keeps_tracking_after_turns():
    """Agent turns through 360° in steps; peak should stay world-stable in intensity
    when we interpret DOA in head frame correctly.

    After total yaw +90 (agent turns left), a world-front source is at head-right
    (−90) in listener frame if field is counter-rotated.
    """
    st = world_from_sources([0.0], [0.0], [1.0])
    # 6 steps of +15° agent yaw
    for _ in range(6):
        st = step_rotate(st, d_yaw=15.0, max_order=3, dense=True)
    assert abs(st.yaw - 90.0) < 1e-9
    az, el = doa_from_intensity(st.hoa)
    # front source, agent yawed +90 → source appears at az=-90 (right)
    err = angular_error_deg(-90.0, 0.0, az, el)
    assert err < 8.0, f"after +90 agent yaw, head-frame DOA=({az},{el}) err={err}"


def test_iterative_path_energy_stable_order1():
    st = world_from_sources([45.0], [0.0], [1.0])
    e0 = float(np.dot(st.hoa[:4], st.hoa[:4]))
    for _ in range(24):
        st = step_rotate(st, d_yaw=15.0, max_order=1, dense=False)
    e1 = float(np.dot(st.hoa[:4], st.hoa[:4]))
    # full 360° of order-1 rotations should preserve energy exactly
    assert abs(e1 - e0) / e0 < 1e-9


if __name__ == "__main__":
    for fn in [
        test_order1_yaw_matches_reencode,
        test_source_dir_rotation_exact_reencode,
        test_rnn_loop_keeps_tracking_after_turns,
        test_iterative_path_energy_stable_order1,
    ]:
        fn()
        print("OK", fn.__name__)