"""
语义通信模块 / Semantic communication module for Co-MADDPG.

实现基于 DeepSC 经验曲线的语义相似度 (SSim) 计算，以及语义通信用户的 QoE 计算。
Implements semantic similarity (SSim) computation based on empirical DeepSC curves,
and QoE calculation for semantic communication users.

作者/Author: Sisyphus-Junior
日期/Date: 2026-02-28
论文引用/Paper Reference: Co-MADDPG based Resource Allocation for Semantic Communication
依赖/Dependencies: numpy
"""

import numpy as np


class SemanticModule:
    """
    语义通信质量模块。
    Semantic communication quality module.

    根据平均 SNR 和压缩率，使用受 DeepSC 文献启发的经验拟合曲线计算语义相似度指数 (SSim)。
    Computes semantic similarity index (SSim) from average SNR and compression
    ratio, using empirical fitting curves inspired by DeepSC literature.

    Parameters
    ----------
    config : dict
        完整的配置字典（必须包含 "env" 部分，且具有 rho_max, rho_min, w1, w2 键）。
        Full configuration dictionary (must contain an "env" section 
        with keys "rho_max", "rho_min", "w1", "w2").
    """

    def __init__(self, config: dict) -> None:
        # 初始化语义参数 / Initialize semantic parameters
        self.config = config
        env = config["env"]
        # 最大压缩率 / Maximum compression ratio ρ_max
        self.rho_max = env.get("rho_max", 1.0)
        # 最小压缩率 / Minimum compression ratio ρ_min
        self.rho_min = env.get("rho_min", 0.05)
        # QoE 权重 1 / QoE Weight w1 (SSim weight)
        self.w1 = env.get("w1", 0.7)
        # QoE 权重 2 / QoE Weight w2 (Resource efficiency weight)
        self.w2 = env.get("w2", 0.3)

    @staticmethod
    def _a(rho: float) -> float:
        """经验曲线参数 a(ρ) = 0.8 / (ρ + 0.1)。 / Empirical curve parameter a(ρ) = 0.8 / (ρ + 0.1)."""
        return 0.8 / (rho + 0.1)

    @staticmethod
    def _b(rho: float) -> float:
        """经验曲线参数 b(ρ) = 0.6 + 0.2 * ρ。 / Empirical curve parameter b(ρ) = 0.6 + 0.2 * ρ."""
        return 0.6 + 0.2 * rho

    def compute_ssim(self, avg_snr, rho: float):
        """
        计算语义相似度指数 (SSim)。
        Compute semantic similarity index (SSim).

        φ(γ̄, ρ) = 1 - exp(-a(ρ) * γ̄^{b(ρ)}) (公式参考 SSim 章节)
        φ(γ̄, ρ) = 1 - exp(-a(ρ) * γ̄^{b(ρ)}) (Refer to SSim section in paper)

        Parameters
        ----------
        avg_snr : float or np.ndarray
            线性尺度的平均 SNR（非 dB）。
            Average SNR in linear scale (not dB).
        rho : float
            压缩率 ρ ∈ [ρ_min, ρ_max]。
            Compression ratio ρ ∈ [ρ_min, ρ_max].

        Returns
        -------
        float or np.ndarray
            [0, 1] 范围内的语义相似度。
            Semantic similarity in [0, 1].
        """
        # 防止 SNR 过小导致数值错误 / Avoid numerical errors with small SNR
        avg_snr = np.maximum(avg_snr, 1e-10)
        # 获取经验参数 a 和 b / Get empirical parameters a and b
        a = self._a(rho)
        b = self._b(rho)
        # 计算 SSim 公式 / Compute SSim formula
        return 1.0 - np.exp(-a * np.power(avg_snr, b))

    def compute_avg_snr(self, snr_per_subcarrier: np.ndarray,
                        allocation_mask: np.ndarray) -> float:
        """
        计算已分配子载波上的平均 SNR。
        Compute average SNR over allocated subcarriers.

        Parameters
        ----------
        snr_per_subcarrier : np.ndarray
            所有子载波的 SNR 值（线性尺度）。
            SNR values for all subcarriers (linear scale).
        allocation_mask : np.ndarray
            指示已分配子载波的二进制掩码。
            Binary mask indicating allocated subcarriers.

        Returns
        -------
        float
            已分配子载波的平均 SNR；若未分配则返回 0。
            Mean SNR over allocated subcarriers; 0 if none allocated.
        """
        # 提取已分配子载波的 SNR / Extract SNR for allocated subcarriers
        allocated = snr_per_subcarrier[allocation_mask > 0]
        # 如果没有子载波被分配 / If no subcarriers are allocated
        if len(allocated) == 0:
            return 0.0
        # 返回平均值 / Return mean value
        return float(np.mean(allocated))

    def compute_semantic_qoe(self, ssim: float, rho: float,
                             w1: float = None, w2: float = None,
                             rho_max: float = None) -> float:
        """
        计算语义通信用户的 QoE。
        Compute QoE for a semantic communication user.

        QoE_s = w1 * SSim + w2 * (1 - ρ / ρ_max) (公式参考 QoE_s)
        QoE_s = w1 * SSim + w2 * (1 - ρ / ρ_max) (Refer to QoE_s formula)

        Parameters
        ----------
        ssim : float
            [0, 1] 范围内的语义相似度指数。
            Semantic similarity index in [0, 1].
        rho : float
            使用的压缩率。
            Compression ratio used.
        w1, w2 : float, optional
            权重（默认为配置中的实例值）。
            Weights (defaults to instance values from config).
        rho_max : float, optional
            最大压缩率（默认为配置中的值）。
            Maximum compression ratio (default from config).

        Returns
        -------
        float
            [0, 1] 范围内的 QoE 值。
            QoE value in [0, 1].
        """
        # 使用默认值或输入值 / Use default or input values
        if w1 is None:
            w1 = self.w1
        if w2 is None:
            w2 = self.w2
        if rho_max is None:
            rho_max = self.rho_max
        # 计算语义 QoE / Calculate semantic QoE
        return float(w1 * ssim + w2 * (1.0 - rho / rho_max))