"""
src/models/semantic_model.py

Implements Module 2: Semantic Similarity Surrogate
"""

import numpy as np


class SemanticSurrogate:
    """
    Simulates the performance of the DeepSC model (Semantic similarity vs SNR and k_n).

    Paper Reference: Figure 2 visual insights.
    - Higher SNR and higher k_n strictly increase similarity.
    """

    def __init__(self, L: int = 12):
        self.L = L  # Average words per sentence

    def get_similarity(self, snr_linear: np.ndarray, k_n: int) -> np.ndarray:
        """
        Calculates the simulated semantic similarity.

        Args:
            snr_linear: Linear SNR values for users/channels
            k_n: Semantic representation symbols per word

        Returns:
            np.ndarray of same shape as snr_linear with semantic similarity
        """
        # Convert SNR to dB for the mapping (logistic curve matching Fig 2)
        snr_db = 10 * np.log10(snr_linear)

        # Logistic function parameters roughly tuned to match DeepSC curves
        # Base plateau increases with k_n
        max_sim = min(0.99, 0.6 + 0.04 * k_n)

        # Shift and scale depends on k_n
        # Higher k_n reaches max similarity at lower SNR
        shift = 10 - 0.5 * k_n
        scale = 2.0

        similarity = max_sim / (1 + np.exp(-(snr_db - shift) / scale))

        # Ensure bounds
        similarity = np.clip(similarity, 0.0, 1.0)

        return similarity
