"""
tests/test_baselines.py
"""

import pytest
import numpy as np
from src.models.baselines import BaselineModels


def test_baselines_initialization():
    models = BaselineModels(mu=19.0, L=12)
    assert models.mu == 19.0


def test_calculate_baseline_sse():
    models = BaselineModels(mu=19.0, L=12)

    snr_linear = np.array([10.0, 100.0, 1000.0])

    sse_ideal = models.calculate_baseline_sse(snr_linear, "ideal")
    sse_5g = models.calculate_baseline_sse(snr_linear, "5G")
    sse_4g = models.calculate_baseline_sse(snr_linear, "4G")

    assert len(sse_ideal) == 3
    assert np.all(sse_ideal > sse_5g)
    assert np.all(sse_5g > sse_4g)
