"""
src/models/environment.py

Implements Module 1: Environment & Channel Simulator
"""

import numpy as np
from typing import NamedTuple, Tuple


class EnvironmentConfig(NamedTuple):
    """Configuration for the channel simulator."""

    num_users: int = 10
    num_channels: int = 10
    bandwidth: float = 1e6  # Hz, 1 MHz per channel
    radius: float = 0.5  # km, cell radius
    shadow_fading_std: float = 6.0  # dB
    noise_psd_dbm: float = -174.0  # dBm/Hz


class ChannelSimulator:
    """
    Simulates the wireless channel environment.

    Paper Reference:
    - Pathloss: 128.1 + 37.6 lg[d(km)] dB
    - Shadow fading: 6 dB
    """

    def __init__(self, config: EnvironmentConfig):
        self.config = config

    def _calculate_pathloss(self, distances: np.ndarray) -> np.ndarray:
        """Calculate pathloss for given distances in km."""
        return 128.1 + 37.6 * np.log10(distances)

    def generate_channels(
        self, transmit_power_dbm: float
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Generate channel conditions (SNR) for all users and channels.

        Args:
            transmit_power_dbm: Transmit power in dBm

        Returns:
            Tuple of (snr_db, snr_linear) with shape (num_users, num_channels)
        """
        N, M = self.config.num_users, self.config.num_channels

        # 1. Distances (randomly distributed between 0.05km and cell radius)
        min_dist = 0.05
        distances = np.random.uniform(min_dist, self.config.radius, size=N)

        # 2. Pathloss
        path_loss_db = self._calculate_pathloss(distances)

        # 3. Shadow fading
        shadowing_db = np.random.normal(0, self.config.shadow_fading_std, size=N)

        # 4. Total large scale fading (dB)
        large_scale_db = path_loss_db + shadowing_db

        # 5. Rayleigh fading (small scale)
        # Power of Rayleigh follows exponential distribution (mean=1)
        small_scale_power = np.random.exponential(1.0, size=(N, M))
        small_scale_db = 10 * np.log10(small_scale_power)

        # 6. Noise power
        # Noise = PSD (dBm/Hz) + 10*log10(BW)
        noise_power_dbm = self.config.noise_psd_dbm + 10 * np.log10(
            self.config.bandwidth
        )

        # 7. Calculate SNR
        # SNR(dB) = Pt(dBm) - LargeScale(dB) + SmallScale(dB) - Noise(dBm)
        snr_db = np.zeros((N, M))
        for n in range(N):
            snr_db[n, :] = (
                transmit_power_dbm
                - large_scale_db[n]
                + small_scale_db[n, :]
                - noise_power_dbm
            )

        snr_linear = 10 ** (snr_db / 10)

        return snr_db, snr_linear