# ------------------------------------------------------------
# AAC Coder/Decoder - Huffman wrappers (Level 3)
#
# Multimedia course at Aristotle University of
# Thessaloniki (AUTh)
#
# Author:
#   Christos Choutouridis (ΑΕΜ 8997)
#   cchoutou@ece.auth.gr
#
# Description:
#   Thin wrappers around the provided Huffman utilities (material/huff_utils.py)
#   so that the API matches the assignment text.
#
# Exposed API (assignment):
#   huff_sec, huff_codebook = aac_encode_huff(coeff_sec, huff_LUT_list, force_codebook)
#   dec_coeffs = aac_decode_huff(huff_sec, huff_codebook, huff_LUT_list)
#
# Notes:
#   - Huffman coding operates on tuples. Therefore, decode(encode(x)) may return
#     extra trailing symbols due to tuple padding. The AAC decoder knows the
#     true section length from side information (band limits) and truncates.
# ------------------------------------------------------------
from __future__ import annotations

from typing import Any
import numpy as np

from material.huff_utils import encode_huff, decode_huff


def aac_encode_huff(
    coeff_sec: np.ndarray,
    huff_LUT_list: list[dict[str, Any]],
    force_codebook: int | None = None,
) -> tuple[str, int]:
    """
    Huffman-encode a section of coefficients (MDCT symbols or scalefactors).

    Parameters
    ----------
    coeff_sec : np.ndarray
        Coefficient section to be encoded. Any shape is accepted; the input
        is flattened and treated as a 1-D sequence of int64 symbols.
    huff_LUT_list : list[dict[str, Any]]
        List of Huffman Look-Up Tables (LUTs) as returned by material.load_LUT().
        Index corresponds to codebook id (typically 1..11, with 0 reserved).
    force_codebook : int | None
        If provided, forces the use of this Huffman codebook. In the assignment,
        scalefactors are encoded with codebook 11. For MDCT coefficients, this
        argument is usually omitted (auto-selection).

    Returns
    -------
    tuple[str, int]
        (huff_sec, huff_codebook)
        - huff_sec: bitstream as a string of '0'/'1'
        - huff_codebook: codebook id used by the encoder
    """
    coeff_sec_arr = np.asarray(coeff_sec, dtype=np.int64).reshape(-1)

    if force_codebook is None:
        # Provided utility returns (bitstream, codebook) in the auto-selection case.
        huff_sec, huff_codebook = encode_huff(coeff_sec_arr, huff_LUT_list)
        return str(huff_sec), int(huff_codebook)

    # Provided utility returns ONLY the bitstream when force_codebook is set.
    cb = int(force_codebook)
    huff_sec = encode_huff(coeff_sec_arr, huff_LUT_list, force_codebook=cb)
    return str(huff_sec), cb


def aac_decode_huff(
    huff_sec: str | np.ndarray,
    huff_codebook: int,
    huff_LUT: list[dict[str, Any]],
) -> np.ndarray:
    """
    Huffman-decode a bitstream using the specified codebook.

    Parameters
    ----------
    huff_sec : str | np.ndarray
        Huffman bitstream. Typically a string of '0'/'1'. If an array is provided,
        it is passed through to the provided decoder.
    huff_codebook : int
        Codebook id that was returned by aac_encode_huff.
        Codebook 0 represents an all-zero section.
    huff_LUT : list[dict[str, Any]]
        Huffman LUT list as returned by material.load_LUT().

    Returns
    -------
    np.ndarray
        Decoded coefficients as a 1-D np.int64 array.

        Note: Due to tuple coding, the decoded array may contain extra trailing
        padding symbols. The caller must truncate to the known section length.
    """
    cb = int(huff_codebook)

    if cb == 0:
        # Codebook 0 represents an all-zero section. The decoded length is not
        # recoverable from the bitstream alone; the caller must expand/truncate.
        return np.zeros((0,), dtype=np.int64)

    if cb < 0 or cb >= len(huff_LUT):
        raise ValueError(f"Invalid Huffman codebook index: {cb}")

    lut = huff_LUT[cb]
    dec = decode_huff(huff_sec, lut)
    return np.asarray(dec, dtype=np.int64).reshape(-1)