# ------------------------------------------------------------
# AAC Coder/Decoder - Quantizer Tests
#
# Multimedia course at Aristotle University of
# Thessaloniki (AUTh)
#
# Author:
#   Christos Choutouridis (ΑΕΜ 8997)
#   cchoutou@ece.auth.gr
#
# Description:
#   Tests for Quantizer / iQuantizer module.
#
#   These tests are deliberately "contract-oriented":
#   - They validate shapes, dtypes and invariants that downstream stages
#     (e.g., Huffman coding) depend on.
#   - They do not attempt to validate psychoacoustic optimality (that would
#     require a reference implementation and careful numerical baselines).
#
# Validates:
#   - I/O shapes for long and ESH modes
#   - DPCM scalefactor coding consistency (sfc)
#   - ESH packing order of quantized symbols (128x8 <-> 1024)
#   - Edge cases (zeros / near silence)
#   - Sanity (finite outputs, no extreme numerical blow-up)
# ------------------------------------------------------------
from __future__ import annotations

import numpy as np
import pytest

from core.aac_quantizer import aac_quantizer, aac_i_quantizer
from core.aac_utils import get_table, band_limits
from core.aac_types import FrameType


# Small epsilon to avoid divisions by zero in sanity ratios
EPS = 1e-12


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

def _nbands(frame_type: FrameType) -> int:
    """
    Return number of scalefactor bands for the given frame type.

    This is derived from TableB219 (psycho tables) via aac_utils helpers,
    so the tests remain consistent even if tables are updated.
    """
    table, _nfft = get_table(frame_type)
    wlow, _whigh, _bval, _qthr = band_limits(table)
    return int(len(wlow))


def _make_smr(frame_type: FrameType, seed: int = 0) -> np.ndarray:
    """
    Create a strictly positive SMR array with the correct shape.

    These tests are not about psycho correctness; they only need SMR > 0
    to avoid division by zero and to make the quantizer's threshold logic
    behave deterministically.
    """
    rng = np.random.default_rng(seed)
    NB = _nbands(frame_type)

    if frame_type == "ESH":
        # ESH uses 8 short windows, thus SMR has 8 columns.
        return (1.0 + np.abs(rng.normal(size=(NB, 8)))).astype(np.float64)

    # Long frames: use a column vector (NB, 1).
    return (1.0 + np.abs(rng.normal(size=(NB, 1)))).astype(np.float64)


def _reconstruct_alpha_from_sfc(sfc: np.ndarray) -> np.ndarray:
    """
    Reconstruct alpha(b) from DPCM-coded scalefactors sfc(b).

    By definition in the assignment:
        sfc(0) = alpha(0)
        alpha(b) = alpha(b-1) + sfc(b)  for b > 0

    This reconstruction is useful to validate the internal consistency
    of the produced scalefactor information.
    """
    sfc = np.asarray(sfc, dtype=np.int64)

    # Long frames: sfc shape (NB, 1)
    if sfc.ndim == 2 and sfc.shape[1] == 1:
        NB = sfc.shape[0]
        alpha = np.zeros((NB,), dtype=np.int64)
        alpha[0] = int(sfc[0, 0])
        for b in range(1, NB):
            alpha[b] = int(alpha[b - 1] + sfc[b, 0])
        return alpha

    # ESH frames: sfc shape (NB, 8)
    if sfc.ndim == 2 and sfc.shape[1] == 8:
        NB = sfc.shape[0]
        alpha = np.zeros((NB, 8), dtype=np.int64)
        alpha[0, :] = sfc[0, :]
        for b in range(1, NB):
            alpha[b, :] = alpha[b - 1, :] + sfc[b, :]
        return alpha

    raise ValueError("Unsupported sfc shape.")


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

@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_quantizer_shapes_long(frame_type: FrameType) -> None:
    """
    Contract test for long frames:
    - Input: MDCT coefficients shape (1024, 1)
    - Output S: always (1024, 1)
    - Output sfc: (NB, 1)
    - G: scalar float for long frames
    - iQuantizer output: (1024, 1)
    """
    NB = _nbands(frame_type)
    rng = np.random.default_rng(1)

    X = rng.normal(size=(1024, 1)).astype(np.float64)
    SMR = _make_smr(frame_type, seed=2)

    S, sfc, G = aac_quantizer(X, frame_type, SMR)

    assert S.shape == (1024, 1)
    assert sfc.shape == (NB, 1)
    assert isinstance(G, (float, np.floating))

    Xhat = aac_i_quantizer(S, sfc, G, frame_type)
    assert Xhat.shape == (1024, 1)


def test_quantizer_shapes_esh() -> None:
    """
    Contract test for ESH frames:
    - Input: MDCT coefficients shape (128, 8)
    - Output S: packed to (1024, 1)
    - Output sfc: (NB, 8)
    - G: array shape (1, 8) for ESH (one gain per short window)
    - iQuantizer output: (128, 8)
    """
    frame_type: FrameType = "ESH"
    NB = _nbands(frame_type)
    rng = np.random.default_rng(3)

    X = rng.normal(size=(128, 8)).astype(np.float64)
    SMR = _make_smr(frame_type, seed=4)

    S, sfc, G = aac_quantizer(X, frame_type, SMR)

    assert S.shape == (1024, 1)
    assert sfc.shape == (NB, 8)
    assert isinstance(G, np.ndarray)
    assert G.shape == (1, 8)

    Xhat = aac_i_quantizer(S, sfc, G, frame_type)
    assert Xhat.shape == (128, 8)


# -----------------------------------------------------------------------------
# DPCM consistency tests
# -----------------------------------------------------------------------------

@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_quantizer_dpcm_reconstructs_alpha_long(frame_type: FrameType) -> None:
    """
    Verify the DPCM coding rule for long frames.

    The quantizer returns:
      sfc(0) = alpha(0)
      sfc(b) = alpha(b) - alpha(b-1), b>0

    Reconstruct alpha from sfc and check:
      alpha(0) == sfc(0) == G
    """
    rng = np.random.default_rng(5)

    X = rng.normal(size=(1024, 1)).astype(np.float64)
    SMR = _make_smr(frame_type, seed=6)

    _S, sfc, G = aac_quantizer(X, frame_type, SMR)

    alpha = _reconstruct_alpha_from_sfc(sfc)

    assert int(sfc[0, 0]) == int(alpha[0])
    assert float(alpha[0]) == float(G)


def test_quantizer_dpcm_reconstructs_alpha_esh() -> None:
    """
    Verify the DPCM coding rule for ESH frames.

    For each short window j:
      sfc(0, j) = alpha(0, j) == G(0, j)
    """
    frame_type: FrameType = "ESH"
    rng = np.random.default_rng(7)

    X = rng.normal(size=(128, 8)).astype(np.float64)
    SMR = _make_smr(frame_type, seed=8)

    _S, sfc, G = aac_quantizer(X, frame_type, SMR)

    alpha = _reconstruct_alpha_from_sfc(sfc)

    assert np.all(alpha[0, :] == sfc[0, :])
    assert np.all(alpha[0, :] == G.reshape(-1))


# -----------------------------------------------------------------------------
# ESH packing order test
# -----------------------------------------------------------------------------

def test_quantizer_esh_packing_order_matches_iquantizer_layout() -> None:
    """
    Verify ESH packing order.

    The quantizer outputs S in packed shape (1024, 1). The expected packing
    is column-major concatenation of the 8 short subframes.

    This test constructs a deterministic input where each subframe column
    has a distinct constant value. After quantize+inverse-quantize, the
    reconstructed columns should remain distinguishable in the same order.

    This primarily tests ordering, not exact numerical values.
    """
    frame_type: FrameType = "ESH"
    NB = _nbands(frame_type)

    # Create 8 distinct subframes: column j is constant (j+1)
    X = np.zeros((128, 8), dtype=np.float64)
    for j in range(8):
        X[:, j] = float(j + 1)

    # Use very large SMR so thresholds are permissive and alpha changes are
    # minimal. This helps keep the ordering signal strong.
    SMR = np.ones((NB, 8), dtype=np.float64) * 1e6

    S, sfc, G = aac_quantizer(X, frame_type, SMR)
    Xhat = aac_i_quantizer(S, sfc, G, frame_type)

    # The average magnitude per column must be increasing with the original order.
    col_means = np.mean(Xhat, axis=0)
    assert np.all(np.diff(col_means) > 0.0)


# -----------------------------------------------------------------------------
# Edge cases: zeros and near-silence
# -----------------------------------------------------------------------------

@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_quantizer_zero_input_long_is_finite(frame_type: FrameType) -> None:
    """
    Edge case: zero MDCT coefficients should not produce NaN/Inf.

    We do not require identity here (quantizer is lossy), but we require
    the pipeline to remain numerically safe and produce finite outputs.
    """
    NB = _nbands(frame_type)

    X = np.zeros((1024, 1), dtype=np.float64)
    SMR = np.ones((NB, 1), dtype=np.float64)

    S, sfc, G = aac_quantizer(X, frame_type, SMR)
    assert np.isfinite(S).all()
    assert np.isfinite(sfc).all()
    assert isinstance(G, (float, np.floating))

    Xhat = aac_i_quantizer(S, sfc, G, frame_type)
    assert np.isfinite(Xhat).all()


def test_quantizer_zero_input_esh_is_finite() -> None:
    """
    Edge case: same as above, for ESH mode.
    """
    frame_type: FrameType = "ESH"
    NB = _nbands(frame_type)

    X = np.zeros((128, 8), dtype=np.float64)
    SMR = np.ones((NB, 8), dtype=np.float64)

    S, sfc, G = aac_quantizer(X, frame_type, SMR)
    assert np.isfinite(S).all()
    assert np.isfinite(sfc).all()
    assert np.isfinite(G).all()

    Xhat = aac_i_quantizer(S, sfc, G, frame_type)
    assert np.isfinite(Xhat).all()


@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_quantizer_near_silence_long_is_finite(frame_type: FrameType) -> None:
    """
    Edge case: extremely small values.

    This stresses numerical guards (EPS usage) and ensures no invalid operations.
    """
    NB = _nbands(frame_type)

    X = (1e-15 * np.ones((1024, 1), dtype=np.float64))
    SMR = np.ones((NB, 1), dtype=np.float64)

    S, sfc, G = aac_quantizer(X, frame_type, SMR)
    assert np.isfinite(S).all()
    assert np.isfinite(sfc).all()

    Xhat = aac_i_quantizer(S, sfc, G, frame_type)
    assert np.isfinite(Xhat).all()


def test_quantizer_near_silence_esh_is_finite() -> None:
    """
    Edge case: extremely small values, ESH mode.
    """
    frame_type: FrameType = "ESH"
    NB = _nbands(frame_type)

    X = (1e-15 * np.ones((128, 8), dtype=np.float64))
    SMR = np.ones((NB, 8), dtype=np.float64)

    S, sfc, G = aac_quantizer(X, frame_type, SMR)
    assert np.isfinite(S).all()
    assert np.isfinite(sfc).all()
    assert np.isfinite(G).all()

    Xhat = aac_i_quantizer(S, sfc, G, frame_type)
    assert np.isfinite(Xhat).all()


# -----------------------------------------------------------------------------
# Sanity: avoid catastrophic numerical blow-up
# -----------------------------------------------------------------------------

@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_quantizer_sanity_no_extreme_blowup_long(frame_type: FrameType) -> None:
    """
    Loose sanity guard.

    The quantizer is lossy, but it should not produce reconstructions with
    catastrophic peak/energy growth compared to the input.
    """
    NB = _nbands(frame_type)
    rng = np.random.default_rng(11)

    X = rng.normal(size=(1024, 1)).astype(np.float64)
    SMR = np.ones((NB, 1), dtype=np.float64) * 10.0

    S, sfc, G = aac_quantizer(X, frame_type, SMR)
    Xhat = aac_i_quantizer(S, sfc, G, frame_type)

    in_peak = float(np.max(np.abs(X)))
    out_peak = float(np.max(np.abs(Xhat)))
    peak_ratio = out_peak / (in_peak + EPS)

    in_energy = float(np.sum(X * X))
    out_energy = float(np.sum(Xhat * Xhat))
    energy_ratio = out_energy / (in_energy + EPS)

    # Very loose thresholds: only catch severe regressions.
    assert peak_ratio < 100.0
    assert energy_ratio < 1e4


def test_quantizer_sanity_no_extreme_blowup_esh() -> None:
    """
    Same loose sanity guard for ESH mode.
    """
    frame_type: FrameType = "ESH"
    NB = _nbands(frame_type)
    rng = np.random.default_rng(12)

    X = rng.normal(size=(128, 8)).astype(np.float64)
    SMR = np.ones((NB, 8), dtype=np.float64) * 10.0

    S, sfc, G = aac_quantizer(X, frame_type, SMR)
    Xhat = aac_i_quantizer(S, sfc, G, frame_type)

    in_peak = float(np.max(np.abs(X)))
    out_peak = float(np.max(np.abs(Xhat)))
    peak_ratio = out_peak / (in_peak + EPS)

    in_energy = float(np.sum(X * X))
    out_energy = float(np.sum(Xhat * Xhat))
    energy_ratio = out_energy / (in_energy + EPS)

    assert peak_ratio < 100.0
    assert energy_ratio < 1e4
