# ------------------------------------------------------------
# AAC Coder/Decoder - Psychoacoustic Model Tests
#
# Multimedia course at Aristotle University of
# Thessaloniki (AUTh)
#
# Author:
#   Christos Choutouridis (ΑΕΜ 8997)
#
# Description:
#   Contract + sanity tests for the psychoacoustic model (core.aac_psycho).
#
#   These tests focus on:
#   - output shapes per frame_type (long vs ESH),
#   - numerical sanity (finite, non-negative),
#   - deterministic behavior,
#   - ESH central-region dependency (outer regions must not affect result),
#   - basic input validation (length checks).
#
#   We intentionally avoid asserting exact numeric values, because the model
#   includes FFT operations and table-driven psychoacoustic parameters.
# ------------------------------------------------------------
from __future__ import annotations

import numpy as np
import pytest

from core.aac_psycho import aac_psycho
from core.aac_types import FrameChannelT, FrameType


# -----------------------------------------------------------------------------
# Helpers
# -----------------------------------------------------------------------------

def _make_frames(
    *,
    kind: str,
    amp: float = 1.0,
    seed: int = 0,
) -> tuple[FrameChannelT, FrameChannelT, FrameChannelT]:
    """
    Create (current, prev1, prev2) 2048-sample frames for one channel.

    Parameters
    ----------
    kind : str
        "noise" or "tone".
    amp : float
        Amplitude scaling (applied to all frames).
    seed : int
        RNG seed for reproducibility (noise case).

    Returns
    -------
    tuple[FrameChannelT, FrameChannelT, FrameChannelT]
        Three arrays of shape (2048,), dtype float64.
    """
    if kind == "noise":
        rng = np.random.default_rng(seed)
        x2 = amp * rng.normal(size=2048).astype(np.float64)
        x1 = amp * rng.normal(size=2048).astype(np.float64)
        x0 = amp * rng.normal(size=2048).astype(np.float64)
        return x0, x1, x2

    if kind == "tone":
        # A simple sinusoid which is identical across frames (highly predictable).
        n = np.arange(2048, dtype=np.float64)
        f0 = 13.0  # arbitrary normalized-bin-ish tone (not critical for these tests)
        tone = amp * np.sin(2.0 * np.pi * f0 * n / 2048.0).astype(np.float64)
        return tone, tone.copy(), tone.copy()

    raise ValueError(f"Unknown kind: {kind!r}")


def _assert_finite_nonnegative(x: np.ndarray) -> None:
    """Utility assertions for psycho outputs."""
    assert np.isfinite(x).all()
    # SMR is a ratio of energies, it should not be negative.
    assert np.min(x) >= 0.0


# -----------------------------------------------------------------------------
# Shape / contract tests
# -----------------------------------------------------------------------------

@pytest.mark.parametrize("frame_type", ["OLS", "LSS", "LPS"])
def test_psycho_long_shapes(frame_type: FrameType) -> None:
    """
    Contract test:
    For long frame types, psycho returns SMR shape (69,).
    """
    x0, x1, x2 = _make_frames(kind="noise", seed=1, amp=1.0)

    smr = aac_psycho(x0, frame_type, x1, x2)
    assert isinstance(smr, np.ndarray)
    assert smr.shape == (69,)
    _assert_finite_nonnegative(smr)


def test_psycho_esh_shape() -> None:
    """
    Contract test:
    For ESH, psycho returns SMR shape (42, 8).
    """
    x0, x1, x2 = _make_frames(kind="noise", seed=2, amp=1.0)

    smr = aac_psycho(x0, "ESH", x1, x2)
    assert isinstance(smr, np.ndarray)
    assert smr.shape == (42, 8)
    _assert_finite_nonnegative(smr)


def test_psycho_is_deterministic_for_same_inputs() -> None:
    """
    Determinism test:
    Psycho must return the same output for identical inputs.
    """
    x0, x1, x2 = _make_frames(kind="noise", seed=3, amp=1.0)

    smr1 = aac_psycho(x0, "OLS", x1, x2)
    smr2 = aac_psycho(x0, "OLS", x1, x2)

    np.testing.assert_allclose(smr1, smr2, rtol=0.0, atol=0.0)


# -----------------------------------------------------------------------------
# ESH-specific behavior tests
# -----------------------------------------------------------------------------

def test_psycho_esh_ignores_outer_regions() -> None:
    """
    Spec-driven behavior test:
    In this project, ESH uses only the central region of the 2048-sample frame to
    derive the 8 overlapping 256-sample subframes:
        start = 448 + 128*j, j=0..7

    Therefore, changing samples outside [448, 1600) must not affect the output.
    """
    rng = np.random.default_rng(10)

    # Build base frames (current and prev1) with identical central region.
    center_cur = rng.normal(size=1152).astype(np.float64)
    center_prev1 = rng.normal(size=1152).astype(np.float64)

    cur_a = np.zeros(2048, dtype=np.float64)
    cur_b = np.zeros(2048, dtype=np.float64)
    prev1_a = np.zeros(2048, dtype=np.float64)
    prev1_b = np.zeros(2048, dtype=np.float64)

    cur_a[448:1600] = center_cur
    cur_b[448:1600] = center_cur
    prev1_a[448:1600] = center_prev1
    prev1_b[448:1600] = center_prev1

    # Modify only outer regions in the *_b variants.
    cur_b[:448] = rng.normal(size=448)
    cur_b[1600:] = rng.normal(size=448)
    prev1_b[:448] = rng.normal(size=448)
    prev1_b[1600:] = rng.normal(size=448)

    # prev2 is irrelevant for the chosen ESH history convention; keep it fixed.
    prev2 = rng.normal(size=2048).astype(np.float64)

    smr_a = aac_psycho(cur_a, "ESH", prev1_a, prev2)
    smr_b = aac_psycho(cur_b, "ESH", prev1_b, prev2)

    np.testing.assert_allclose(smr_a, smr_b, rtol=0.0, atol=0.0)


def test_psycho_esh_columns_are_not_all_identical_for_random_input() -> None:
    """
    Sanity test:
    For random input, different ESH subframes should typically produce
    different SMR columns (not a strict requirement, but a strong sanity signal).

    We check that at least one column differs from another beyond a tiny tolerance.
    """
    x0, x1, x2 = _make_frames(kind="noise", seed=11, amp=1.0)

    smr = aac_psycho(x0, "ESH", x1, x2)
    # Compare column 0 vs column 7; for random signals they should differ.
    diff = np.max(np.abs(smr[:, 0] - smr[:, 7]))
    assert diff > 1e-12


# -----------------------------------------------------------------------------
# Scaling sanity (avoid fragile numeric targets)
# -----------------------------------------------------------------------------

def test_psycho_long_smr_is_mostly_monotone_with_amplitude() -> None:
    """
    Sanity test:
    Increasing signal amplitude should not reduce the SMR for the vast majority
    of Bark bands.

    Due to the use of max(nb, qthr), a small fraction of bands close to the
    threshold-in-quiet boundary may violate strict monotonicity. This is expected
    behavior, so we test a percentage-based criterion instead of a strict one.
    """
    x0, x1, x2 = _make_frames(kind="noise", seed=20, amp=1e3)
    y0, y1, y2 = (2.0 * x0, 2.0 * x1, 2.0 * x2)

    smr1 = aac_psycho(x0, "OLS", x1, x2)
    smr2 = aac_psycho(y0, "OLS", y1, y2)

    eps = 1e-12
    nondecreasing = np.sum(smr2 + eps >= smr1)

    ratio = nondecreasing / smr1.size

    # Expect monotonic behavior for the overwhelming majority of bands.
    assert ratio >= 0.95

def test_psycho_long_is_approximately_scale_invariant_at_high_level() -> None:
    """
    Sanity test (robust):
    At high levels, SMR should be approximately scale-invariant for most bands.
    Some bands may deviate due to the max(nb, qthr) branch.
    """
    x0, x1, x2 = _make_frames(kind="noise", seed=20, amp=1e3)
    y0, y1, y2 = (2.0 * x0, 2.0 * x1, 2.0 * x2)

    smr1 = aac_psycho(x0, "OLS", x1, x2)
    smr2 = aac_psycho(y0, "OLS", y1, y2)

    rel = np.abs(smr2 - smr1) / np.maximum(np.abs(smr1), 1e-12)

    # Most bands should be close (<= 5%), but allow a small number of outliers.
    close = np.sum(rel <= 5e-2)
    assert close >= (smr1.size - 2)  # allow up to 2 bands to deviate


# -----------------------------------------------------------------------------
# Input validation tests
# -----------------------------------------------------------------------------

def test_psycho_rejects_wrong_lengths() -> None:
    """
    Contract test:
    aac_psycho requires 2048-sample frames for current/prev1/prev2.
    """
    x = np.zeros(2048, dtype=np.float64)
    bad = np.zeros(2047, dtype=np.float64)

    with pytest.raises(ValueError):
        _ = aac_psycho(bad, "OLS", x, x)

    with pytest.raises(ValueError):
        _ = aac_psycho(x, "OLS", bad, x)

    with pytest.raises(ValueError):
        _ = aac_psycho(x, "OLS", x, bad)
