# ------------------------------------------------------------
# AAC Coder/Decoder - Sequence Segmentation Control Tests
#
# Multimedia course at Aristotle University of
# Thessaloniki (AUTh)
#
# Author:
#   Christos Choutouridis (ΑΕΜ 8997)
#   cchoutou@ece.auth.gr
#
# Description:
#   Tests for Sequence Segmentation Control module (SSC).
# ------------------------------------------------------------

from __future__ import annotations

import numpy as np

from core.aac_ssc import aac_ssc
from core.aac_types import FrameT

# -----------------------------------------------------------------------------
# Helper fixtures for SSC
# -----------------------------------------------------------------------------

def _next_frame_no_attack() -> FrameT:
    """
    Build a next_frame_T that must NOT trigger ESH detection.

    Uses exact zeros so all segment energies are zero and the condition
    s[l] > 1e-3 cannot hold for any l.
    """
    return np.zeros((2048, 2), dtype=np.float64)


def _next_frame_strong_attack(
    *,
    attack_left: bool,
    attack_right: bool,
    segment_l: int = 4,
    baseline: float = 1e-6,
    burst_amp: float = 1.0,
) -> FrameT:
    """
    Build a next_frame_T (2048x2) that should trigger ESH detection on selected channels.

    Attack criterion (spec):
      Attack exists if there exists l in {1..7} such that:
        s[l] > 1e-3  and  ds[l] > 10,
      where s[l] is the energy of segment l (length 128) after high-pass filtering,
      and ds[l] = s[l] / s[l-1].

    Construction:
    - A small baseline is added everywhere to avoid relying on the epsilon guard in ds,
      keeping ds behavior stable/reproducible.
    - A strong burst is added inside a chosen segment l in 1..7.
    """
    if not (1 <= segment_l <= 7):
        raise ValueError(f"segment_l must be in [1, 7], got {segment_l}.")

    x = np.full((2048, 2), baseline, dtype=np.float64)

    a = segment_l * 128
    b = (segment_l + 1) * 128

    if attack_left:
        x[a:b, 0] += burst_amp
    if attack_right:
        x[a:b, 1] += burst_amp

    return x


def _next_frame_below_s_threshold(
    *,
    left: bool,
    right: bool,
    segment_l: int = 4,
    impulse_amp: float = 0.01,
) -> FrameT:
    """
    Construct a next_frame_T where s[l] is below 1e-3, so ESH must NOT be triggered,
    even if the ratio ds[l] could be large.

    We place a single impulse of amplitude 'impulse_amp' inside one segment.
    Approx. segment energy: s[l] ~= impulse_amp^2.

    Example:
      impulse_amp = 0.01 => s[l] ~= 1e-4 < 1e-3
    """
    if not (1 <= segment_l <= 7):
        raise ValueError(f"segment_l must be in [1, 7], got {segment_l}.")

    x = np.zeros((2048, 2), dtype=np.float64)

    idx = segment_l * 128 + 10  # inside segment l
    if left:
        x[idx, 0] = impulse_amp
    if right:
        x[idx, 1] = impulse_amp

    return x


# -----------------------------------------------------------------------------
# 1) Fixed/mandatory cases (prev frame type forces current type)
# -----------------------------------------------------------------------------

def test_ssc_fixed_cases_prev_lss_and_lps() -> None:
    """
    Spec:
      - If prev was LSS => current MUST be ESH
      - If prev was LPS => current MUST be OLS
    independent of attack detection on (i+1).
    """
    frame_t: FrameT = np.zeros((2048, 2), dtype=np.float64)

    next_attack = _next_frame_strong_attack(attack_left=True, attack_right=True)

    out1 = aac_ssc(frame_t, next_attack, "LSS")
    assert out1 == "ESH"

    out2 = aac_ssc(frame_t, next_attack, "LPS")
    assert out2 == "OLS"


# -----------------------------------------------------------------------------
# 2) Cases requiring next-frame ESH prediction (attack computation)
# -----------------------------------------------------------------------------

def test_prev_ols_next_not_esh_returns_ols() -> None:
    """
    If prev=OLS, current is:
      - LSS iff (i+1) is predicted ESH
      - else OLS
    Here: no attack => expect OLS.
    """
    frame_t: FrameT = np.zeros((2048, 2), dtype=np.float64)
    next_t = _next_frame_no_attack()

    out = aac_ssc(frame_t, next_t, "OLS")
    assert out == "OLS"


def test_prev_ols_next_esh_both_channels_returns_lss() -> None:
    """
    prev=OLS and next predicted ESH for both channels:
      per-channel: LSS, LSS
      merged: LSS
    """
    frame_t: FrameT = np.zeros((2048, 2), dtype=np.float64)
    next_t = _next_frame_strong_attack(attack_left=True, attack_right=True)

    out = aac_ssc(frame_t, next_t, "OLS")
    assert out == "LSS"


def test_prev_ols_next_esh_one_channel_returns_lss() -> None:
    """
    prev=OLS:
      - one channel predicts ESH => LSS
      - other channel predicts not ESH => OLS
    Merge table: OLS + LSS => LSS (either side).
    """
    frame_t: FrameT = np.zeros((2048, 2), dtype=np.float64)

    next1_t = _next_frame_strong_attack(attack_left=True, attack_right=False)
    out1 = aac_ssc(frame_t, next1_t, "OLS")
    assert out1 == "LSS"

    next2_t = _next_frame_strong_attack(attack_left=False, attack_right=True)
    out2 = aac_ssc(frame_t, next2_t, "OLS")
    assert out2 == "LSS"


def test_prev_esh_next_esh_both_channels_returns_esh() -> None:
    """
    prev=ESH and next predicted ESH for both channels:
      per-channel: ESH, ESH
      merged: ESH
    """
    frame_t: FrameT = np.zeros((2048, 2), dtype=np.float64)
    next_t = _next_frame_strong_attack(attack_left=True, attack_right=True)

    out = aac_ssc(frame_t, next_t, "ESH")
    assert out == "ESH"


def test_prev_esh_next_not_esh_both_channels_returns_lps() -> None:
    """
    prev=ESH and next not predicted ESH for both channels:
      per-channel: LPS, LPS
      merged: LPS
    """
    frame_t: FrameT = np.zeros((2048, 2), dtype=np.float64)
    next_t = _next_frame_no_attack()

    out = aac_ssc(frame_t, next_t, "ESH")
    assert out == "LPS"


def test_prev_esh_next_esh_one_channel_merged_is_esh() -> None:
    """
    prev=ESH:
      - one channel predicts ESH => ESH
      - other channel predicts not ESH => LPS
    Merge table: ESH + LPS => ESH (either side).
    """
    frame_t: FrameT = np.zeros((2048, 2), dtype=np.float64)

    next1_t = _next_frame_strong_attack(attack_left=True, attack_right=False)
    out1 = aac_ssc(frame_t, next1_t, "ESH")
    assert out1 == "ESH"

    next2_t = _next_frame_strong_attack(attack_left=False, attack_right=True)
    out2 = aac_ssc(frame_t, next2_t, "ESH")
    assert out2 == "ESH"


def test_threshold_s_must_exceed_1e_3() -> None:
    """
    Spec: next frame is predicted ESH only if:
      s[l] > 1e-3  AND  ds[l] > 10
    for some l in 1..7.

    This test checks the necessity of the s[l] threshold:
      - Create a frame with s[l] ~= 1e-4 < 1e-3 (single impulse with amp 0.01).
      - Expect: not classified as ESH -> for prev=OLS return OLS.
    """
    frame_t: FrameT = np.zeros((2048, 2), dtype=np.float64)
    next_t = _next_frame_below_s_threshold(left=True, right=True, impulse_amp=0.01)

    out = aac_ssc(frame_t, next_t, "OLS")
    assert out == "OLS"
