classdef Optical_Multiplex < handle
    % Takes a cell array of signals
    % returns a total field signal
    % WDM spacing is given in wavelength plan OR via delta_F

    % The grid is stored in the output signal -> the demux will ideally
    % look this up and use this as the demux frequencies...

    % signal_cell = {Opt_sig_1, Opt_sig_2};
    % Opt_sig_wdm = Optical_Multiplex("fs_in",Opt_sig.fs,"fs_out",4*Opt_sig.fs,...
    %     "lambda_center",1310,"random_key",0,"filtype",1,"B",200e9,"delta_f",400e9).process(signal_cell);

    properties(Access=public)
        fs_in
        fs_out
        lambda_center
        delta_f
        random_key
        attenuation
        B
        mgauss
        filtype

        c = physconst('lightspeed')
        f_center
        f_T
        lambda_T
        df_T
    end

    methods (Access=public)
        function obj = Optical_Multiplex(options)
            %NAME Construct an instance of this class
            %   Detailed explanation goes here

            arguments
                options.fs_in
                options.fs_out
                options.lambda_center
                options.B = 200e9
                options.mgauss = 3
                options.filtype = 2

                options.delta_f = 0
                options.random_key
                options.attenuation = 0;
            end

            %
            fn = fieldnames(options);
            for n = 1:numel(fn)
                try
                    obj.(fn{n}) = options.(fn{n});
                end
            end


        end

        function signalclass_out = process(obj,signalclasses_in)

            % actual processing of the signal (steps 1. - 3.)
            signalclass_out = obj.process_(signalclasses_in);

            % append to logbook
            lbdesc = ['Opt. Mux. '];
            signalclass_out  = signalclass_out.logbookentry(lbdesc);

        end

        function data_out = process_(obj,data_in)
            %METHOD1 Summary of this method goes here
            %   Detailed explanation goes here
            arguments(Input)
                obj
                data_in cell
            end

            % assert(data_in{1}.fs == obj.fs_in,'Sampling rate');
            att           = 1/10^(obj.attenuation/10);
            N = numel(data_in);
            w = obj.fs_out/data_in{1}.fs;
            blocklen_in = length(data_in{1});
            blocklen_out = w*blocklen_in;
            freqaxis = linspace(-obj.fs_out/2, obj.fs_out/2, blocklen_out+1);
            obj.f_center = obj.c/(obj.lambda_center.*1e-9);

            if obj.random_key ~= 0
                res = freqaxis(2)-freqaxis(1);
                R = RandStream("twister","Seed",obj.random_key);
                laser_frequency_imperfection = res .* round(R.randn(N,1)*10); %in mutliples of the fft resolution, i.e. the distance between two freq. bins
            else
                laser_frequency_imperfection = zeros(blocklen_in,1);
            end

            obj.f_T = [];
            obj.df_T = [];
            polrots = [];
            for o = 1:N

                if obj.delta_f ~= 0
                    % user defined a channel spacing in GHz. Build plan
                    % left and right from zero
                    obj.df_T(o) = (-length(data_in)/2-0.5+o) .* obj.delta_f;
                    obj.df_T(o) = obj.df_T(o)+ laser_frequency_imperfection(o);

                    obj.f_T = [obj.f_T obj.f_center+obj.df_T(o)];
                else
                    %center frequencies of channels
                    obj.f_T = [obj.f_T obj.c/(data_in{o}.lambda)];
                    obj.lambda_T = [obj.lambda_T data_in{o}.lambda];

                    %difference between mid frequency of MUX and channels
                    obj.df_T   = [obj.df_T obj.f_center - obj.f_T(o)];
                end

                % adapt frequency shifts to match the FFT grid! Find nearest grid point
                [glitch(o),pos] = min(abs( freqaxis-obj.df_T(o) ));
                obj.df_T(o) = freqaxis(pos);

                polrots =  [polrots, data_in{o}.polrot];
            end

            obj.lambda_T = obj.c ./ (obj.f_center-obj.df_T);

            obj.B = 200e9; %200GHz
            faxis = linspace(-obj.fs_out/2,obj.fs_out/2, blocklen_out+1);%generates arow vector faxis of blocklen+1 points linearly spaced between and including -para.fs/2 and para.fs/2
            faxis = ifftshift(faxis(1:end-1));

            switch obj.filtype
                case 1
                    H =exp(-(faxis/obj.B).^(2*obj.mgauss)*log(2)*2^(2*obj.mgauss-1)).';
                case 2
                    H=zeros(length(faxis),1);
                    H(find(abs(faxis)<=obj.B/2))=1;
                case 3
                    H = 1;
            end

            x_envelopes = NaN([blocklen_out N]);
            y_envelopes = x_envelopes;

            % ---- OPTIMIZED: Pre-compute all LO phases as [blocklen_out × N] matrix ----
            time_idx = (0:blocklen_out-1).';  % [blocklen_out × 1]
            lo_phases = mod(2*pi * time_idx * obj.df_T / obj.fs_out, 2*pi);  % [blocklen_out × N]
            lo_all = cos(lo_phases) + 1i*sin(lo_phases);  % [blocklen_out × N]

            % Collect all resampled signals first (still requires loop due to cell array)
            x_signals = zeros(blocklen_out, N);
            y_signals = zeros(blocklen_out, N);

            for o = 1:N
                data_in_resampled = data_in{o}.resample("fs_out", obj.fs_out);
                x_signals(:, o) = data_in_resampled.signal(:, 1);
                y_signals(:, o) = data_in_resampled.signal(:, 2);
            end

            % ---- VECTORIZED: Batched FFT/IFFT for all channels ----
            % Apply filter to all channels at once
            x_filtered = ifft(fft(x_signals) .* H);  % [blocklen_out × N]
            y_filtered = ifft(fft(y_signals) .* H);  % [blocklen_out × N]

            % Apply attenuation and LO shift to all channels
            x_envelopes = att .* x_filtered .* lo_all;  % [blocklen_out × N]
            y_envelopes = att .* y_filtered .* lo_all;  % [blocklen_out × N]

            data_out = data_in_resampled;
            data_out.signal = [sum(x_envelopes,2), sum(y_envelopes,2)];
            data_out.lambda = obj.lambda_T;
            data_out.polrot = polrots;

        end

    end

    methods (Access=private)
        % Cant be seen from outside! So put all your functions here that can/
        % shall not be called from outside


    end
end