classdef Optical_Demultiplex < handle
    % Dual-Polarization optical demultiplexer
    % - Input:  total-field signal
    % - Output: single-channel dual-pol signal objects in cell array
    %
    % Notes:
    % Opt_sig_wdm_demux = Optical_Demultiplex("attenuation",0,"B",200e9,"filtype",1,"fs_out",Opt_sig_wdm_rx.fs/4,"fs_in",Opt_sig_wdm_rx.fs,"lambda_center",1310).process(Opt_sig_wdm_rx);
    % Opt_sig_wdm_demux{1}.spectrum("fignum",1100,"displayname",'bla','normalizeTo0dB',0,'max_num_lines',4);
    % Opt_sig_wdm_demux{2}.spectrum("fignum",1100,"displayname",'bla','normalizeTo0dB',0,'max_num_lines',4);

    properties (Access=public)
        fs_in              % [Hz]  (optional; inferred from data_in.fs if omitted)
        fs_out             % [Hz]
        lambda_center      % [nm] center wavelength of the WDM grid
        wavelengthplan     % [nm]
        attenuation = 0    % [dB] insertion loss
        filtype = 1        % 1=Gaussian, 2=Rectangle, 3=No filter
        B       = 200e9   % [Hz] 3 dB bandwidth (Gaussian) or width (Rect)
        mgauss  = 3        % Gaussian order (multiple of 1/2)

        % Derived/utility
        c = physconst('lightspeed') % [m/s]
    end

    methods (Access=public)
        function obj = Optical_Demultiplex(options)
            arguments
                options.fs_in           = []
                options.fs_out
                options.lambda_center
                options.wavelengthplan
                options.attenuation     = 0
                options.filtype         = 1
                options.B               = 2.5e10
                options.mgauss          = 3
            end
            fn = fieldnames(options);
            for n = 1:numel(fn)
                try obj.(fn{n}) = options.(fn{n}); end
            end
        end

        function signalclasses_out = process(obj, signalclass_in)

            % ---- Infer wavelength: either given or from input total signal 
            if isempty(obj.wavelengthplan)
                obj.wavelengthplan = signalclass_in.lambda; %meter
            else
                if all(500e-9 < obj.wavelengthplan) && all(obj.wavelengthplan < 1500e-9) %check if given in nm
                    obj.wavelengthplan = obj.wavelengthplan.*1e-9;
                end
            end

            % ---- Infer input sampling rates
            if isempty(obj.fs_in)
                assert(isprop(signalclass_in,'fs') && ~isempty(signalclass_in.fs), ...
                    'Dual_Pol_Demultiplexer: data_in.fs missing and options.fs_in not provided.');
                obj.fs_in = signalclass_in.fs;
            end

            % Runs demultiplexing in one go and appends a logbook entry.
            [x_envelopes,y_envelopes] = obj.process_(signalclass_in.signal);

            for n = 1:min(size(x_envelopes))
                signalclasses_out{n}        = signalclass_in;
                signalclasses_out{n}.signal = [x_envelopes(:,n), y_envelopes(:,n)];
                signalclasses_out{n}        = signalclasses_out{n}.resample("fs_in",obj.fs_in,"fs_out",obj.fs_out);
                signalclasses_out{n}.lambda = obj.wavelengthplan(n);
                lbdesc = ['Opt. Demux ', num2str( obj.wavelengthplan(n)),' nm'];
                signalclasses_out{n} = signalclasses_out{n}.logbookentry(lbdesc);
            end

        end

        function [x_envelopes,y_envelopes] = process_(obj, signal_in)
            % Core demux:
            %  - frequency translate target channel to baseband
            %  - apply optical filter H
            %  - resample to fs_out
            arguments (Input)
                obj
                signal_in
            end
            
            w = obj.fs_out ./ obj.fs_in ;
            blocklen_in = length(signal_in);
            blocklen_out = w*blocklen_in;

            att = 1/10^(obj.attenuation/10);
            faxis=linspace( -obj.fs_in/2 , obj.fs_in/2 , blocklen_in+1 );
            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
                    %all zero filter
                    H=zeros(1,length(faxis)).';
                    %set filter = 1 inside bandwidth -B/2 <-> B/2
                    H(abs(faxis)<=obj.B/2)=1;
                case 3
                    H = 1;
            end

            f_mid = obj.c/(obj.lambda_center*1e-9); % center frequency of WDM grid [Hz]

            f_channels = obj.c./(obj.wavelengthplan) ;

            N = numel(f_channels);

            df_T   = f_mid - f_channels;

            pha = mod(-2*pi*(0:blocklen_in-1).'.*df_T/obj.fs_in, 2*pi);
            lo  =  cos(pha)+1i*sin(pha);

            % x_envelopes = ifft(fft(att.*signal_in(:,1).*lo).*H);
            % y_envelopes = ifft(fft(att.*signal_in(:,2).*lo).*H);

            N = size(lo,1);
            C = size(lo,2);

            x_envelopes = zeros(N, C, 'like', signal_in);
            y_envelopes = zeros(N, C, 'like', signal_in);

            s1 = signal_in(:,1);
            s2 = signal_in(:,2);

            % Reusable work buffers (avoid reallocations)
            wrk_time = zeros(N,1, 'like', signal_in);
            wrk_freq = zeros(N,1, 'like', signal_in);

            for c = 1:C
                % ---- X branch ----
                wrk_time(:) = att .* s1 .* lo(:,c);   % N×1
                wrk_freq(:) = fft(wrk_time);          % N×1
                wrk_freq(:) = wrk_freq .* H;          % N×1
                x_envelopes(:,c) = ifft(wrk_freq);     % N×1

                % ---- Y branch ----
                wrk_time(:) = att .* s2 .* lo(:,c);
                wrk_freq(:) = fft(wrk_time);
                wrk_freq(:) = wrk_freq .* H;
                y_envelopes(:,c) = ifft(wrk_freq);
            end
 
        end
    end
end