# ------------------------------------------------------------
# AAC Coder/Decoder - Filterbank module
#
# Multimedia course at Aristotle University of
# Thessaloniki (AUTh)
#
# Author:
#   Christos Choutouridis (ΑΕΜ 8997)
#   cchoutou@ece.auth.gr
#
# Description:
#   Filterbank stage (MDCT/IMDCT), windowing, ESH packing/unpacking
#
# ------------------------------------------------------------
from __future__ import annotations

from core.aac_utils import mdct, imdct
from core.aac_types import *

from scipy.signal.windows import kaiser

# Private helpers for Filterbank
# ------------------------------------------------------------

def _sin_window(N: int) -> Window:
    """
    Build a sinusoidal (SIN) window of length N.

    The AAC sinusoid window is:
        w[n] = sin(pi/N * (n + 0.5)),  for 0 <= n < N

    Parameters
    ----------
    N : int
        Window length in samples.

    Returns
    -------
    Window
        1-D array of shape (N, ) with dtype float64.
    """
    n = np.arange(N, dtype=np.float64)
    return np.sin((np.pi / N) * (n + 0.5))


def _kbd_window(N: int, alpha: float) -> Window:
    """
    Build a Kaiser-Bessel-Derived (KBD) window of length N.

    This follows the standard KBD construction used in AAC:
      1) Build a Kaiser kernel of length (N/2 + 1).
      2) Form the left half by cumulative summation, normalization, and sqrt.
      3) Mirror the left half to form the right half (symmetric full-length window).

    Notes
    -----
    - N must be even (AAC uses N=2048 for long and N=256 for short).
    - The assignment specifies alpha=6 for long windows and alpha=4 for short windows.
    - The Kaiser beta parameter is commonly taken as beta = pi * alpha for this context.

    Parameters
    ----------
    N : int
        Window length in samples (must be even).
    alpha : float
        KBD alpha parameter.

    Returns
    -------
    Window
        1-D array of shape (N,) with dtype float64.
    """
    half = N // 2

    # Kaiser kernel length: half + 1 samples (0 .. half)
    # beta = pi * alpha per the usual correspondence with the ISO definition
    kernel = kaiser(half + 1, beta=np.pi * alpha).astype(np.float64)

    csum = np.cumsum(kernel)
    denom = csum[-1]

    w_left = np.sqrt(csum[:-1] / denom)  # length half, n = 0 .. half-1
    w_right = w_left[::-1]               # mirror for second half

    return np.concatenate([w_left, w_right])


def _long_window(win_type: WinType) -> Window:
    """
    Return the long AAC window (length 2048) for the selected window family.

    Parameters
    ----------
    win_type : WinType
        Either "SIN" or "KBD".

    Returns
    -------
    Window
        1-D array of shape (2048,) with dtype float64.
    """
    if win_type == "SIN":
        return _sin_window(2048)
    if win_type == "KBD":
        # Assignment-specific alpha values
        return _kbd_window(2048, alpha=6.0)
    raise ValueError(f"Invalid win_type: {win_type!r}")


def _short_window(win_type: WinType) -> Window:
    """
    Return the short AAC window (length 256) for the selected window family.

    Parameters
    ----------
    win_type : WinType
        Either "SIN" or "KBD".

    Returns
    -------
    Window
        1-D array of shape (256,) with dtype float64.
    """
    if win_type == "SIN":
        return _sin_window(256)
    if win_type == "KBD":
        # Assignment-specific alpha values
        return _kbd_window(256, alpha=4.0)
    raise ValueError(f"Invalid win_type: {win_type!r}")


def _window_sequence(frame_type: FrameType, win_type: WinType) -> Window:
    """
    Build the 2048-sample analysis/synthesis window for OLS/LSS/LPS.

    In this assignment we assume a single window family is used globally
    (no mixed KBD/SIN halves). Therefore, both the long and short windows
    are drawn from the same family.

    For frame_type:
    - "OLS": return the long window Wl (2048).
    - "LSS": construct [Wl_left(1024), ones(448), Ws_right(128), zeros(448)].
    - "LPS": construct [zeros(448), Ws_left(128), ones(448), Wl_right(1024)].

    Parameters
    ----------
    frame_type : FrameType
        One of "OLS", "LSS", "LPS".
    win_type : WinType
        Either "SIN" or "KBD".

    Returns
    -------
    Window
        1-D array of shape (2048,) with dtype float64.
    """
    wL = _long_window(win_type)   # length 2048
    wS = _short_window(win_type)  # length 256

    if frame_type == "OLS":
        return wL

    if frame_type == "LSS":
        # 0..1023: left half of long window
        # 1024..1471: ones (448 samples)
        # 1472..1599: right half of short window (128 samples)
        # 1600..2047: zeros (448 samples)
        out = np.zeros(2048, dtype=np.float64)
        out[0:1024] = wL[0:1024]
        out[1024:1472] = 1.0
        out[1472:1600] = wS[128:256]
        out[1600:2048] = 0.0
        return out

    if frame_type == "LPS":
        # 0..447: zeros (448)
        # 448..575: left half of short window (128)
        # 576..1023: ones (448)
        # 1024..2047: right half of long window (1024)
        out = np.zeros(2048, dtype=np.float64)
        out[0:448] = 0.0
        out[448:576] = wS[0:128]
        out[576:1024] = 1.0
        out[1024:2048] = wL[1024:2048]
        return out

    raise ValueError(f"Invalid frame_type for long window sequence: {frame_type!r}")


def _filter_bank_esh_channel(x_ch: FrameChannelT, win_type: WinType) -> FrameChannelF:
    """
    ESH analysis for one channel.

    Parameters
    ----------
    x_ch : FrameChannelT
        Time-domain channel frame (expected shape: (2048,)).
    win_type : WinType
        Window family ("KBD" or "SIN").

    Returns
    -------
    FrameChannelF
        Array of shape (128, 8). Column j contains the 128 MDCT coefficients
        of the j-th short window.
    """
    wS = _short_window(win_type)  # (256,)
    X_esh = np.empty((128, 8), dtype=np.float64)

    # ESH subwindows are taken from the central region:
    # start positions: 448 + 128*j, j = 0..7
    for j in range(8):
        start = 448 + 128 * j
        seg = x_ch[start:start + 256] * wS  # (256,)
        X_esh[:, j] = mdct(seg)           # (128,)

    return X_esh


def _unpack_esh(frame_F: FrameF) -> tuple[FrameChannelF, FrameChannelF]:
    """
    Unpack ESH spectrum from shape (128, 16) into per-channel arrays (128, 8).

    Parameters
    ----------
    frame_F : FrameF
        Packed ESH spectrum (expected shape: (128, 16)).

    Returns
    -------
    left : FrameChannelF
        Left channel spectrum, shape (128, 8).
    right : FrameChannelF
        Right channel spectrum, shape (128, 8).

    Notes
    -----
    Inverse mapping of the packing used in aac_filter_bank():
      packed[:, 2*j]   = left[:, j]
      packed[:, 2*j+1] = right[:, j]
    """
    if frame_F.shape != (128, 16):
        raise ValueError("ESH frame_F must have shape (128, 16).")

    left = np.empty((128, 8), dtype=np.float64)
    right = np.empty((128, 8), dtype=np.float64)
    for j in range(8):
        left[:, j] = frame_F[:, 2 * j + 0]
        right[:, j] = frame_F[:, 2 * j + 1]
    return left, right


def _i_filter_bank_esh_channel(X_esh: FrameChannelF, win_type: WinType) -> FrameChannelT:
    """
    ESH synthesis for one channel.

    Parameters
    ----------
    X_esh : FrameChannelF
        MDCT coefficients for 8 short windows (expected shape: (128, 8)).
    win_type : WinType
        Window family ("KBD" or "SIN").

    Returns
    -------
    FrameChannelT
        Time-domain channel contribution, shape (2048,).
        This is already overlap-added internally for the 8 short blocks and
        ready for OLA at the caller level.
    """
    if X_esh.shape != (128, 8):
        raise ValueError("X_esh must have shape (128, 8).")

    wS = _short_window(win_type)  # (256,)
    out = np.zeros(2048, dtype=np.float64)

    # Each short IMDCT returns 256 samples. Place them at:
    # start = 448 + 128*j, j=0..7 (50% overlap)
    for j in range(8):
        seg = imdct(X_esh[:, j]) * wS  # (256,)
        start = 448 + 128 * j
        out[start:start + 256] += seg

    return out


# -----------------------------------------------------------------------------
# Public Function prototypes
# -----------------------------------------------------------------------------

def aac_filter_bank(frame_T: FrameT, frame_type: FrameType, win_type: WinType) -> FrameF:
    """
    Filterbank stage (MDCT analysis).

    Parameters
    ----------
    frame_T : FrameT
        Time-domain frame, stereo, shape (2048, 2).
    frame_type : FrameType
        Type of the frame under encoding ("OLS"|"LSS"|"ESH"|"LPS").
    win_type : WinType
        Window type ("KBD" or "SIN") used for the current frame.

    Returns
    -------
    frame_F : FrameF
        Frequency-domain MDCT coefficients:
        - If frame_type in {"OLS","LSS","LPS"}: array shape (1024, 2)
          containing MDCT coefficients for both channels.
        - If frame_type == "ESH": contains 8 subframes, each subframe has shape (128,2),
          placed in columns according to subframe order, i.e. overall shape (128, 16).
    """
    if frame_T.shape != (2048, 2):
        raise ValueError("frame_T must have shape (2048, 2).")

    xL :FrameChannelT = frame_T[:, 0].astype(np.float64, copy=False)
    xR :FrameChannelT = frame_T[:, 1].astype(np.float64, copy=False)

    if frame_type in ("OLS", "LSS", "LPS"):
        w = _window_sequence(frame_type, win_type)  # length 2048
        XL = mdct(xL * w)  # length 1024
        XR = mdct(xR * w)  # length 1024
        out = np.empty((1024, 2), dtype=np.float64)
        out[:, 0] = XL
        out[:, 1] = XR
        return out

    if frame_type == "ESH":
        Xl = _filter_bank_esh_channel(xL, win_type)  # (128, 8)
        Xr = _filter_bank_esh_channel(xR, win_type)  # (128, 8)

        # Pack into (128, 16): each subframe as (128,2) placed in columns
        out = np.empty((128, 16), dtype=np.float64)
        for j in range(8):
            out[:, 2 * j + 0] = Xl[:, j]
            out[:, 2 * j + 1] = Xr[:, j]
        return out

    raise ValueError(f"Invalid frame_type: {frame_type!r}")


def aac_i_filter_bank(frame_F: FrameF, frame_type: FrameType, win_type: WinType) -> FrameT:
    """
    Inverse filterbank (IMDCT synthesis).

    Parameters
    ----------
    frame_F : FrameF
        Frequency-domain MDCT coefficients as produced by filter_bank().
    frame_type : FrameType
        Frame type ("OLS"|"LSS"|"ESH"|"LPS").
    win_type : WinType
        Window type ("KBD" or "SIN").

    Returns
    -------
    frame_T : FrameT
        Reconstructed time-domain frame, stereo, shape (2048, 2).
    """
    if frame_type in ("OLS", "LSS", "LPS"):
        if frame_F.shape != (1024, 2):
            raise ValueError("For OLS/LSS/LPS, frame_F must have shape (1024, 2).")

        w = _window_sequence(frame_type, win_type)

        xL = imdct(frame_F[:, 0]) * w
        xR = imdct(frame_F[:, 1]) * w

        out = np.empty((2048, 2), dtype=np.float64)
        out[:, 0] = xL
        out[:, 1] = xR
        return out

    if frame_type == "ESH":
        if frame_F.shape != (128, 16):
            raise ValueError("For ESH, frame_F must have shape (128, 16).")

        Xl, Xr = _unpack_esh(frame_F)
        xL = _i_filter_bank_esh_channel(Xl, win_type)
        xR = _i_filter_bank_esh_channel(Xr, win_type)

        out = np.empty((2048, 2), dtype=np.float64)
        out[:, 0] = xL
        out[:, 1] = xR
        return out

    raise ValueError(f"Invalid frame_type: {frame_type!r}")
