import numpy as np
import scipy.constants as Constant
import amp_mode 
import gain_mode
import nase_mode

class Amplifier:
    def __init__(self, options=None):
        self.amp_mode = options.amp_mode if options and hasattr(options, 'amp_mode') else amp_mode.ideal_no_noise
        self.gain_mode = options.gain_mode if options and hasattr(options, 'gain_mode') else gain_mode.output_power
        self.nase_mode = options.nase_mode if options and hasattr(options, 'nase_mode') else nase_mode.pass_ase
        self.amplification_db = options.amplification_db if options and hasattr(options, 'amplification_db') else 0
        self.noifig = options.noifig if options and hasattr(options, 'noifig') else 0
        self.fsimu = options.fsimu if options and hasattr(options, 'fsimu') else None

    def process(self, signalclass_in):
        signalclass_in = self.process_(signalclass_in)
        lbdesc = 'Amp '
        signalclass_in = signalclass_in.logbookentry(lbdesc)
        signalclass_out = signalclass_in
        return signalclass_out

    def process_(self, X_in):
        a_lin = self.calculateGain(X_in.signal)
        X_in.signal = a_lin * X_in.signal

        if isinstance(X_in, Opticalsignal):
            if self.amp_mode == amp_mode.ideal_no_noise:
                X_in.nase = self.onlyAmplifyAse(X_in.nase, a_lin)
            elif self.amp_mode == amp_mode.edfa_increase_nase:
                X_in.nase = self.increaseAse(X_in.nase, a_lin, self.noifig, X_in.lambda_)
            elif self.amp_mode == amp_mode.edfa_replace_nase:
                X_in.nase = self.replaceAse(X_in.nase, a_lin, self.noifig, X_in.lambda_)

            if self.nase_mode == nase_mode.generate_ase:
                nase = X_in.nase
                fs = X_in.fs
                dimension = X_in.signal.shape
                nase_numeric = self.generateAseNoise(nase, fs, dimension)
                X_in.signal, osnr = self.applyAseToSignal(X_in.signal, nase_numeric)
                X_in.nase = 0
            elif self.nase_mode == nase_mode.pass_ase:
                X_in.nase = X_in.nase

        X_out = X_in
        return X_out

    def calculateGain(self, xin):
        if self.gain_mode == gain_mode.output_power:
            pow_in = np.mean(np.abs(xin) ** 2)
            pow_out = 10 ** (self.amplification_db / 10 - 3)
            a_lin = np.sqrt(pow_out / pow_in)
        elif self.gain_mode == gain_mode.gain:
            a_lin = 10 ** (self.amplification_db / 20)
        return a_lin

    def onlyAmplifyAse(self, nase_old, a_lin):
        nase_amped = a_lin ** 2 * nase_old
        return nase_amped

    def calculateAseFromThisAmp(self, a_lin, noisefig, lambda_):
        h = Constant.Planck
        c = Constant.speed_of_light
        nase_new = 0.5 * 10 ** (noisefig / 10) * h * c / lambda_ * (a_lin ** 2 - 1)
        return nase_new

    def increaseAse(self, nase_old, a_lin, noisefig, lambda_):
        nase_fromThisAmp = self.calculateAseFromThisAmp(a_lin, noisefig, lambda_)
        nase_old = a_lin ** 2 * nase_old
        ase_increased = nase_old + nase_fromThisAmp
        return ase_increased

    def replaceAse(self, nase_old, a_lin, noisefig, lambda_):
        nase_fromThisAmp = self.calculateAseFromThisAmp(a_lin, noisefig, lambda_)
        nase_new = nase_fromThisAmp
        return nase_new

    def generateAseNoise(self, nase, fs, dimension):
        nase_numeric = (np.random.randn(*dimension) + 1j * np.random.randn(*dimension)) * np.sqrt(nase / 2 * fs)
        return nase_numeric

    def applyAseToSignal(self, opticalsignal, nase_numeric):
        noisy_sig = opticalsignal + nase_numeric
        osnr = pow(10 * np.log10(np.mean(np.abs(opticalsignal) ** 2) / np.mean(np.abs(nase_numeric) ** 2)))
        print(osnr)
        return noisy_sig, osnr