import torch.nn as nn
from torch.distributions.normal import Normal
import math


class GaussianChannel(nn.Module):
    """
    Simulated communication channel that assumes a Gaussian noise model for
    taking in account interference.
    """
    def __init__(self):
        super().__init__()

        # Initialize channel parameters
        sys_rate = 32e9
        r = 0.05
        dispersion = 16.48e-6
        B_2 = dispersion
        non_linear_index = 1.3e3
        gam = non_linear_index
        loss = 10**20
        alpha = loss
        span_count = 20
        N_s = span_count
        span_length = 10e5  # (km)
        L_s = span_length
        noise_figure = 10 ** 0.5  # (dB)
        h = 6.6261e-34
        v = 299792458
        B_WDM = sys_rate * (1 + r)
        B_N = 0.1

        P_ASE_1 = h * v * B_N * (loss * span_length * noise_figure - 1)
        P_ASE = P_ASE_1 * span_count
        L_eff = 1 - math.exp(-loss * span_length) / 2 / alpha

        eps = 0.3 * math.log(
            1 + (6 / L_s) * (
                L_eff / math.asinh(
                    (math.pi ** 2)
                    / 3
                    * B_2
                    * L_eff
                    * (B_WDM ** 2)
                )
            )
        )

        b = P_ASE_1 / (
            2
            * (N_s ** eps)
            * B_N
            * (gam ** 2)
            * L_eff * math.asinh(
                (math.pi ** 2) / 3
                * B_2
                * L_eff
                * (B_WDM ** 2)
            )
        )

        P_ch = sys_rate * (((27 * math.pi * B_2 / 16) * b) ** (1 / 3))
        OSNR = (2 * P_ch / 3 / P_ASE)

        OSNR_dB = 10 * math.log10(OSNR)

        p_N_dB = -OSNR_dB
        p_N_watt = 10**(p_N_dB/10)

        self.noise_std = math.sqrt(p_N_watt * 5000)

    def get_noise(self, rows):
        """
        Generate Gaussian random noise according to the channel’s parameters.

        :param rows: Number of noise vectors to generate.
        :return: Matrix of shape `rows` × 2.
        """
        return Normal(0, self.noise_std).sample((rows, 2))

    def forward(self, x):
        return x + self.get_noise(len(x))