"""
src/models/allocator.py

Implements Module 3: Resource Allocation Optimizer
"""

import numpy as np
from scipy.optimize import linear_sum_assignment
from typing import Tuple, List, Optional
from src.models.semantic_model import SemanticSurrogate
from src.models.environment import EnvironmentConfig


class ResourceAllocator:
    """
    Implements the optimization algorithms from the paper.
    - Exhaustive search for optimal k_n
    - Hungarian algorithm for bipartite channel assignment
    """

    def __init__(
        self,
        surrogate: SemanticSurrogate,
        env_config: EnvironmentConfig,
        K_max: int = 8,
    ):
        self.surrogate = surrogate
        self.env_config = env_config
        self.K_max = K_max
        self.B = env_config.bandwidth

    def calculate_sse(self, snr_linear: np.ndarray, k_n: int) -> np.ndarray:
        """
        Calculate Semantic Spectral Efficiency (S-SE).

        S-SE = (L * xi) / (k_n * L) * (W / B) = xi * W / (k_n * B)
        where W is the channel capacity (Shannon rate).

        Args:
            snr_linear: Linear SNR
            k_n: Semantic symbols per word

        Returns:
            S-SE matrix of shape (N, M)
        """
        # Channel capacity W = B * log2(1 + SNR)
        w = self.B * np.log2(1 + snr_linear)

        # Semantic similarity xi
        xi = self.surrogate.get_similarity(snr_linear, k_n)

        # S-SE = xi * W / (k_n * L) * (L / B) = xi * W / (k_n * B)
        sse = (xi * w) / (k_n * self.B)

        return sse

    def optimize_semantic_aware(
        self, snr_linear: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray, float]:
        """
        Runs the semantic-aware resource allocation algorithm.

        Args:
            snr_linear: Linear SNR matrix of shape (N, M)

        Returns:
            Tuple of (optimal_k_n array, assignment_matrix, total_sse)
        """
        N, M = snr_linear.shape

        # 1. Exhaustive search for optimal k_n for each user-channel pair
        max_sse_matrix = np.zeros((N, M))
        optimal_k_matrix = np.zeros((N, M), dtype=int)

        for k in range(1, self.K_max + 1):
            sse_k = self.calculate_sse(snr_linear, k)

            # Update where this k provides better S-SE
            better_mask = sse_k > max_sse_matrix
            max_sse_matrix[better_mask] = sse_k[better_mask]
            optimal_k_matrix[better_mask] = k

        # 2. Hungarian algorithm for bipartite matching
        # scipy's linear_sum_assignment finds minimum weight matching
        # We want maximum weight, so we use negative S-SE
        cost_matrix = -max_sse_matrix

        row_ind, col_ind = linear_sum_assignment(cost_matrix)

        # Create assignment matrix
        assignment = np.zeros((N, M), dtype=int)
        assignment[row_ind, col_ind] = 1

        # Collect optimal k_n for assigned channels
        optimal_k = np.zeros(N, dtype=int)
        for i, j in zip(row_ind, col_ind):
            optimal_k[i] = optimal_k_matrix[i, j]

        # Calculate total S-SE
        total_sse = np.sum(max_sse_matrix[row_ind, col_ind])

        return optimal_k, assignment, total_sse

    def evaluate_fixed_k(self, snr_linear: np.ndarray, k_n: int) -> float:
        """
        Evaluate performance with a fixed k_n using Hungarian matching.

        Args:
            snr_linear: Linear SNR matrix (N, M)
            k_n: Fixed k_n to use

        Returns:
            Total S-SE
        """
        N, M = snr_linear.shape
        sse_matrix = self.calculate_sse(snr_linear, k_n)

        cost_matrix = -sse_matrix
        row_ind, col_ind = linear_sum_assignment(cost_matrix)

        total_sse = np.sum(sse_matrix[row_ind, col_ind])
        return total_sse
