% Problem 2

close all;
clear all;

% Load the received signal and the training symbols
load rx_signal.mat
load training_symbols.mat

% Define parameters
N = 256; % number of carriers
L = 25; % length of the cyclic prefix


% Pre-process in an appropriate manner the training symbols, such that we
% can use them for finding the start of the cyclic prefix of the training
% OFDM block in the received samples.

ts_IFFT = ifft(training_symbols, N); % take the IFFT
ts_IFFT_with_CP = [ts_IFFT(end - L + 1 : end), ts_IFFT]; % add the cyclic prefix

% Find the start of the cyclic prefix of the training OFDM block in
% the received samples. Check point: you should find the beginning of the
% cyclic prefix at index 160 in the received samples.

R = xcorr(rx_signal, ts_IFFT_with_CP);
R = R(length(rx_signal) : end - length(ts_IFFT_with_CP) + 1); % remove ramp up/down
[~, index] = max(abs(R))

% plot for checking
figure;
stem(abs(R)); grid on;
xlabel('Index'); ylabel('|R|'); title('Correlation result')

% Implement the OFDM receiver for the portion of the received data starting
% with the training OFDM block. Note that you might need to trim the
% received data such that you obtain an integer number of OFDM symbols.

rx = rx_signal(index:end); % start at the training symbols
rx = rx(1:end - mod(length(rx), N + L)); % trim to obtain an integer number of OFDM symbols
rx = reshape(rx, N + L, []); 
rx = rx(L+1:end,:); % remove the cyclic prefix
rxFFT = fft(rx, N); 

% Using the LS approximation, estimate the vector lambda_LS of channel
% coefficients

lambda_LS = rxFFT(:,1).' ./ training_symbols; % LS is the component-wise division

% Plot the absolute value of the channel coefficients lambda_LS obtained above.
% The symbol-level impulse response of the channel is h = [1.621, 0.8448, 0.0283]. 
% Plot the corresponding lambdas (absolute value) on the same figure. 
% Check point: lambda_LS should be a noisy version of the lambdas obtained from h.

h = [1.621, 0.8448, 0.0283];
hFFT = fft(h, N); % lamdbas = DFT(h)

figure;
plot(abs(lambda_LS)); hold on;
plot(abs(hFFT)); grid on;
xlabel('Frequency index'); ylabel('|Lambdas|'); title('Lambda estimates');
legend('Lambdas with LS','Lambdas with DFT(h)');

% Using lambda_LS, equalize the OFDM symbols which follow the training
% symbols. Plot (scatterplot) the result. Check point: you should obtain a
% noisy 4-QAM constellation (4 distinct clouds, no rotation).

rxFFT = rxFFT(:, 2:end); % remove the training symbols
rxEq = rxFFT ./ (repmat(lambda_LS.', 1, size(rxFFT, 2)));

figure;
plot(rxEq(:), '*'); grid on;
xlabel('Re'); ylabel('Im'); title('Received constellation after OFDM demodulation and equalization');
