"""
tests/test_allocator.py
"""

import pytest
import numpy as np
from src.models.environment import EnvironmentConfig
from src.models.semantic_model import SemanticSurrogate
from src.models.allocator import ResourceAllocator


def test_allocator_initialization():
    env_config = EnvironmentConfig(num_users=3, num_channels=4, bandwidth=1e6)
    surrogate = SemanticSurrogate(L=12)
    allocator = ResourceAllocator(surrogate, env_config, K_max=8)

    assert allocator.K_max == 8


def test_optimize_semantic_aware():
    env_config = EnvironmentConfig(num_users=2, num_channels=3, bandwidth=1e6)
    surrogate = SemanticSurrogate(L=12)
    allocator = ResourceAllocator(surrogate, env_config, K_max=4)

    # Mock SNR values
    snr_linear = np.array([[10.0, 100.0, 1000.0], [50.0, 20.0, 500.0]])

    optimal_k, assignment, total_sse = allocator.optimize_semantic_aware(snr_linear)

    assert optimal_k.shape == (2,)
    assert assignment.shape == (2, 3)
    assert np.all(assignment.sum(axis=1) == 1)  # Each user gets 1 channel
    assert np.all(assignment.sum(axis=0) <= 1)  # Each channel used by max 1 user


def test_evaluate_fixed_k():
    env_config = EnvironmentConfig(num_users=2, num_channels=2, bandwidth=1e6)
    surrogate = SemanticSurrogate(L=12)
    allocator = ResourceAllocator(surrogate, env_config, K_max=4)

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

    total_sse = allocator.evaluate_fixed_k(snr_linear, k_n=4)
    assert total_sse > 0