# ------------------------------------------------------------
# AAC Coder/Decoder - SNR dB Tests
#
# Multimedia course at Aristotle University of
# Thessaloniki (AUTh)
#
# Author:
#   Christos Choutouridis (ΑΕΜ 8997)
#   cchoutou@ece.auth.gr
#
# Description:
#   - Basic tests for SNR calculation utility.
#   - Contract and sanity tests for TableB219-related utilities.
#
#     These tests do NOT validate the numerical correctness of the
#     Bark tables themselves (they are given by the AAC spec),
#     but instead ensure:
#       - correct loading from disk,
#       - correct table selection per frame type,
#       - internal consistency of band limits,
#       - correct caching behavior.
# ------------------------------------------------------------
from __future__ import annotations

import numpy as np
import pytest

from core.aac_utils import mdct, imdct, snr_db, load_b219_tables, get_table, band_limits
from core.aac_types import *

tolerance = 1e-10

# mdct / imdct
# ------------------------------------------------------------

def _assert_allclose(a: FloatArray, b: FloatArray, *, rtol: float, atol: float) -> None:
    """
    Helper for consistent tolerances across tests.
    """
    np.testing.assert_allclose(a, b, rtol=rtol, atol=atol)


def _estimate_gain(y: MdctCoeffs, x: MdctCoeffs) -> float:
    """
    Estimate scalar gain g such that y ~= g*x in least-squares sense.
    """
    denom = float(np.dot(x, x))
    if denom == 0.0:
        return 0.0
    return float(np.dot(y, x) / denom)



@pytest.mark.parametrize("N", [256, 2048])
def test_mdct_imdct_mdct_identity_up_to_gain(N: int) -> None:
    """
    Consistency test in coefficient domain:
      mdct(imdct(X)) ~= g * X

    For the chosen (non-orthonormal) scaling, g is expected to be close to 2.
    """
    rng = np.random.default_rng(0)
    K = N // 2

    X: MdctCoeffs = rng.normal(size=K).astype(np.float64)
    x: TimeSignal = imdct(X)
    X_hat: MdctCoeffs = mdct(x)

    g = _estimate_gain(X_hat, X)
    _assert_allclose(X_hat, g * X, rtol=tolerance, atol=tolerance)
    _assert_allclose(np.array([g], dtype=np.float64), np.array([2.0], dtype=np.float64), rtol=tolerance, atol=tolerance)


@pytest.mark.parametrize("N", [256, 2048])
def test_mdct_linearity(N: int) -> None:
    """
    Linearity test:
      mdct(a*x + b*y) == a*mdct(x) + b*mdct(y)
    """
    rng = np.random.default_rng(1)
    x: TimeSignal = rng.normal(size=N).astype(np.float64)
    y: TimeSignal = rng.normal(size=N).astype(np.float64)

    a = 0.37
    b = -1.12

    left: MdctCoeffs = mdct(a * x + b * y)
    right: MdctCoeffs = a * mdct(x) + b * mdct(y)

    _assert_allclose(left, right, rtol=tolerance, atol=tolerance)


@pytest.mark.parametrize("N", [256, 2048])
def test_imdct_linearity(N: int) -> None:
    """
    Linearity test for IMDCT:
      imdct(a*X + b*Y) == a*imdct(X) + b*imdct(Y)
    """
    rng = np.random.default_rng(2)
    K = N // 2

    X: MdctCoeffs = rng.normal(size=K).astype(np.float64)
    Y: MdctCoeffs = rng.normal(size=K).astype(np.float64)

    a = -0.5
    b = 2.0

    left: TimeSignal = imdct(a * X + b * Y)
    right: TimeSignal = a * imdct(X) + b * imdct(Y)

    _assert_allclose(left, right, rtol=tolerance, atol=tolerance)


@pytest.mark.parametrize("N", [256, 2048])
def test_mdct_imdct_outputs_are_finite(N: int) -> None:
    """
    Sanity test: no NaN/inf on random inputs.
    """
    rng = np.random.default_rng(3)
    K = N // 2

    x: TimeSignal = rng.normal(size=N).astype(np.float64)
    X: MdctCoeffs = rng.normal(size=K).astype(np.float64)

    X1 = mdct(x)
    x1 = imdct(X)

    assert np.isfinite(X1).all()
    assert np.isfinite(x1).all()



# SNR
# ------------------------------------------------------------

def test_snr_perfect_reconstruction_returns_inf() -> None:
    """
    If x_hat == x_ref exactly, noise power is zero and SNR must be +inf.
    """
    rng = np.random.default_rng(0)
    x: StereoSignal = rng.normal(size=(1024, 2)).astype(np.float64)

    snr = snr_db(x, x)
    assert snr == float("inf")


def test_snr_zero_reference_returns_minus_inf() -> None:
    """
    If reference signal is identically zero, signal power is zero
    and SNR must be -inf (unless noise is also zero, which is degenerate).
    """
    x_ref: StereoSignal = np.zeros((1024, 2), dtype=np.float64)
    x_hat: StereoSignal = np.ones((1024, 2), dtype=np.float64)

    snr = snr_db(x_ref, x_hat)
    assert snr == float("-inf")


def test_snr_known_noise_level_matches_expected_value() -> None:
    """
    Deterministic test with known signal and noise power.

    Let:
      x_ref = ones
      x_hat = ones + noise

    With noise variance sigma^2, expected SNR:
      10 * log10(Ps / Pn)
    """
    n = 1000
    sigma = 0.1

    x_ref: StereoSignal = np.ones((n, 2), dtype=np.float64)
    noise = sigma * np.ones((n, 2), dtype=np.float64)
    x_hat: StereoSignal = x_ref + noise

    ps = float(np.sum(x_ref * x_ref))
    pn = float(np.sum(noise * noise))
    expected = 10.0 * np.log10(ps / pn)

    snr = snr_db(x_ref, x_hat)
    assert np.isclose(snr, expected, rtol=1e-12, atol=1e-12)


def test_snr_aligns_different_lengths_and_channels() -> None:
    """
    The function must:
    - align to minimum length
    - align to minimum channel count
    without crashing.
    """
    rng = np.random.default_rng(1)

    x_ref: StereoSignal = rng.normal(size=(1000, 2)).astype(np.float64)
    x_hat: StereoSignal = rng.normal(size=(800, 1)).astype(np.float64)

    snr = snr_db(x_ref, x_hat)
    assert np.isfinite(snr)


def test_snr_accepts_1d_inputs() -> None:
    """
    1-D inputs must be accepted and treated as single-channel signals.
    """
    rng = np.random.default_rng(2)

    x_ref = rng.normal(size=1024).astype(np.float64)
    x_hat = x_ref + 0.01 * rng.normal(size=1024).astype(np.float64)

    snr = snr_db(x_ref, x_hat)
    assert np.isfinite(snr)



# Table219b
# ------------------------------------------------------------

def test_load_b219_tables_returns_expected_keys() -> None:
    """
    Contract test:
    TableB219.mat must load successfully and expose both tables
    required by the psychoacoustic model.

    The AAC spec defines:
      - B219a: long-frame Bark bands
      - B219b: short-frame Bark bands
    """
    tables = load_b219_tables()

    assert isinstance(tables, dict)
    assert "B219a" in tables
    assert "B219b" in tables


def test_b219_table_shapes_are_correct() -> None:
    """
    Sanity test:
    Verify that the Bark tables have the expected number of bands
    and sufficient columns.

    Expected from AAC spec:
      - B219a: 69 bands (long frames)
      - B219b: 42 bands (short frames)
      - At least 6 columns (as accessed by band_limits()).
    """
    tables = load_b219_tables()

    B219a = tables["B219a"]
    assert B219a.ndim == 2
    assert B219a.shape[0] == 69
    assert B219a.shape[1] >= 6

    B219b = tables["B219b"]
    assert B219b.ndim == 2
    assert B219b.shape[0] == 42
    assert B219b.shape[1] >= 6


def test_get_table_returns_correct_fft_size() -> None:
    """
    Interface test:
    get_table(frame_type) must return both:
      - the correct Bark table
      - the correct FFT size N

    This mapping is fundamental for the psychoacoustic model.
    """
    table_long, N_long = get_table("OLS")
    assert N_long == 2048
    assert table_long.shape[0] == 69

    table_short, N_short = get_table("ESH")
    assert N_short == 256
    assert table_short.shape[0] == 42


def test_band_limits_are_consistent_for_long_table() -> None:
    """
    Sanity test for band limits (long frames):

    For each Bark band:
      - wlow <= whigh
      - frequency indices stay within [0, N/2)
      - all returned arrays have consistent lengths
    """
    table, N = get_table("OLS")
    wlow, whigh, bval, qthr = band_limits(table)

    B = table.shape[0]
    assert len(wlow) == B
    assert len(whigh) == B
    assert len(bval) == B
    assert len(qthr) == B

    for b in range(B):
        assert 0 <= wlow[b] <= whigh[b]
        assert whigh[b] < N // 2


def test_band_limits_are_consistent_for_short_table() -> None:
    """
    Sanity test for band limits (short frames / ESH).

    Same invariants as for long frames, but with FFT size N=256.
    """
    table, N = get_table("ESH")
    wlow, whigh, bval, qthr = band_limits(table)

    B = table.shape[0]
    assert len(wlow) == B
    assert len(whigh) == B

    for b in range(B):
        assert 0 <= wlow[b] <= whigh[b]
        assert whigh[b] < N // 2


def test_b219_tables_are_cached() -> None:
    """
    Implementation test:
    load_b219_tables() should cache the loaded tables so that
    subsequent calls return the same object (identity check).

    This avoids repeated disk I/O during psychoacoustic analysis.
    """
    t1 = load_b219_tables()
    t2 = load_b219_tables()

    assert t1 is t2