%%% discoRX.m --- 
%% 
%% Filename: discoRX.m
%% Description: Matlab/Octave model for illustrating discontinuous
%%               symbol-synchronized RX processing 
%% Author: Juha Yli-Kaakinen
%% Maintainer: 
%% Created: Wed Aug 14 13:24:52 2019
%% Version: 1.0
%% Last-Updated: Tue Nov 10 17:32:00 2020
%%           By: Juha Yli-Kaakinen
%%     Update #: 1885
%% Keywords: 
%% References: 
%% [1] Juha Yli-Kaakinen, AlaaEddin Loulou, Toni Levanen, Kari Pajukoski, Arto Palin, 
%%     Markku Renfors, and Mikko Valkama, "Frequency-domain signal processing for 
%%     spectrally-enhanced CP-OFDM waveforms in 5G new radio," submitted IEEE
%%     Transactions on Wireless Communications, 2020.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function discoRX

isOctave = exist('OCTAVE_VERSION', 'builtin') ~= 0; 

vec = @(x) x(:); % Vectorization function

ran_seed = 10;
if isOctave
  rand('seed', ran_seed);
  randn('seed', ran_seed);
  
else
  rng(ran_seed); 
 
end
hextodec = @(x) hex2dec(reshape(x,2,3).').'/255;

i0 = 1;                       % Index of the 1st element of the vector

figure(1), clf

Lact = 1*12;                  % Number of active subcarriers (one PRB)
bOFDM = 14;                   % Number of OFDM symbols 

% Input data on active subcarries. Shift around DC
Xact = 1/sqrt(2)*(sign(randn(Lact, bOFDM)) + 1j*sign(randn(Lact, bOFDM)));  

Bin0 = 63;                    % Center bin of the subband
N = 1024;                     % Long FC SFB transform (IFFT) size  
OLF = 1/2;                    % Overlap factor in FC processing

assert(OLF == 1/2, 'In this model, the overlap factor has to be equal to half')

k = 0;
%% ======================================================================
%% TX PROCESSING: Plain CP-OFDM processing
%% ======================================================================

L = N;                       % Short FC transform (FFT) size
I = N/L;                     % Interpolation factor [Eq. (4)]
L_OFDM = L;                  % Low-rate OFDM transform size
N_OFDM = I*L_OFDM;           % High-rate OFDM transform size
  
% Determine the FC processing block partition sizes [based on Eqs. (2) and (3)]
[Ls, Ll, Ns, Nl] = FCpartitioning(L, N, OLF);

% Determine the CP lengths (resempling the LTE and 5G-NR CP lengths for 15 kHz SCS)
[L_CP, N_CP] = CPlenghts(L_OFDM, N_OFDM, bOFDM);

% Number of samples to be processed 
nSamp_highrate = sum(N_CP, 1) + N_OFDM.*bOFDM; % Samples on the output rate

% Map Lact active subcarries to L_OFDM bins of the OFDM IFFT
X = circshift([fftshift(Xact, 1); zeros(L_OFDM(1)-Lact, bOFDM)], ...
	      [-floor(Lact/2) 0]); 

SCS = 15e3;                   % OFDM subcarrier spacing (SCS)
Fs = SCS*N_OFDM;              % Output sample rate
BS_FC = SCS;                  % Bin spacing in FC processing

% OFDM modulation, CP inclusion, and zero padding to length 3/2L
[Z, x_cp] = OFDMmod(X, Lact, bOFDM, L_OFDM, N_CP, N, Ns, Nl);  

prePad = Nl-N_CP(1); postPad = L-Ls; 

% Modulate the CP-OFDM waveforms to desired center frequency
w_m = x_cp.*exp(j*2*pi*Bin0.*1/N_OFDM*((0:nSamp_highrate-1)+prePad).');

% Active subcarries on TX side
iAct_TX = mod([1:Lact]-floor(Lact/2)-1, N_OFDM)+1;

for iModel = 1:3 
  
  %% ======================================================================
  %% RX PROCESSING
  %% ======================================================================
  switch iModel
   case 1, % 'Continuous FC-OFDM RX'
    L = 128; % Short FC transform (FFT) size

    lineColor = hextodec('0072bd'); lineStyle = 'o-';
    legStr = ['Continuous FC-F-OFDM RX, L = L\_OFDM = ' int2str(L)];
    
   case 2, % 'Discontinuous FC-OFDM RX'
    L = 128; % Short FC transform (FFT) size

    lineColor = hextodec('77ac30'); lineStyle = 'o-';
    legStr = ['Discontinuous FC-F-OFDM RX, L = L\_OFDM = ' int2str(L)];
    
   case 3, % 'Extrapolating FC-OFDM RX'
    L = 48; % Short FC transform (FFT) size

    lineColor = hextodec('a2142f'); lineStyle = 'x-';
    legStr = ['Extrapolating FC-F-OFDM RX, L = L\_OFDM = ' int2str(L)];
    
   otherwise
    error('Unsupported')
    
  end

  L_OFDM = L;                 % Low-rate OFDM transform size
  N_OFDM = N/L*L_OFDM;        % High-rate OFDM transform size
  
  % Determine the CP lengths (resempling the 5G-NR CP lengths for 15 kHz SCS)
  [L_CP, N_CP] = CPlenghts(L_OFDM, N_OFDM, bOFDM);
  
  disp([legStr ':'])
  disp(['High-rate IFFT, N_OFDM = [' int2str(N_OFDM.') ']'])
  disp([' Low-rate IFFT, L_OFDM = [' int2str(L_OFDM.') ']'])
  disp(['    High-rate CP, N_CP = [' sprintf('%2i  ', N_CP) ']'])
  disp(['     Low-rate CP, L_CP = [' sprintf('%2i  ', L_CP) ']'])

  % Determine the FC processing block partition sizes [based on Eqs. (2) and (3)]
  [Ls, Ll, Ns, Nl] = FCpartitioning(L, N, OLF);
  
  % Phase rotation [Eqs. (9b), (14b), and (14c)]
  theta = @(r) exp(j*2*pi*r*Bin0*Ls/L); 
  phi   = @(n) exp(j*2*pi*Bin0*sum(N_CP((1:n)+i0))/N);

  % Matrix 
  W_N  = 1/sqrt(L)*fft(eye(N)); % DFT matrix of size N x N
  W_L  = 1/sqrt(N)*fft(eye(L)); % DFT matrix of size L x L
  WI_N = sqrt(N)*ifft(eye(N));  % IFFT matrix of size N x N
  WI_L = sqrt(L)*ifft(eye(L));  % IFFT matrix of size L x L
  P    = fftshift(eye(L), 1);   % FFT shift matrix 
  
  % Time-domain analysis window used for overlap-and-add (OSA) processing [Eq. (10a)]
  A_mr = sparse(Ll+(1:Ls), Ll+(1:Ls), ones(1, Ls), L, L);
  
  % Mapping matrix maps L input bins to N output bins [Eq. (9a)]
  M_LN = @(theta) sparse(mod(ceil(-L/2:L/2-1)+Bin0, N)+1, 1:L, ...
			 theta*ones(1, L), N, L);
  
  % Frequency-domain window function: Ones on the passband and zeros on the stopband 
  % as well as 'nTB' transition-band weights on both sides of the passband
  nTB = 12;                   % Number of transition band weights in FC processing
  excessBW = 3;               % Passband extension in frequency-domain bins
  win = genWin(nTB, excessBW, Lact, 1, L);
  
  switch iModel
   case {2,3}
    %% DISCONTINUOUS SYMBOL-SYNCHRONIZED FC RX PROCESSING

    % Length of the output vector (Eq. 16d)
    Q = N/2 + bOFDM*N_OFDM + sum(N_CP(2:bOFDM));

    w_hat_m = [zeros(Nl-N_CP(1), 1); w_m(1:nSamp_highrate); zeros(N-Ns-Nl, 1)];
    
    % Segmentation of the received symbols [Eq. (20a)]
    for n = 0:bOFDM-1
      % Starting index of the nth symbol [Eq. (16b)]
      sigma_n = n*N_OFDM(1) + sum(N_CP((1:n)+i0));

      % Alignment [Eq. (16c)]
      Gamma = [ 
	  zeros(sigma_n, 3/2*N)
	  eye(3/2*N)
	  zeros(Q-3/2*N-sigma_n, 3/2*N)
	      ];
      
      % Overlap-and-save
      w_hat_mn(:, n+i0) = Gamma.'*w_hat_m;
    
    end % END for n
            
    R_Nm0 = [eye(N); zeros(N/2, N)]; 
    R_Nm1 = [zeros(N/2, N); eye(N)];

    z_hat_0 = zeros(L, bOFDM);
    z_hat_1 = zeros(L, bOFDM);
    
    % Selection matrices 
    for n = 0:bOFDM-1
      % Overlap-and-add of the filtered halfs [Eq. (20b)]
      w_hat_m0 = R_Nm0.'*w_hat_mn(:, n+i0); 
      w_hat_m1 = R_Nm1.'*w_hat_mn(:, n+i0); 

      % Time-domain synthesis matrix for 1st block in OLA processing [Eq. (13)]
      A_m0 = sparse(Ll-L_CP(n+i0)+(1:Ls+L_CP(n+i0)), ...  % Selection matrix
		    Ll-L_CP(n+i0)+(1:Ls+L_CP(n+i0)), ones(1, Ls+L_CP(n+i0)), L, L);
      
      % Block-wise OLA processing matrices of size Ns x L [Eq. (7a)]
      G_m0 = sqrt(L/N)*A_m0*WI_L*P*diag(win)*conj(M_LN(theta(0)*phi(n))).'*W_N; 
      G_m1 = sqrt(L/N)*A_mr*WI_L*P*diag(win)*conj(M_LN(theta(1)*phi(n))).'*W_N; 
      
      % Filtering [Eq. (20c)]
      z_hat_0(:, n+i0) = G_m0*w_hat_m0;
      z_hat_1(:, n+i0) = G_m1*w_hat_m1;

    end % END for n
    
    X_hat = zeros(L_OFDM, bOFDM);
    R_Lsm0 = [eye(L); zeros(L/2, L)];
    R_Lsm1 = [zeros(L/2, L); eye(L)];
    
    % Concatenation of the filtered blocks [Eq. (20d)]    
    for n = 0:bOFDM-1
      x_hat_OFDM = R_Lsm0*z_hat_0(:, n+i0) + R_Lsm1*z_hat_1(:, n+i0);
      
      % Here timing offset corresponding half the CP duration is used on the RX side
      time_offset = ceil(L_CP(n+i0)/2);  
      ZP = [ 
	  zeros(L_OFDM, L/4-time_offset) eye(L_OFDM) zeros(L_OFDM, L/4+time_offset);
	   ];
      x_hat = ZP*x_hat_OFDM;

      X_hat(:, n+i0) = 1/sqrt(L_OFDM)*fft(circshift(x_hat, [-time_offset 0]));
      
    end
     
    % Active subcarries on RX sides
    iAct_RX = mod([1:Lact]-floor(Lact/2)-1, L_OFDM)+1;
   
   case 1
    % Zero-padding needed in FC processing
    prePad = N-Ns; postPad = N-Ns; 

    nSamp_lowrate  = sum(L_CP, 1) + L_OFDM.*bOFDM; % Samples on the input rate
    
    % Total number of samples including the zero padding
    nTotSamp = nSamp_highrate + prePad + postPad;
    
    % Resulting number of FC prosessing blocks
    nBlks = ceil((nTotSamp-N)./Ns)+1;
    
    % Construction of block diagonal matrix with overlapping blocks
    Gm = zeros(L+Ls*(nBlks-1), N+Ns*(nBlks-1));    

    % Phase rotation syncing the phases of the TX and RX processing
    theta = @(r) exp(j*2*pi*Bin0*(r*Ls/L-(N-Ns-(Nl-N_CP(1)))/N_OFDM)); 
    
    for r = 0:nBlks-1 % Do for all processing blocks
      % Block-wise OLA processing matrix of size Ns x L [Eq. (7a)]
      Gmr = A_mr*WI_L*P*diag(win)*conj(M_LN(theta(r))).'*W_N;
      
      % Block-diagonal matrix with overlaping blocks [Eq. (5b)]
      Gm(r*Ls+(1:L), r*Ns+(1:N)) = Gm(r*Ls+(1:L), r*Ns+(1:N)) + Gmr;
      
    end % END for r
    
    x_hat = sqrt(L/N)*Gm*[zeros(prePad,1); w_m; zeros(postPad,1)];
    
    X_hat = OFDMdem(x_hat(L-Ls+1:end), bOFDM, L_OFDM, L_CP);
    
    % Active subcarries on RX sides
    iAct_RX = mod([1:Lact]-floor(Lact/2)-1, L_OFDM)+1;

  end
  
  % Difference between the transmitted and received symbols
  E = X_hat(iAct_RX, :) - X(iAct_TX, :);

  % Mean-squared error on active subcarries
  subplot(122)
  h = plot(10*log10(MSE(E)), lineStyle, 'Color', lineColor, 'MarkerFaceColor', 'w', 'DisplayName', legStr);
  hold on
  
end % END for iModel

legend('Location', 'North')
hold off
ylabel('MSE error in dB'); xlabel('Active subcarrier n'), grid on
title('MSE error on active subcarriers')
grid on

% Power spectral density estimate
subplot(121)  
nGridPoints = 2^15;
[Y, freq] = PSDeval(w_m, Fs, nSamp_highrate, nGridPoints);
h = plot(freq/1e6, 10*log10(Y), 'DisplayName', 'CP-OFDM TX waveform');

title({['PSD of the TX signal, center frequency = ' int2str(SCS*Bin0/1e3) ' kHz']})
xlabel('Frequency in MHz'), ylabel('PSD in dB')
legend('Location', 'South')
ylim([-80 10])
xlim([-Fs Fs]/2/1e6)
grid on

end

% SUBFUNCTIONS
% ------------------------------------------------------------------------------------
function mse = MSE(Err)

  % Mean-squared error over active subcarriers

  mse = sum(abs(Err).^2, 2)/size(Err, 2);

end

% ------------------------------------------------------------------------------------
function [X_CP_OFDM, x_CP_OFDM] = OFDMmod(X, Lact, bOFDM, L_OFDM, L_CP, L, Ls, Ll)

  % OFDM modulation, CP inclusion, and zero padding to length 3/2xL

  % CP-OFDM symbol vector (used for basic FC processing and PSD evaluation)
  x_CP_OFDM = zeros(bOFDM*L_OFDM+sum(L_CP), 1); 

  % CP-OFDM symbol matrix of size (3/2xL x bOFDM) 
  X_CP_OFDM = zeros(3/2*L, bOFDM);
  
  K = @(N, K) ([zeros(K, N-K) eye(K); eye(N)]); % CP insertion matrix
    
  for iSymb = 1:bOFDM
    % OFDM modulation [Eq. (1a)]
    x_OFDMn = sqrt(L_OFDM)*ifft(X(:, iSymb));

    % Add cyclic prefix [Eq. (1b)]
    x_CP_OFDMn = K(L_OFDM, L_CP(iSymb))*x_OFDMn;
    
    % Indices of the CP-OFDM symbol
    inx = (iSymb-1)*L_OFDM + sum(L_CP(1:iSymb-1)) + (1:L_OFDM+L_CP(iSymb)); 

    x_CP_OFDM(inx) = x_CP_OFDMn;
    
    % Zero-padding to length 3/2xL (Eq. 10)
    S_L = Ll-L_CP(iSymb);
    S_T = L-Ll-Ls;
    ZP = [ 
	zeros(S_L, L_OFDM+L_CP(iSymb));
	eye(L_OFDM+L_CP(iSymb));
	zeros(S_T, L_OFDM+L_CP(iSymb));
	 ];
    X_CP_OFDM(:, iSymb) = ZP*x_CP_OFDMn;
      
  end % END for iSymb

end

% ------------------------------------------------------------------------------------
function Y = OFDMdem(z, bOFDM, L_OFDM, L_CP)

  % Removal of the CP and OFDM demodulation

  Y = zeros(L_OFDM, bOFDM); 
  for iSymb = 1:bOFDM
    % Here, the reference symbol timing is in the middle of the CP
    time_offset = ceil(L_CP(iSymb)/2);
    
    OFDMsymb = z(L_OFDM*(iSymb-1)+sum(L_CP(1:iSymb))-time_offset+(1:L_OFDM));
    
    Y(:, iSymb) = 1/sqrt(L_OFDM)*fft(circshift(OFDMsymb, [-time_offset 0]));    
  
  end % END for iSymb

end

% ------------------------------------------------------------------------------------
function win = genWin(nTB, excessBW, Lact, SCS_fac, L)

  % Frequency-domain window function
  
  % Optimized transition-band weights. Here, the optimization 
  % is carried out following the techniques described in [2]

  tap = [ % JYK; 27.11.19, 14:25
	  -0.000340952946487
	  +0.000769433263641
	  +0.003043703626772
	  -0.012311509605436
	  -0.052593175378407
	  -0.081557396098516
	  -0.018092949343188
	  +0.183410321328965
	  +0.481282119474811
	  +0.759888460602338
	  +0.929330128417243
	  +0.991705584048036
	];

  win = [ones(ceil(Lact/2*SCS_fac+excessBW), 1);  
	 tap(nTB:-1:1); 
	 zeros(L-Lact*SCS_fac-2*excessBW-2*nTB, 1); 
	 tap(1:nTB); 
	 ones(floor(Lact/2*SCS_fac+excessBW), 1)];

  win = fftshift(win);
  
end

% ------------------------------------------------------------------------------------
function [Y, freq] = PSDeval(y, Fs, nSamp, nGridPoints)

  % Evaluate power spectral density estimate (magnitude-squared response)

  Y = 1/sum(nSamp)*abs(fftshift(fft(y, nGridPoints))).^2;

  freq = ((0:nGridPoints-1)/nGridPoints-0.5)*Fs;
  
end

% ------------------------------------------------------------------------------------
function [L_CP, N_CP] = CPlenghts(L_OFDM, N_OFDM, bOFDM)

  % CP lengths on low- and high-rate sides

  L_CP = repmat(floor(9/128*L_OFDM), [bOFDM 1]); % Low-rate CP lengths
  N_CP = repmat(floor(9/128*N_OFDM), [bOFDM 1]); % High-rate CP lengths
  L_CP(1:7:end) = floor(10/128*L_OFDM);          % Longer CP for 1st symbols in a slot
  N_CP(1:7:end) = floor(10/128*N_OFDM);          % Longer CP for 1st symbols in a slot

end
% ------------------------------------------------------------------------------------
function [Ls, Ll, Ns, Nl] = FCpartitioning(L, N, OLF)

  % FC processing block partition sizes

  Ls = round(L*(1 - OLF)); % Non-overlapping samples for short FC processing transf.
  Ll = ceil((L - Ls)/2);   % Leading overlapping samples for short FC process. transf.
  Ns = round(N*(1 - OLF)); % Non-overlapping samples for long FC processing transf.
  Nl = ceil((N - Ns)/2);   % Leading overlapping samples for long FC process. transf.

end
