# ------------------------------------------------------------
# AAC Coder/Decoder - Public Type Aliases
#
# Multimedia course at Aristotle University of
# Thessaloniki (AUTh)
#
# Author:
#   Christos Choutouridis (ΑΕΜ 8997)
#   cchoutou@ece.auth.gr
#
# Description:
#   This module implements Public Type aliases
# ------------------------------------------------------------
from __future__ import annotations

from typing import List, Literal, TypeAlias, TypedDict
import numpy as np
from numpy.typing import NDArray

# -----------------------------------------------------------------------------
# Code enums (for readability; not intended to enforce shapes/lengths)
# -----------------------------------------------------------------------------

FrameType: TypeAlias = Literal["OLS", "LSS", "ESH", "LPS"]
"""
Frame type codes (AAC):
- "OLS": ONLY_LONG_SEQUENCE
- "LSS": LONG_START_SEQUENCE
- "ESH": EIGHT_SHORT_SEQUENCE
- "LPS": LONG_STOP_SEQUENCE
"""

WinType: TypeAlias = Literal["KBD", "SIN"]
"""
Window type codes (AAC):
- "KBD": Kaiser-Bessel-Derived
- "SIN": sinusoid
"""

ChannelKey: TypeAlias = Literal["chl", "chr"]
"""Channel dictionary keys used in Level payloads."""


# -----------------------------------------------------------------------------
# Array “semantic” aliases
#
# Goal: communicate meaning (time/frequency/window, stereo/channel) without
# forcing strict shapes in the type system.
# -----------------------------------------------------------------------------

FloatArray: TypeAlias = NDArray[np.float64]
"""
Generic float64 NumPy array.

Note:
- We standardize internal numeric computations to float64 for stability and
  reproducibility. External I/O can still be float32, but we convert at the
  boundaries.
"""

Window: TypeAlias = FloatArray
"""
Time-domain window (weighting sequence), 1-D.

Typical lengths in this assignment:
- Long: 2048
- Short: 256
- Window sequences for LSS/LPS are also 2048

Expected shape: (N,)
dtype: float64
"""

TimeSignal: TypeAlias = FloatArray
"""
Time-domain signal samples, typically 1-D.

Examples:
- Windowed MDCT input: shape (N,)
- IMDCT output: shape (N,)

dtype: float64
"""

StereoSignal: TypeAlias = FloatArray
"""
Time-domain stereo signal stream.

Expected (typical) shape: (N, 2)
- axis 0: time samples
- axis 1: channels [L, R]

dtype: float64
"""

MdctCoeffs: TypeAlias = FloatArray
"""
MDCT coefficient vector, typically 1-D.

Examples:
- Long: shape (1024,)
- Short: shape (128,)

dtype: float64
"""

MdctFrameChannel: TypeAlias = FloatArray
"""
Per-channel MDCT container used in Level-1/2 sequences.

Typical shapes:
- If frame_type in {"OLS","LSS","LPS"}: (1024, 1) or (1024,)
- If frame_type == "ESH": (128, 8)  (8 short subframes for one channel)

dtype: float64

Notes
-----
Some parts of the assignment store long-frame coefficients as a column vector
(1024, 1) to match MATLAB conventions. Internally you may also use (1024,)
when convenient, but the semantic meaning is identical.
"""

TnsCoeffs: TypeAlias = FloatArray
"""
Quantized TNS predictor coefficients (one channel).

Typical shapes (Level 2):
- If frame_type == "ESH": (4, 8)   (order p=4 for each of the 8 short subframes)
- Else: (4, 1)            (order p=4 for the long frame)

dtype: float64

Notes
-----
The assignment uses a 4-bit uniform symmetric quantizer with step size 0.1.
We store the quantized coefficient values as float64 (typically multiples of 0.1)
to keep the pipeline simple and readable.
"""


FrameT: TypeAlias = FloatArray
"""
Time-domain frame (stereo), as used by the filterbank input/output.

Expected (typical) shape for stereo: (2048, 2)
- axis 0: time samples
- axis 1: channels [L, R]

dtype: float64
"""

FrameChannelT: TypeAlias = FloatArray
"""
Time-domain single-channel frame.

Expected (typical) shape: (2048,)

dtype: float64
"""

FrameF: TypeAlias = FloatArray
"""
Frequency-domain frame (MDCT coefficients), stereo container.

Typical shapes (Level 1):
- If frame_type in {"OLS","LSS","LPS"}: (1024, 2)
- If frame_type == "ESH": (128, 16)

Rationale for ESH (128, 16):
- 8 short subframes per channel => 8 * 2 = 16 columns total
- Each short subframe per stereo is (128, 2), flattened into columns
  in subframe order: [sf0_L, sf0_R, sf1_L, sf1_R, ..., sf7_L, sf7_R]

dtype: float64
"""

FrameChannelF: TypeAlias = MdctFrameChannel
"""
Frequency-domain single-channel MDCT coefficients.

Typical shapes (Level 1/2):
- If frame_type in {"OLS","LSS","LPS"}: (1024, 1) or (1024,)
- If frame_type == "ESH": (128, 8)

dtype: float64
"""

BandRanges: TypeAlias = list[tuple[int, int]]
"""
Bark-band index ranges [start, end] (inclusive) for MDCT lines.

Used by TNS to map MDCT indices k to Bark bands.
"""

BarkTable: TypeAlias = FloatArray
"""
Psychoacoustic Bark band table loaded from TableB219.mat.

Typical shapes:
- Long: (69, 6)
- Short: (42, 6)
"""

BandIndexArray: TypeAlias = NDArray[np.int_]
"""
Array of FFT bin indices per psychoacoustic band.
"""

BandValueArray: TypeAlias = FloatArray
"""
Per-band psychoacoustic values (e.g. Bark position, thresholds).
"""


# Quantizer-related semantic aliases

QuantizedSymbols: TypeAlias = NDArray[np.generic]
"""
Quantized MDCT symbols S(k).

Shapes:
- Always (1024, 1) at the quantizer output (ESH packed to 1024 symbols).
"""

ScaleFactors: TypeAlias = NDArray[np.generic]
"""
DPCM-coded scalefactors sfc(b) = alpha(b) - alpha(b-1).

Shapes:
- Long frames: (NB, 1)
- ESH frames: (NB, 8)
"""

GlobalGain: TypeAlias = float | NDArray[np.generic]
"""
Global gain G = alpha(0).

- Long frames: scalar float
- ESH frames: array shape (1, 8)
"""

# Huffman semantic aliases

HuffmanBitstream: TypeAlias = str
"""Huffman-coded bitstream stored as a string of '0'/'1'."""

HuffmanCodebook: TypeAlias = int
"""Huffman codebook id (e.g., 0..11)."""

# -----------------------------------------------------------------------------
# Level 1 AAC sequence payload types
# -----------------------------------------------------------------------------

class AACChannelFrameF(TypedDict):
    """
    Per-channel payload for aac_seq_1[i]["chl"] or ["chr"] (Level 1).

    Keys
    ----
    frame_F:
        The MDCT coefficients for ONE channel.
        Typical shapes:
        - ESH: (128, 8)   (8 short subframes)
        - else: (1024, 1) or (1024,)
    """
    frame_F: FrameChannelF


class AACSeq1Frame(TypedDict):
    """
    One frame dictionary element of aac_seq_1 (Level 1).
    """
    frame_type: FrameType
    win_type: WinType
    chl: AACChannelFrameF
    chr: AACChannelFrameF


AACSeq1: TypeAlias = List[AACSeq1Frame]
"""
AAC sequence for Level 1:
List of length K (K = number of frames).

Each element is a dict with keys:
- "frame_type", "win_type", "chl", "chr"
"""


# -----------------------------------------------------------------------------
# Level 2 AAC sequence payload types (TNS)
# -----------------------------------------------------------------------------

class AACChannelFrameF2(TypedDict):
    """
    Per-channel payload for aac_seq_2[i]["chl"] or ["chr"] (Level 2).

    Keys
    ----
    frame_F:
        The TNS-processed MDCT coefficients for ONE channel.
        Typical shapes:
        - ESH: (128, 8)
        - else: (1024, 1) or (1024,)
    tns_coeffs:
        Quantized TNS predictor coefficients for ONE channel.
        Typical shapes:
        - ESH: (PRED_ORDER, 8)
        - else: (PRED_ORDER, 1)
    """
    frame_F: FrameChannelF
    tns_coeffs: TnsCoeffs


class AACSeq2Frame(TypedDict):
    """
    One frame dictionary element of aac_seq_2 (Level 2).
    """
    frame_type: FrameType
    win_type: WinType
    chl: AACChannelFrameF2
    chr: AACChannelFrameF2


AACSeq2: TypeAlias = List[AACSeq2Frame]
"""
AAC sequence for Level 2:
List of length K (K = number of frames).

Each element is a dict with keys:
- "frame_type", "win_type", "chl", "chr"

Level 2 adds:
- per-channel "tns_coeffs"
and stores:
- per-channel "frame_F" after applying TNS.
"""

# -----------------------------------------------------------------------------
# Level 3 AAC sequence payload types (Quantizer + Huffman)
# -----------------------------------------------------------------------------

class AACChannelFrameF3(TypedDict):
    """
    Per-channel payload for aac_seq_3[i]["chl"] or ["chr"] (Level 3).

    Keys
    ----
    tns_coeffs:
        Quantized TNS predictor coefficients for ONE channel.
        Shapes:
        - ESH: (PRED_ORDER, 8)
        - else: (PRED_ORDER, 1)

    T:
        Psychoacoustic thresholds per band.
        Shapes:
        - ESH: (NB, 8)
        - else: (NB, 1)
        Note: Stored for completeness / debugging; not entropy-coded.

    G:
        Quantized global gains.
        Shapes:
        - ESH: (1, 8)  (one per short subframe)
        - else: scalar (or compatible np scalar)

    sfc:
        Huffman-coded scalefactor differences (DPCM sequence).

    stream:
        Huffman-coded MDCT quantized symbols S(k) (packed to 1024 symbols).

    codebook:
        Huffman codebook id used for MDCT symbols (stream).
        (Scalefactors typically use fixed codebook 11 and do not need to store it.)
    """
    tns_coeffs: TnsCoeffs
    T: FloatArray
    G: FloatArray | float
    sfc: HuffmanBitstream
    stream: HuffmanBitstream
    codebook: HuffmanCodebook


class AACSeq3Frame(TypedDict):
    """
    One frame dictionary element of aac_seq_3 (Level 3).
    """
    frame_type: FrameType
    win_type: WinType
    chl: AACChannelFrameF3
    chr: AACChannelFrameF3


AACSeq3: TypeAlias = List[AACSeq3Frame]
"""
AAC sequence for Level 3:
List of length K (K = number of frames).

Each element is a dict with keys:
- "frame_type", "win_type", "chl", "chr"

Level 3 adds (per channel):
- "tns_coeffs"
- "T" thresholds (not entropy-coded)
- "G" global gain(s)
- "sfc" Huffman-coded scalefactor differences
- "stream" Huffman-coded MDCT quantized symbols
- "codebook" Huffman codebook for MDCT symbols
"""
