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

from pathlib import Path

import numpy as np
import pytest
import soundfile as sf

from core.aac_coder import aac_coder_1, aac_coder_2, aac_coder_3, aac_read_wav_stereo_48k
from core.aac_decoder import aac_decoder_1, aac_decoder_2, aac_decoder_3, aac_remove_padding
from core.aac_utils import snr_db
from core.aac_types import *


# -----------------------------------------------------------------------------
# Fixtures (small wav logic)
# -----------------------------------------------------------------------------

@pytest.fixture(scope="session")
def wav_in_path() -> Path:
    """
    Provided input WAV used for end-to-end tests.

    Expected project layout:
      source/material/LicorDeCalandraca.wav
    """
    return Path(__file__).resolve().parents[2] / "material" / "LicorDeCalandraca.wav"


@pytest.fixture()
def mk_random_stereo_wav(tmp_path: Path, request: pytest.FixtureRequest) -> Path:
    """
    Create a temporary 48 kHz stereo WAV with random samples.

    Length (in seconds) must be provided via indirect parametrization.
    """
    length = float(request.param)
    rng = np.random.default_rng(123)
    fs = 48000

    n = int(fs * length)
    x: StereoSignal = rng.normal(size=(n, 2)).astype(np.float64)

    wav_path = tmp_path / "in_random.wav"
    sf.write(str(wav_path), x, fs)
    return wav_path


@pytest.fixture()
def mk_actual_stereo_wav(tmp_path: Path, wav_in_path: Path, request: pytest.FixtureRequest) -> Path:
    """
    Create a temporary 48 kHz stereo WAV by chopping from the provided material WAV.
    Length can be overridden via indirect parametrization.
    """
    length = float(getattr(request, "param", 0.25))  # seconds (default: small)

    x, fs = aac_read_wav_stereo_48k(wav_in_path)
    n = int(fs * length)
    x_short = x[:n, :]

    wav_path = tmp_path / "in_actual.wav"
    sf.write(str(wav_path), x_short, fs)
    return wav_path


# -----------------------------------------------------------------------------
# Helper-function tests
# -----------------------------------------------------------------------------

@pytest.mark.parametrize("mk_random_stereo_wav", [2.0], indirect=True)
def test_aac_read_wav_stereo_48k_roundtrip(mk_random_stereo_wav: Path) -> None:
    x, fs = aac_read_wav_stereo_48k(mk_random_stereo_wav)

    assert int(fs) == 48000
    assert isinstance(x, np.ndarray)
    assert x.dtype == np.float64
    assert x.ndim == 2
    assert x.shape[1] == 2
    assert x.shape[0] > 0


def test_aac_remove_padding_removes_hop_from_both_ends() -> None:
    hop = 1024
    n = 10000

    y_pad: StereoSignal = np.zeros((n, 2), dtype=np.float64)
    y: StereoSignal = aac_remove_padding(y_pad, hop=hop)

    assert y.shape == (n - 2 * hop, 2)
    assert y.dtype == np.float64


def test_aac_remove_padding_errors_on_too_short_input() -> None:
    hop = 1024
    y_pad: StereoSignal = np.zeros((2 * hop - 1, 2), dtype=np.float64)

    with pytest.raises(ValueError):
        _ = aac_remove_padding(y_pad, hop=hop)


# -----------------------------------------------------------------------------
# Level 1 tests
# -----------------------------------------------------------------------------

@pytest.mark.parametrize("mk_random_stereo_wav", [0.5], indirect=True)
def test_aac_coder_seq_schema_and_shapes(mk_random_stereo_wav: Path) -> None:
    aac_seq: AACSeq1 = aac_coder_1(mk_random_stereo_wav)

    assert isinstance(aac_seq, list)
    assert len(aac_seq) > 0

    for fr in aac_seq:
        assert isinstance(fr, dict)

        assert "frame_type" in fr
        assert "win_type" in fr
        assert "chl" in fr
        assert "chr" in fr

        frame_type = fr["frame_type"]
        win_type = fr["win_type"]

        assert frame_type in ("OLS", "LSS", "ESH", "LPS")
        assert win_type in ("SIN", "KBD")

        assert isinstance(fr["chl"], dict)
        assert isinstance(fr["chr"], dict)
        assert "frame_F" in fr["chl"]
        assert "frame_F" in fr["chr"]

        chl_f = np.asarray(fr["chl"]["frame_F"], dtype=np.float64)
        chr_f = np.asarray(fr["chr"]["frame_F"], dtype=np.float64)

        if frame_type == "ESH":
            assert chl_f.shape == (128, 8)
            assert chr_f.shape == (128, 8)
        else:
            assert chl_f.shape == (1024, 1)
            assert chr_f.shape == (1024, 1)


@pytest.mark.parametrize("mk_random_stereo_wav", [0.5], indirect=True)
def test_end_to_end_aac_coder_decoder_high_snr(mk_random_stereo_wav: Path, tmp_path: Path) -> None:
    x_ref, fs = sf.read(str(mk_random_stereo_wav), always_2d=True)
    x_ref = np.asarray(x_ref, dtype=np.float64)
    assert int(fs) == 48000

    out_wav = tmp_path / "out.wav"

    aac_seq = aac_coder_1(mk_random_stereo_wav)
    x_hat: StereoSignal = aac_decoder_1(aac_seq, out_wav)

    assert out_wav.exists()
    _, fs_hat = sf.read(str(out_wav), always_2d=True)
    assert int(fs_hat) == 48000

    snr = snr_db(x_ref, x_hat)
    assert snr > 80.0


# -----------------------------------------------------------------------------
# Level 2 tests
# -----------------------------------------------------------------------------

@pytest.mark.parametrize("mk_random_stereo_wav", [0.5], indirect=True)
def test_aac_coder_2_seq_schema_and_shapes(mk_random_stereo_wav: Path) -> None:
    aac_seq: AACSeq2 = aac_coder_2(mk_random_stereo_wav)

    assert isinstance(aac_seq, list)
    assert len(aac_seq) > 0

    for fr in aac_seq:
        assert "frame_type" in fr
        assert "win_type" in fr
        assert "chl" in fr
        assert "chr" in fr

        frame_type: FrameType = fr["frame_type"]
        assert frame_type in ("OLS", "LSS", "ESH", "LPS")

        for ch_key in ("chl", "chr"):
            ch = fr[ch_key]
            assert "frame_F" in ch
            assert "tns_coeffs" in ch

            frame_f = np.asarray(ch["frame_F"], dtype=np.float64)
            coeffs = np.asarray(ch["tns_coeffs"], dtype=np.float64)

            if frame_type == "ESH":
                assert frame_f.shape == (128, 8)
                assert coeffs.shape == (4, 8)
            else:
                assert frame_f.shape == (1024, 1)
                assert coeffs.shape == (4, 1)


@pytest.mark.parametrize("mk_random_stereo_wav", [0.5], indirect=True)
def test_end_to_end_level_2_high_snr(mk_random_stereo_wav: Path, tmp_path: Path) -> None:
    x_ref, fs = sf.read(str(mk_random_stereo_wav), always_2d=True)
    x_ref = np.asarray(x_ref, dtype=np.float64)
    assert int(fs) == 48000

    out_wav = tmp_path / "out_l2.wav"
    aac_seq = aac_coder_2(mk_random_stereo_wav)
    x_hat: StereoSignal = aac_decoder_2(aac_seq, out_wav)

    assert out_wav.exists()
    _, fs_hat = sf.read(str(out_wav), always_2d=True)
    assert int(fs_hat) == 48000

    snr = snr_db(x_ref, x_hat)
    assert snr > 80.0


# -----------------------------------------------------------------------------
# Level 3 tests (Quantizer + Huffman)
# -----------------------------------------------------------------------------

def _assert_level3_frame_schema(frame: AACSeq3Frame) -> None:
    assert "frame_type" in frame
    assert "win_type" in frame
    assert "chl" in frame
    assert "chr" in frame

    for ch_key in ("chl", "chr"):
        ch = frame[ch_key]  # type: ignore[index]
        assert "tns_coeffs" in ch
        assert "T" in ch
        assert "G" in ch
        assert "sfc" in ch
        assert "stream" in ch
        assert "codebook" in ch

        assert isinstance(ch["sfc"], str)
        assert isinstance(ch["stream"], str)
        assert isinstance(ch["codebook"], int)

        assert isinstance(ch["tns_coeffs"], np.ndarray)
        assert isinstance(ch["T"], np.ndarray)
        assert np.isscalar(ch["G"]) or isinstance(ch["G"], np.ndarray)


@pytest.mark.parametrize("mk_actual_stereo_wav", [0.5], indirect=True)
def test_aac_coder_3_seq_schema_and_shapes(mk_actual_stereo_wav: Path) -> None:
    """
    Uses a short WAV excerpt produced by mk_actual_stereo_wav.
    0.11s is enough for a few frames at 48 kHz.
    """
    aac_seq_3: AACSeq3 = aac_coder_3(mk_actual_stereo_wav)

    assert isinstance(aac_seq_3, list)
    assert len(aac_seq_3) > 0

    for fr in aac_seq_3:
        _assert_level3_frame_schema(fr)

        frame_type = fr["frame_type"]
        for ch_key in ("chl", "chr"):
            ch = fr[ch_key]  # type: ignore[index]

            tns = np.asarray(ch["tns_coeffs"])
            if frame_type == "ESH":
                assert tns.ndim == 2
                assert tns.shape[1] == 8
            else:
                assert tns.ndim == 2
                assert tns.shape[1] == 1

            T = np.asarray(ch["T"])
            if frame_type == "ESH":
                assert T.ndim == 2
                assert T.shape[1] == 8
            else:
                assert T.ndim == 2
                assert T.shape[1] == 1

            G = ch["G"]
            if frame_type == "ESH":
                assert isinstance(G, np.ndarray)
                assert np.asarray(G).shape == (1, 8)
            else:
                assert np.isscalar(G)


@pytest.mark.parametrize("mk_actual_stereo_wav", [0.5], indirect=True)
def test_end_to_end_level_3_high_snr(mk_actual_stereo_wav: Path, tmp_path: Path) -> None:
    """
    End-to-end Level 3 using a small WAV excerpt produced by mk_actual_stereo_wav.
    """
    x_ref, fs = aac_read_wav_stereo_48k(mk_actual_stereo_wav)
    assert int(fs) == 48000

    out_wav = tmp_path / "decoded_level3.wav"

    aac_seq_3: AACSeq3 = aac_coder_3(mk_actual_stereo_wav)
    y_hat: StereoSignal = aac_decoder_3(aac_seq_3, out_wav)

    n = min(x_ref.shape[0], y_hat.shape[0])
    s = snr_db(x_ref[:n, :], y_hat[:n, :])
    print(f"SNR={s}")
    assert s > 2.0
