# ------------------------------------------------------------
# AAC Coder/Decoder - Huffman Wrapper Tests (Level 3)
#
# Multimedia course at Aristotle University of
# Thessaloniki (AUTh)
#
# Author:
#   Christos Choutouridis (ΑΕΜ 8997)
#   cchoutou@ece.auth.gr
#
# Description:
#   Contract tests for the Huffman coding stage, using the provided
#   Huffman utilities (material/huff_utils.py).
#
#   The Huffman encoder/decoder itself is GIVEN by the assignment and
#   is not re-implemented here. These tests only verify that:
#
#   - The wrapper functions (aac_encode_huff / aac_decode_huff) expose
#     the API described in the assignment.
#   - Forced codebook selection works as expected (e.g. scalefactors).
#   - Tuple-based Huffman coding semantics are respected.
#
# Notes on tuple coding:
#   Huffman coding operates on tuples of symbols. As a result,
#   decode(encode(x)) may return extra trailing symbols due to padding.
#   The AAC decoder always knows the true section length (from band limits)
#   and truncates accordingly. Therefore, these tests only enforce that
#   the decoded PREFIX matches the original data.
# ------------------------------------------------------------
from __future__ import annotations

import numpy as np
import pytest

from core.aac_huffman import aac_encode_huff, aac_decode_huff
from material.huff_utils import load_LUT


# -----------------------------------------------------------------------------
# Fixtures
# -----------------------------------------------------------------------------

@pytest.fixture(scope="module")
def huff_LUT():
    """
    Load Huffman Look-Up Tables (LUTs) once per test module.

    The LUTs are provided by the assignment (huffCodebooks.mat) via
    material.huff_utils.load_LUT().
    """
    return load_LUT()


# -----------------------------------------------------------------------------
# Roundtrip (prefix) tests
# -----------------------------------------------------------------------------

@pytest.mark.parametrize(
    "coeff_sec",
    [
        np.array([1, -1, 2, -2, 0, 0, 3], dtype=np.int64),
        np.array([0, 0, 0, 0], dtype=np.int64),
        np.array([5], dtype=np.int64),
        np.array([-3, -3, -3, -3], dtype=np.int64),
    ],
)
def test_huffman_roundtrip_prefix_matches(
    coeff_sec: np.ndarray,
    huff_LUT,
) -> None:
    """
    Contract test for Huffman encode/decode.

    Guarantees:
    - Encoding followed by decoding does not crash.
    - The decoded output has at least as many symbols as the input.
    - The prefix of the decoded output matches the original coefficients.

    Rationale:
    Huffman tuple coding may introduce padding, so exact length equality
    is NOT required or expected.
    """
    huff_sec, cb = aac_encode_huff(coeff_sec, huff_LUT)
    dec = aac_decode_huff(huff_sec, cb, huff_LUT)

    if cb == 0:
        # Codebook 0 represents an all-zero section.
        assert np.all(coeff_sec == 0)
        assert dec.size == 0
        return

    assert dec.size >= coeff_sec.size
    np.testing.assert_array_equal(dec[: coeff_sec.size], coeff_sec)


# -----------------------------------------------------------------------------
# Forced codebook tests
# -----------------------------------------------------------------------------

def test_huffman_force_codebook_returns_requested_codebook(huff_LUT) -> None:
    """
    Verify forced codebook selection.

    According to the assignment, scalefactors must be encoded using
    Huffman codebook 11. This test checks that:
    - The requested codebook is actually used.
    - The decoded prefix matches the original scalefactors.
    """
    scalefactors = np.array([10, -2, 1, 0, -1, 3], dtype=np.int64)

    huff_sec, cb = aac_encode_huff(
        scalefactors,
        huff_LUT,
        force_codebook=11,
    )

    assert cb == 11
    assert isinstance(huff_sec, str)

    dec = aac_decode_huff(huff_sec, cb, huff_LUT)

    assert dec.size >= scalefactors.size
    np.testing.assert_array_equal(dec[: scalefactors.size], scalefactors)


# -----------------------------------------------------------------------------
# Error handling
# -----------------------------------------------------------------------------

def test_huffman_invalid_codebook_raises(huff_LUT) -> None:
    """
    Decoding with an invalid Huffman codebook index must raise an error.
    """
    with pytest.raises(Exception):
        _ = aac_decode_huff(
            huff_sec="010101",
            huff_codebook=99,
            huff_LUT=huff_LUT,
        )
