# ------------------------------------------------------------
# AAC Coder/Decoder - TNS Tests
#
# Multimedia course at Aristotle University of
# Thessaloniki (AUTh)
#
# Author:
#   Christos Choutouridis (ΑΕΜ 8997)
#   cchoutou@ece.auth.gr
#
# Description:
#   Tests for Temporal Noise Shaping (TNS) module (Level 2).
#
#   Validates:
#   - I/O shapes for long and ESH modes
#   - Quantization grid and clamping of predictor coefficients
#   - Inverse-filter stability (all poles inside unit circle)
#   - Functional correctness: iTNS(TNS(X)) ≈ X
# ------------------------------------------------------------

from __future__ import annotations

import pytest

from core.aac_configuration import PRED_ORDER, QUANT_MAX, QUANT_STEP
from core.aac_tns import aac_tns, aac_i_tns
from core.aac_types import *

EPS = 1e-12

# -----------------------------------------------------------------------------
# Helper utilities
# -----------------------------------------------------------------------------

def _is_inverse_stable_from_coeffs(a_q: MdctCoeffs) -> bool:
    """
    Check stability of the inverse TNS filter H_TNS^{-1}.

    Poles are roots of:
        z^p - a1 z^{p-1} - ... - ap = 0

    Stability condition:
        |root| < 1 for all roots.
    """
    a_q = np.asarray(a_q, dtype=np.float64).reshape(-1)
    p = int(a_q.shape[0])

    poly = np.empty(p + 1, dtype=np.float64)
    poly[0] = 1.0
    poly[1:] = -a_q

    roots = np.roots(poly)
    margin = 1e-12
    return bool(np.all(np.abs(roots) < (1.0 - margin)))


def _assert_quantized_and_clamped(a_q: MdctCoeffs) -> None:
    """
    Assert that coefficients:
    - lie on the QUANT_STEP grid
    - are clamped to [-QUANT_MAX, +QUANT_MAX]
    """
    a_q = np.asarray(a_q, dtype=np.float64)

    assert np.max(np.abs(a_q)) <= (QUANT_MAX + 1e-12)

    grid = a_q / float(QUANT_STEP)
    assert np.max(np.abs(grid - np.round(grid))) < 1e-12


# -----------------------------------------------------------------------------
# Shape / contract tests
# -----------------------------------------------------------------------------

@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_tns_shapes_long_sequences(frame_type: FrameType) -> None:
    """
    Contract test (long frames):
    - Input shape: (1024, 1)
    - Output shape preserved
    - Predictor coeffs shape: (PRED_ORDER, 1)
    """
    rng = np.random.default_rng(0)
    frame_F_in: FrameChannelF = rng.normal(size=(1024, 1)).astype(np.float64)

    frame_F_out, tns_coeffs = aac_tns(frame_F_in, frame_type)

    assert frame_F_out.shape == frame_F_in.shape
    assert tns_coeffs.shape == (PRED_ORDER, 1)


def test_tns_shapes_esh() -> None:
    """
    Contract test (ESH):
    - Input shape: (128, 8)
    - Output shape preserved
    - Predictor coeffs shape: (PRED_ORDER, 8)
    """
    rng = np.random.default_rng(1)
    frame_F_in: FrameChannelF = rng.normal(size=(128, 8)).astype(np.float64)

    frame_F_out, tns_coeffs = aac_tns(frame_F_in, "ESH")

    assert frame_F_out.shape == (128, 8)
    assert tns_coeffs.shape == (PRED_ORDER, 8)


# -----------------------------------------------------------------------------
# Coefficient properties
# -----------------------------------------------------------------------------

@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_tns_coeffs_quantized_clamped_and_stable_long(frame_type: FrameType) -> None:
    """
    Long-frame predictor coefficients must be:
    - quantized on QUANT_STEP grid
    - clamped to [-QUANT_MAX, +QUANT_MAX]
    - stable for inverse filtering
    """
    rng = np.random.default_rng(2)
    frame_F_in: FrameChannelF = rng.normal(size=(1024, 1)).astype(np.float64)

    _, tns_coeffs = aac_tns(frame_F_in, frame_type)

    a_q: MdctCoeffs = tns_coeffs[:, 0]
    _assert_quantized_and_clamped(a_q)
    assert _is_inverse_stable_from_coeffs(a_q)


def test_tns_coeffs_quantized_clamped_and_stable_esh() -> None:
    """
    ESH predictor coefficients must satisfy quantization and stability
    independently for each of the 8 short subframes.
    """
    rng = np.random.default_rng(3)
    frame_F_in: FrameChannelF = rng.normal(size=(128, 8)).astype(np.float64)

    _, tns_coeffs = aac_tns(frame_F_in, "ESH")

    for j in range(8):
        a_q: MdctCoeffs = tns_coeffs[:, j]
        _assert_quantized_and_clamped(a_q)
        assert _is_inverse_stable_from_coeffs(a_q)


# -----------------------------------------------------------------------------
# Functional correctness (round-trip)
# -----------------------------------------------------------------------------

@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_tns_roundtrip_long_is_close(frame_type: FrameType) -> None:
    """
    Functional test:
        iTNS(TNS(X)) ≈ X   for long frames.
    """
    rng = np.random.default_rng(4)
    frame_F_in: FrameChannelF = rng.normal(size=(1024, 1)).astype(np.float64)

    frame_F_tns, tns_coeffs = aac_tns(frame_F_in, frame_type)
    frame_F_hat = aac_i_tns(frame_F_tns, frame_type, tns_coeffs)

    np.testing.assert_allclose(frame_F_hat, frame_F_in, rtol=1e-9, atol=1e-9)


def test_tns_roundtrip_esh_is_close() -> None:
    """
    Functional test:
        iTNS(TNS(X)) ≈ X   for ESH frames (8 independent subframes).
    """
    rng = np.random.default_rng(5)
    frame_F_in: FrameChannelF = rng.normal(size=(128, 8)).astype(np.float64)

    frame_F_tns, tns_coeffs = aac_tns(frame_F_in, "ESH")
    frame_F_hat = aac_i_tns(frame_F_tns, "ESH", tns_coeffs)

    np.testing.assert_allclose(frame_F_hat, frame_F_in, rtol=1e-9, atol=1e-9)


# -----------------------------------------------------------------------------
# Sanity
# -----------------------------------------------------------------------------

def test_tns_outputs_are_finite() -> None:
    """
    Sanity test: no NaN or inf in outputs.
    """
    rng = np.random.default_rng(6)

    frame_F_long: FrameChannelF = rng.normal(size=(1024, 1)).astype(np.float64)
    out_long, coeffs_long = aac_tns(frame_F_long, "OLS")
    assert np.isfinite(out_long).all()
    assert np.isfinite(coeffs_long).all()

    frame_F_esh: FrameChannelF = rng.normal(size=(128, 8)).astype(np.float64)
    out_esh, coeffs_esh = aac_tns(frame_F_esh, "ESH")
    assert np.isfinite(out_esh).all()
    assert np.isfinite(coeffs_esh).all()


@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_tns_zero_input_is_identity_long(frame_type: FrameType) -> None:
    """
    Edge case: zero MDCT coefficients should remain zero after TNS and iTNS.
    This checks that no NaN/Inf appears and the pipeline is numerically safe.
    """
    frame_F_in = np.zeros((1024, 1), dtype=np.float64)

    frame_F_tns, tns_coeffs = aac_tns(frame_F_in, frame_type)
    assert np.isfinite(frame_F_tns).all()
    assert np.isfinite(tns_coeffs).all()
    assert np.all(frame_F_tns == 0.0)

    frame_F_hat = aac_i_tns(frame_F_tns, frame_type, tns_coeffs)
    assert np.isfinite(frame_F_hat).all()
    assert np.all(frame_F_hat == 0.0)


def test_tns_zero_input_is_identity_esh() -> None:
    """
    Edge case: zero MDCT coefficients should remain zero for ESH too.
    """
    frame_F_in = np.zeros((128, 8), dtype=np.float64)

    frame_F_tns, tns_coeffs = aac_tns(frame_F_in, "ESH")
    assert np.isfinite(frame_F_tns).all()
    assert np.isfinite(tns_coeffs).all()
    assert np.all(frame_F_tns == 0.0)

    frame_F_hat = aac_i_tns(frame_F_tns, "ESH", tns_coeffs)
    assert np.isfinite(frame_F_hat).all()
    assert np.all(frame_F_hat == 0.0)


@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_tns_near_silence_is_finite_and_roundtrips(frame_type: FrameType) -> None:
    """
    Edge case: extremely small values should not cause NaN/Inf,
    and round-trip should remain close.
    """
    frame_F_in = (1e-15 * np.ones((1024, 1), dtype=np.float64))

    frame_F_tns, tns_coeffs = aac_tns(frame_F_in, frame_type)
    assert np.isfinite(frame_F_tns).all()
    assert np.isfinite(tns_coeffs).all()

    frame_F_hat = aac_i_tns(frame_F_tns, frame_type, tns_coeffs)
    assert np.isfinite(frame_F_hat).all()

    np.testing.assert_allclose(frame_F_hat, frame_F_in, rtol=1e-6, atol=1e-12)


def test_tns_near_silence_esh_is_finite_and_roundtrips() -> None:
    """
    Near-silence test for ESH mode.
    """
    frame_F_in = (1e-15 * np.ones((128, 8), dtype=np.float64))

    frame_F_tns, tns_coeffs = aac_tns(frame_F_in, "ESH")
    assert np.isfinite(frame_F_tns).all()
    assert np.isfinite(tns_coeffs).all()

    frame_F_hat = aac_i_tns(frame_F_tns, "ESH", tns_coeffs)
    assert np.isfinite(frame_F_hat).all()

    np.testing.assert_allclose(frame_F_hat, frame_F_in, rtol=1e-6, atol=1e-12)


@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_tns_accepts_flat_vector_shape_long(frame_type: FrameType) -> None:
    """
    Contract test: for non-ESH, aac_tns must accept input shape (1024,)
    in addition to (1024, 1), and preserve the shape convention.
    """
    rng = np.random.default_rng(7)
    frame_F_in = rng.normal(size=(1024,)).astype(np.float64)

    frame_F_out, tns_coeffs = aac_tns(frame_F_in, frame_type)

    assert frame_F_out.shape == (1024,)
    assert tns_coeffs.shape == (PRED_ORDER, 1)


@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_tns_does_not_explode_peak_or_energy_long(frame_type: FrameType) -> None:
    """
    Sanity: TNS should not cause extreme peak/energy blow-up on typical inputs.
    This is a loose guard to catch regressions.
    """
    rng = np.random.default_rng(8)
    frame_F_in = rng.normal(size=(1024, 1)).astype(np.float64)

    in_peak = float(np.max(np.abs(frame_F_in)))
    in_energy = float(np.sum(frame_F_in * frame_F_in))

    frame_F_out, _ = aac_tns(frame_F_in, frame_type)

    out_peak = float(np.max(np.abs(frame_F_out)))
    out_energy = float(np.sum(frame_F_out * frame_F_out))

    peak_ratio = out_peak / (in_peak + EPS)
    energy_ratio = out_energy / (in_energy + EPS)

    assert peak_ratio < 50.0
    assert energy_ratio < 2500.0


def test_tns_does_not_explode_peak_or_energy_esh() -> None:
    """
    Sanity: same blow-up guard for ESH mode.
    """
    rng = np.random.default_rng(9)
    frame_F_in = rng.normal(size=(128, 8)).astype(np.float64)

    in_peak = float(np.max(np.abs(frame_F_in)))
    in_energy = float(np.sum(frame_F_in * frame_F_in))

    frame_F_out, _ = aac_tns(frame_F_in, "ESH")

    out_peak = float(np.max(np.abs(frame_F_out)))
    out_energy = float(np.sum(frame_F_out * frame_F_out))

    peak_ratio = out_peak / (in_peak + EPS)
    energy_ratio = out_energy / (in_energy + EPS)

    assert peak_ratio < 50.0
    assert energy_ratio < 2500.0
