"""
src/models/baselines.py

Implements Module 4: Transform Method & Baselines
"""

import numpy as np
from typing import Tuple


class BaselineModels:
    """
    Implements 4G/5G baselines and Ideal Shannon Limit using the transform method.
    """

    def __init__(self, mu: float = 19.0, L: int = 12):
        self.mu = mu  # Transforming factor (bits/word)
        self.L = L  # Average words per sentence

    def _cqi_mapping(self, snr_db: np.ndarray, generation: str = "5G") -> np.ndarray:
        """
        Approximate 3GPP CQI mapping from SNR to Spectral Efficiency (bps/Hz)
        """
        if generation == "5G":
            # 5G supports higher modulation schemes (e.g. 256 QAM)
            max_se = 7.4  # up to ~7.4 bps/Hz
            shift = 15.0
            scale = 4.0
        else:  # 4G
            # 4G supports up to 64 QAM typically in this context
            max_se = 4.8  # up to ~4.8 bps/Hz
            shift = 18.0
            scale = 5.0

        # Sigmoid approximation of discrete CQI steps
        se = max_se / (1 + np.exp(-(snr_db - shift) / scale))

        # Ensure minimum SE for low SNR
        se = np.maximum(se, 0.1)

        return se

    def calculate_baseline_sse(
        self, snr_linear: np.ndarray, model_type: str
    ) -> np.ndarray:
        """
        Calculate equivalent Semantic Spectral Efficiency for baselines.
        """
        snr_db = 10 * np.log10(np.maximum(snr_linear, 1e-10))

        if model_type.lower() == "ideal":
            # Ideal Shannon capacity W / B = log2(1 + SNR)
            spectral_efficiency = np.log2(1 + snr_linear)
        elif model_type.upper() == "5G":
            spectral_efficiency = self._cqi_mapping(snr_db, "5G")
        elif model_type.upper() == "4G":
            spectral_efficiency = self._cqi_mapping(snr_db, "4G")
        else:
            raise ValueError(f"Unknown model_type: {model_type}")

        # Transform method: equivalent S-SE = SE / mu
        # Because S-SE is measured in words/sec/Hz, and SE is bits/sec/Hz
        # mu is bits/word
        sse = spectral_efficiency / self.mu

        return sse