# ------------------------------------------------------------
# AAC Coder/Decoder - Level 1 Wrappers + Demo
#
# Multimedia course at Aristotle University of
# Thessaloniki (AUTh)
#
# Author:
#   Christos Choutouridis (ΑΕΜ 8997)
#   cchoutou@ece.auth.gr
#
# Description:
#   Level 1 wrapper module.
#
#   This file provides:
#   - Thin wrappers for Level 1 API functions (encode/decode) that delegate
#     to the corresponding core implementations.
#   - A demo function that runs end-to-end and computes SNR.
#   - A small CLI entrypoint for convenience.
# ------------------------------------------------------------
from __future__ import annotations

from pathlib import Path
from typing import Union

import numpy as np
import soundfile as sf

from core.aac_types   import AACSeq1, StereoSignal
from core.aac_coder   import aac_coder_1 as core_aac_coder_1
from core.aac_coder   import aac_read_wav_stereo_48k
from core.aac_decoder import aac_decoder_1 as core_aac_decoder_1


# -----------------------------------------------------------------------------
# Public Level 1 API (wrappers)
# -----------------------------------------------------------------------------

def aac_coder_1(filename_in: Union[str, Path]) -> AACSeq1:
    """
    Level-1 AAC encoder (wrapper).

    Delegates to core implementation.

    Parameters
    ----------
    filename_in : Union[str, Path]
        Input WAV filename.
        Assumption: stereo audio, sampling rate 48 kHz.

    Returns
    -------
    AACSeq1
        List of encoded frames (Level 1 schema).
    """
    return core_aac_coder_1(filename_in)


def aac_decoder_1(aac_seq_1: AACSeq1, filename_out: Union[str, Path]) -> StereoSignal:
    """
    Level-1 AAC decoder (wrapper).

    Delegates to core implementation.

    Parameters
    ----------
    aac_seq_1 : AACSeq1
        Encoded sequence as produced by aac_coder_1().
    filename_out : Union[str, Path]
        Output WAV filename. Assumption: 48 kHz, stereo.

    Returns
    -------
    StereoSignal
        Decoded audio samples (time-domain), stereo, shape (N, 2), dtype float64.
    """
    return core_aac_decoder_1(aac_seq_1, filename_out)


# -----------------------------------------------------------------------------
# Demo (Level 1)
# -----------------------------------------------------------------------------

def _snr_db(x_ref: StereoSignal, x_hat: StereoSignal) -> float:
    """
    Compute overall SNR (dB) over all samples and channels after aligning lengths.

    Parameters
    ----------
    x_ref : StereoSignal
        Reference stereo stream.
    x_hat : StereoSignal
        Reconstructed stereo stream.

    Returns
    -------
    float
        SNR in dB.
        - Returns +inf if noise power is zero.
        - Returns -inf if signal power is zero.
    """
    x_ref = np.asarray(x_ref, dtype=np.float64)
    x_hat = np.asarray(x_hat, dtype=np.float64)

    if x_ref.ndim == 1:
        x_ref = x_ref.reshape(-1, 1)
    if x_hat.ndim == 1:
        x_hat = x_hat.reshape(-1, 1)

    n = min(x_ref.shape[0], x_hat.shape[0])
    c = min(x_ref.shape[1], x_hat.shape[1])

    x_ref = x_ref[:n, :c]
    x_hat = x_hat[:n, :c]

    err = x_ref - x_hat
    ps = float(np.sum(x_ref * x_ref))
    pn = float(np.sum(err * err))

    if pn <= 0.0:
        return float("inf")
    if ps <= 0.0:
        return float("-inf")

    return float(10.0 * np.log10(ps / pn))


def demo_aac_1(filename_in: Union[str, Path], filename_out: Union[str, Path]) -> float:
    """
    Demonstration for the Level-1 codec.

    Runs:
    - aac_coder_1(filename_in)
    - aac_decoder_1(aac_seq_1, filename_out)
    and computes total SNR between original and decoded audio.

    Parameters
    ----------
    filename_in : Union[str, Path]
        Input WAV filename (stereo, 48 kHz).
    filename_out : Union[str, Path]
        Output WAV filename (stereo, 48 kHz).

    Returns
    -------
    float
        Overall SNR in dB.
    """
    filename_in = Path(filename_in)
    filename_out = Path(filename_out)

    # Read original audio (reference) with the same validation as the codec.
    x_ref, fs_ref = aac_read_wav_stereo_48k(filename_in)
    if int(fs_ref) != 48000:
        raise ValueError("Input sampling rate must be 48 kHz.")

    # Encode / decode
    aac_seq_1 = aac_coder_1(filename_in)
    x_hat = aac_decoder_1(aac_seq_1, filename_out)

    # Optional sanity: ensure output file exists and is readable
    x_hat_file, fs_hat = sf.read(str(filename_out), always_2d=True)
    _ = x_hat_file
    if int(fs_hat) != 48000:
        raise ValueError("Decoded output sampling rate must be 48 kHz.")

    return _snr_db(x_ref, x_hat)


# -----------------------------------------------------------------------------
# CLI
# -----------------------------------------------------------------------------

if __name__ == "__main__":
    # Example:
    #   python -m level_1.level_1 input.wav output.wav
    import sys

    if len(sys.argv) != 3:
        raise SystemExit("Usage: python -m level_1.level_1 <input.wav> <output.wav>")

    in_wav = Path(sys.argv[1])
    out_wav = Path(sys.argv[2])

    print(f"Encoding/Decoding {in_wav} to {out_wav}")
    snr = demo_aac_1(in_wav, out_wav)
    print(f"SNR = {snr:.3f} dB")