import numpy as np

class OpticalNoiseModel:
    def __init__(self, bit_depth=8, snr_db=30):
        self.bit_depth = bit_depth
        self.snr_db = snr_db

    def apply_noise(self, signal):
        """
        Applies:
        1. Shot Noise (Poisson) / Thermal Noise (Gaussian)
        2. ADC/DAC Quantization
        """
        # 1. Gaussian Noise based on SNR
        # Signal Power
        sig_power = np.mean(signal ** 2)
        sig_power_db = 10 * np.log10(sig_power + 1e-12)
        noise_power_db = sig_power_db - self.snr_db
        noise_power = 10 ** (noise_power_db / 10)
        noise_std = np.sqrt(noise_power)

        noise = np.random.normal(0, noise_std, signal.shape)
        noisy_signal = signal + noise

        # 2. Quantization (DAC/ADC)
        # Normalize to [0, 1] for quantization simulation
        min_val = np.min(noisy_signal)
        max_val = np.max(noisy_signal)
        r = max_val - min_val + 1e-9

        norm = (noisy_signal - min_val) / r
        levels = 2**self.bit_depth
        quantized = np.round(norm * (levels - 1)) / (levels - 1)

        # Denormalize
        out = quantized * r + min_val
        return out

class OpticalDispersionModel:
    def __init__(self, crosstalk_db=-40, dispersion_ps_nm_km=17):
        self.crosstalk_coeff = 10**(crosstalk_db/10)

    def apply_dispersion(self, signal_matrix):
        """
        Simulates Crosstalk (leaking energy to adjacent channels/elements).
        """
        # Simple neighbour crosstalk convolution
        kernel = np.array([[0, self.crosstalk_coeff, 0],
                           [self.crosstalk_coeff, 1.0 - 4*self.crosstalk_coeff, self.crosstalk_coeff],
                           [0, self.crosstalk_coeff, 0]])

        # We assume signal_matrix is 2D spatial representation
        from scipy.signal import convolve2d
        return convolve2d(signal_matrix, kernel, mode='same')
