# ------------------------------------------------------------
# AAC Coder/Decoder - Filterbank Tests
#
# Multimedia course at Aristotle University of
# Thessaloniki (AUTh)
#
# Author:
#   Christos Choutouridis (ΑΕΜ 8997)
#   cchoutou@ece.auth.gr
#
# Description:
#   Tests for Filterbank module.
# ------------------------------------------------------------
from __future__ import annotations

from typing import Sequence
import pytest

from core.aac_filterbank import aac_filter_bank, aac_i_filter_bank
from core.aac_utils import snr_db
from core.aac_types import *

# Helper fixtures for filterbank
# -----------------------------------------------------------------------------

def _ola_reconstruct(x: StereoSignal, frame_types: Sequence[FrameType], win_type: WinType) -> StereoSignal:
    """
    Analyze-synthesize each frame and overlap-add with hop=1024.

    Parameters
    ----------
    x : StereoSignal
        Input stereo stream, expected shape (N, 2).
    frame_types : Sequence[FrameType]
        Length K sequence of frame types for frames starting at i*1024.
    win_type : WinType
        Window type ("SIN" or "KBD").

    Returns
    -------
    StereoSignal
        Reconstructed stereo stream, same shape as x (N, 2).
    """
    hop = 1024
    win = 2048
    K = len(frame_types)

    y: StereoSignal = np.zeros_like(x, dtype=np.float64)

    for i in range(K):
        start = i * hop
        frame_t: FrameT = x[start:start + win, :]
        frame_f: FrameF = aac_filter_bank(frame_t, frame_types[i], win_type)
        frame_t_hat: FrameT = aac_i_filter_bank(frame_f, frame_types[i], win_type)
        y[start:start + win, :] += frame_t_hat

    return y


# -----------------------------------------------------------------------------
# Forward filterbank tests
# -----------------------------------------------------------------------------

@pytest.mark.parametrize("win_type", ["SIN", "KBD"])
@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_filterbank_shapes_long_sequences(frame_type: FrameType, win_type: WinType) -> None:
    """
    Contract test: for OLS/LSS/LPS, aac_filter_bank returns shape (1024, 2).
    """
    frame_t: FrameT = np.zeros((2048, 2), dtype=np.float64)
    frame_f = aac_filter_bank(frame_t, frame_type, win_type)
    assert frame_f.shape == (1024, 2)


@pytest.mark.parametrize("win_type", ["SIN", "KBD"])
def test_filterbank_shapes_esh(win_type: WinType) -> None:
    """
    Contract test: for ESH, aac_filter_bank returns shape (128, 16).
    """
    frame_t: FrameT = np.zeros((2048, 2), dtype=np.float64)
    frame_f = aac_filter_bank(frame_t, "ESH", win_type)
    assert frame_f.shape == (128, 16)


@pytest.mark.parametrize("win_type", ["SIN", "KBD"])
def test_filterbank_channel_isolation_long_sequences(win_type: WinType) -> None:
    """
    Behavior test: for OLS (representative long-sequence), channels are independent.
    If right channel is zero and left is random, right spectrum should be near zero.
    """
    rng = np.random.default_rng(0)
    frame_t: FrameT = np.zeros((2048, 2), dtype=np.float64)
    frame_t[:, 0] = rng.normal(size=2048)

    frame_f = aac_filter_bank(frame_t, "OLS", win_type)

    assert np.max(np.abs(frame_f[:, 1])) < 1e-9


@pytest.mark.parametrize("win_type", ["SIN", "KBD"])
def test_filterbank_channel_isolation_esh(win_type: WinType) -> None:
    """
    Behavior test: for ESH, channels are independent.
    If right channel is zero and left is random, all odd columns (right) should be near zero.
    """
    rng = np.random.default_rng(1)
    frame_t: FrameT = np.zeros((2048, 2), dtype=np.float64)
    frame_t[:, 0] = rng.normal(size=2048)

    frame_f = aac_filter_bank(frame_t, "ESH", win_type)

    right_cols = frame_f[:, 1::2]  # columns 1,3,5,...,15
    assert np.max(np.abs(right_cols)) < 1e-9


@pytest.mark.parametrize("win_type", ["SIN", "KBD"])
def test_filterbank_esh_ignores_outer_regions(win_type: WinType) -> None:
    """
    Spec-driven behavior test:
    ESH uses only the central region [448, 1600), split into 8 overlapping
    windows of length 256 with 50% overlap.

    Therefore, changing samples outside [448, 1600) must not affect the output.
    """
    rng = np.random.default_rng(2)

    frame_a: FrameT = np.zeros((2048, 2), dtype=np.float64)
    frame_b: FrameT = np.zeros((2048, 2), dtype=np.float64)

    center = rng.normal(size=(1152, 2))
    frame_a[448:1600, :] = center
    frame_b[448:1600, :] = center

    frame_b[0:448, :] = rng.normal(size=(448, 2))
    frame_b[1600:2048, :] = rng.normal(size=(448, 2))

    fa = aac_filter_bank(frame_a, "ESH", win_type)
    fb = aac_filter_bank(frame_b, "ESH", win_type)

    # Use a tiny tolerance to avoid flaky failures due to floating-point minutiae.
    np.testing.assert_allclose(fa, fb, rtol=0.0, atol=1e-12)


@pytest.mark.parametrize("win_type", ["SIN", "KBD"])
def test_filterbank_output_is_finite(win_type: WinType) -> None:
    """
    Sanity test: output must not contain NaN or inf for representative cases.
    """
    rng = np.random.default_rng(3)
    frame_t: FrameT = rng.normal(size=(2048, 2)).astype(np.float64)

    for frame_type in ("OLS", "LSS", "ESH", "LPS"):
        frame_f = aac_filter_bank(frame_t, frame_type, win_type)
        assert np.isfinite(frame_f).all()


# -----------------------------------------------------------------------------
# Reverse i_filterbank tests
# -----------------------------------------------------------------------------

@pytest.mark.parametrize("win_type", ["SIN", "KBD"])
def test_ifilterbank_shapes_long_sequences(win_type: WinType) -> None:
    """
    Contract test: for OLS/LSS/LPS, aac_i_filter_bank returns shape (2048, 2).
    """
    frame_f: FrameF = np.zeros((1024, 2), dtype=np.float64)
    for frame_type in ("OLS", "LSS", "LPS"):
        frame_t = aac_i_filter_bank(frame_f, frame_type, win_type)
        assert frame_t.shape == (2048, 2)


@pytest.mark.parametrize("win_type", ["SIN", "KBD"])
def test_ifilterbank_shapes_esh(win_type: WinType) -> None:
    """
    Contract test: for ESH, aac_i_filter_bank returns shape (2048, 2).
    """
    frame_f: FrameF = np.zeros((128, 16), dtype=np.float64)
    frame_t = aac_i_filter_bank(frame_f, "ESH", win_type)
    assert frame_t.shape == (2048, 2)


@pytest.mark.parametrize("win_type", ["SIN", "KBD"])
def test_roundtrip_per_frame_is_finite(win_type: WinType) -> None:
    """
    Sanity test: per-frame analysis+synthesis must produce finite outputs.
    """
    rng = np.random.default_rng(0)
    frame_t: FrameT = rng.normal(size=(2048, 2)).astype(np.float64)

    for frame_type in ("OLS", "LSS", "ESH", "LPS"):
        frame_f = aac_filter_bank(frame_t, frame_type, win_type)
        frame_t_hat = aac_i_filter_bank(frame_f, frame_type, win_type)
        assert np.isfinite(frame_t_hat).all()


@pytest.mark.parametrize("win_type", ["SIN", "KBD"])
def test_ola_reconstruction_ols_high_snr(win_type: WinType) -> None:
    """
    Module-level test:
    OLS analysis+synthesis with hop=1024 must reconstruct with high SNR
    in the steady-state region.
    """
    rng = np.random.default_rng(1)

    K = 6
    N = 1024 * (K + 1)
    x: StereoSignal = rng.normal(size=(N, 2)).astype(np.float64)

    y = _ola_reconstruct(x, ["OLS"] * K, win_type)

    a = 1024
    b = N - 1024
    snr = snr_db(x[a:b, :], y[a:b, :])
    assert snr > 50.0


@pytest.mark.parametrize("win_type", ["SIN", "KBD"])
def test_ola_reconstruction_esh_high_snr(win_type: WinType) -> None:
    """
    Module-level test:
    ESH analysis+synthesis with hop=1024 must reconstruct with high SNR
    in the steady-state region.
    """
    rng = np.random.default_rng(2)

    K = 6
    N = 1024 * (K + 1)
    x: StereoSignal = rng.normal(size=(N, 2)).astype(np.float64)

    y = _ola_reconstruct(x, ["ESH"] * K, win_type)

    a = 1024
    b = N - 1024
    snr = snr_db(x[a:b, :], y[a:b, :])
    assert snr > 45.0


@pytest.mark.parametrize("win_type", ["SIN", "KBD"])
def test_ola_reconstruction_transition_sequence(win_type: WinType) -> None:
    """
    Transition sequence test matching the windowing logic:
      OLS -> LSS -> ESH -> LPS -> OLS -> OLS
    """
    rng = np.random.default_rng(3)

    frame_types: list[FrameType] = ["OLS", "LSS", "ESH", "LPS", "OLS", "OLS"]
    K = len(frame_types)
    N = 1024 * (K + 1)
    x: StereoSignal = rng.normal(size=(N, 2)).astype(np.float64)

    y = _ola_reconstruct(x, frame_types, win_type)

    a = 1024
    b = N - 1024
    snr = snr_db(x[a:b, :], y[a:b, :])
    assert snr > 40.0
