%% 5G NR Downlink Channel Estimation
% Up to now, we have been assuming the channel is known at the receiver.  In this lab, we will build a simple channel estimator at the UE in a 5G downlink.
% In going through the lab, you will learn to:
% Configure and modulate references signals on the DM-RS
% Build a simple channel estimator in frequency and time via kernel regression
% Visualize equalized symbols
% Estimate the noise
% Evaluate the effect of channel estimation error via SNDR and BER
% Also, along with this file, you will have to make modifications for these files:
% NRgNBTxFD.m:  A class for NR gNB TX frequency-domain simulation
% NRUERxFD.m:  A class for NR UE RX frequency-domain simulation
% The parts to be modified are labeled TODO.  There are three other files that you do not need to modify and do not need to be submitted.
% FDChan.m:  A class for frequency-domain channel simulation
% plotChan.m:  A function for plotting channels on the OFDM grid
% kernelReg.m:  Kernel regression function
% Note:  At least on my version of MATLAB, the Live Editor hangs a lot. I have found that you can do the following when this happens.
% Enter something random in the command window like 1+1 .
% The command window will not respond since MATLAB is hung.  But, if you are on the Live Editor tab it will give you a chance to pause the process.
% Hit the pause and then MATLAB will stop where it is stuck.  Sometimes there is an error in your code, and sometimes MATLAB just decides to hang for no reason.
% A this point, you can Quit Debugging and get back control.
% Submission:  Complete this MATLAB live file and the TODO sections in NRgNBTxFD.m and  NRUERxFD.m.  Run the code.  Print this file and NRgNBTxFD.m and  NRUERxFD.m  to PDF.  Merge to a single PDF and submit the PDF for grading.
%
%% Configuring the System Parameters
% In this lab, we will configure the parameters with the following parameters that corresponds to a 10 MHz carrier.

fc = 2.3e9;                % Carrier frequency
SubcarrierSpacing = 15;    % SCS in kHZ
NRB = 51;  % number of resource blocks
nscPerRB = 12;  % number of sub-carriers per RB

%
% Similar to the previous lab, get the carrierConfig and waveformConfig with the nrCarrierConfig and nrOFDMInfo methods in the 5G Toolbox.

carrierConfig = nrCarrierConfig(...
    'NSizeGrid', NRB, 'SubcarrierSpacing', SubcarrierSpacing);
waveformConfig = nrOFDMInfo(carrierConfig);

%
%% Configuring the DM-RS and PDSCH
% The 5G NR DM-RS is highly configurable.  Set the parameters in the following function to configure the DM-RS: 3 REs per symbol at 2 different symbols in time (a total of 6 REs per RB).

% TODO:  Set the DM-RS config
%   dmrsConfig = nrPDSCHDMRSConfig(...
%       'NumCDMGroupsWithoutData', 1, ...  % No unused DM-RS
%       'DMRSAdditionalPosition', ?, ...   % Set this correctly
%       'DMRSConfigurationType', ?);       % Set this correctly
%       ...);

% Given the DM-RS configuration, we set the PDSCH:

% Configure the physical downlink shared channel parameters
pdschConfig = nrPDSCHConfig();
pdschConfig.NSizeBWP = [];   % Empty implies that the value is equal to NSizeGrid
pdschConfig.NStartBWP = [];  % Empty implies that the value is equal to NStartGrid
pdschConfig.PRBSet = (0:NRB-1);   % Allocate the complete carrier
pdschConfig.SymbolAllocation = [0 14]; % Symbol allocation [S L]
pdschConfig.MappingType = 'A'; % PDSCH mapping type ('A' or 'B')
pdschConfig.EnablePTRS = true;
pdschConfig.PTRS =  nrPDSCHPTRSConfig();
pdschConfig.DMRS = dmrsConfig;


%
%% Creating the gNB Transmitter
% Complete the code for the transmitter in the file NRgNBTxFD.m .  The code will:
% Generate random bits for the PDSCH
% Modulate the bits to the PDSCH symbols
% Determine the PT-RS and DM-RS symbols and add them to the TX grid

% TODO:  Finish the code in the NRgNBTxFD.m class

% Once you have completed the code you can create a TX object.

tx = NRgNBTxFD(carrierConfig, pdschConfig);
tx.step();

%
% Use the following command to plot the channels.

plotChan(tx.txGridChan, tx.chanNames);

%
% Now plot the channel types for one RB.

% TODO
%   plotChan(...)

%
% Compute and print the fraction of REs used for each channel.

% TODO

%
%% Simulating the Frequency-Domain OFDM Channel
% A helper funciton at the end of the script creates random channel parameters that we will use repeatedly.

% Channel parameters
dlymean = 200e-9;  % Mean delay spread
aoaAzSpread = 60;  % AoA azimuth spread
aoaElSpread = 20;  % AoA azimuth spread

% Get channel parameters
[gain, dly, aoaAz, aoaEl] = randParam(dlymean, aoaAzSpread, aoaElSpread);

%
% Given the channel parameters, we also set RX velocity.  You can try faster motion as well.

% Mobile velocity vector in m/s
rxVel = [25,0,0]';

% Then we construct the channel

% Parameters for computing the SNR
Etx = 1;        % Average transmitted symbol energy
snrAvg = 20;   % Average SNR

% Create a FD channel
fdchan = FDChan(carrierConfig, 'gain', gain, 'dly', dly, 'aoaAz', aoaAz, 'aoaEl', aoaEl, ...
    'rxVel', rxVel, 'Etx', Etx, 'EsN0Avg', snrAvg, 'fc', fc);

%
% The frequency-domain channel has the identical syntax as the the previous lab as well as the demo.
% To illusrate the channel, we first create a grid of TX symbols.  For now, we will just set them all to ones.

NumLayers = 1;
txGrid = nrResourceGrid(carrierConfig, NumLayers);
txGrid(:) = 1;

% Set the frameNum=0 and slotNum=0 and call the fdchan.step() method to get the channel grid and noise variance.  Compute a channel SNR matrix, chanSnr(i,j) representing the SNR in dB per RE.  Plot chanSnr with the imagesc command.  You should see some frequency selective fading.  But, minimal fast fading.  Nevertheless, there is sufficient variation over time that you would not want to assume the channel is constant over the slot.

frameNum = 0;
slotNum = 0;

% TODO:
%   [...] = fdchan.step(...);

%
% Plot the real  components of the channel at symbol = 8;

isym = 8;

% TODO:  
%   h = ...
%   plot(...);

%
%% Building the a Channel Estimator
% Now complete the chanEst method in the RX class, NRUERxFD.m.  Tthe channel estimator function:
% Extracts the DM-RS symbols
% Finds the sub-carrier indices for each symbol on which there are DM-RS references
% Creates the raw channel estimates on each such symbol
% Interpolates over time to get the final estimate over the entire

% TODO:  Complete the code in the chanEst method in the NRUERxFD class


%
% Once you are done with this you can test the channel estimate

% Create TX and RX
tx = NRgNBTxFD(carrierConfig, pdschConfig);
rx = NRUERxFD(carrierConfig, pdschConfig);

% Create a FD channel
snr = 20;
[gain, dly, aoaAz, aoaEl] = randParam(dlymean, aoaAzSpread, aoaElSpread);
fdchan = FDChan(carrierConfig, 'gain', gain, 'dly', dly, 'aoaAz', aoaAz, 'aoaEl', aoaEl, ...
    'rxVel', rxVel, 'Etx', Etx, 'EsN0Avg', snr, 'fc', fc);
        
% Run the TX
txGrid = tx.step();

% Run the channel
frameNum = 0;
slotNum = 0;
[rxGrid, chanGrid, noiseVar] = fdchan.step(txGrid, frameNum, slotNum);

% Run the channel estimator
rx.chanEst(rxGrid);

%
% For each i, plot the weightings over time, rx.Wtime(i,j) vs. j.  This indicates the weighting of the OFDM symbols on different reference symbols.

% TODO

%
% Get the channel estimate save in rx.chanEstGrid and plot the true and estiamted real values of the channel for the symbols 7 and 12.  Use subplots.   Since we ran the channel estimation at 20 dB, the channel estimation should look quite close on both symbols.

isymPlot = [7,12];

% TODO

%
% Compute and print the MSE of the channel estimate by comapring rx.chanEstGrid to chanGrid.

% TODO:
%    mse = ...

%
%% Using the Channel Estimation for Equalization and Demodulation
% Complete the stepImpl method in NRUERxFD class which performs the equalization and hard decision decoding of the bits.  Note that if we were building the full receiver, we would compute the LLRs and then send these to the decoder as we did in the previous lab.

% TODO:  Complete the TODO sections in stepImpl method of NRUERxFD.m

%
% Once you have done this you can complete the following code to test the receiver.  We will compute the BER using the true channel and the estimated channel and print the results. They should be similar.

% Create TX and RX
tx = NRgNBTxFD(carrierConfig, pdschConfig);
rx = NRUERxFD(carrierConfig, pdschConfig);

% Create a FD channel
snr = 20;
[gain, dly, aoaAz, aoaEl] = randParam(dlymean, aoaAzSpread, aoaElSpread);
fdchan = FDChan(carrierConfig, 'gain', gain, 'dly', dly, 'aoaAz', aoaAz, 'aoaEl', aoaEl, ...
    'rxVel', rxVel, 'Etx', Etx, 'EsN0Avg', snr, 'fc', fc);
        
% Run the TX
txGrid = tx.step();

% Run the channel
frameNum = 0;
slotNum = 0;
[rxGrid, chanGrid, noiseVar] = fdchan.step(txGrid, frameNum, slotNum);

% Run the RX with the estimated channel
rxBits = rx.step(rxGrid);


% TODO:  Compare the rxBits to tx.txBits to compute the BER with channel
% estimation error
%   berChanEst = ...

% TODO:  Run the RX where we supply the true channel and noise
%   berChanTrue = ...

% Print results
fprintf('BER chan est  = %12.4e\n', berChanEst);
fprintf('BER chan true = %12.4e\n', berChanTrue);
    

%
%% Testing over an SNR Range
%
% We conclude by testing over a range of SNRs.  At each SNR, we run a number of slots and measure the BER with the true and estimated channel.  We also estimate the SNDR.


% SNR values to test
snrTest = (0:2:30)';
nsnr = length(snrTest);

% number of trials at each SNR
ntrial = 50;

% Create TX and RX
tx = NRgNBTxFD(carrierConfig, pdschConfig);
rx = NRUERxFD(carrierConfig, pdschConfig);

% Create arrays to store results
nmeth = 2;  % 1=BER with estimated chan, 2=BER with true chan
ber = zeros(ntrial, nsnr, nmeth);
berAvg = zeros(nsnr, nmeth);
sndr = zeros(nsnr, 1);

for isnr = 1:nsnr
    snr = snrTest(isnr);
    
    % Distortion and signal energy in each trial
    distEnergy = zeros(ntrial,1);
    sigEnergy = zeros(ntrial,1);
    
    for itrial = 1:ntrial
        
        % Create a FD channel
        [gain, dly, aoaAz, aoaEl] = randParam(dlymean, aoaAzSpread, aoaElSpread);
        fdchan = FDChan(carrierConfig, 'gain', gain, 'dly', dly, 'aoaAz', aoaAz, 'aoaEl', aoaEl, ...
            'rxVel', rxVel, 'Etx', Etx, 'EsN0Avg', snr, 'fc', fc);
                
        % TODO:  Run the TX
        %   txGrid = ...
        
        
        % TODO:  Run the channel with frameNum = 0 and slotNum = 0
        %   [...] = fdchan.step(...);
        
        
        % TODO:  Run the RX with the channel estimate
        %    rxBits = rx(...);
        
        % TODO:  Get the BER with the channel estimate
        %    ber(itrial, isnr,1) = ...
        
        % TODO:  Measure distortion and signal energy using the PDSCH
        % symbols and channel in rx.pdschSym and rx.pdschChan at the RX and
        % tx.pdschSym at the TX
        %    distEnergy(itrial) = ...
        %    sigEnergy(itrial) = ...
        
        % TODO:  Run the RX with the true channel and noise variance
        %   rxBits = rx(...);
        
        % TODO:  Get the BER
        %    ber(itrial,isnr,2) = ...
                               
    end
    
    % TODO:  Measure the average BER for the SNR
    %    berAvg(isnr,:) = ...
    
    % TODO:  Measure the SNDR
    %    sndr(isnr) = ...
    
    % Plot the progress
    fprintf('SNR=%7.2f SNDR=%7.2f BER = %12.4e %12.4e\n', snr, sndr(isnr), berAvg(isnr,:));
end

%
% Plot the BER with the true and estimated channels.

% TODO.

%
% Plot the SNDR vs. SNR. You should see it saturate arond 20 dB of performance since there is mobility across the channel.

% TODO

%

function [gain, dly, aoaAz, aoaEl] = randParam(dlymean, aoaAzSpread, aoaElSpread)
    % Generates random parameters for a cluster of paths with
    % angular and delay spread
    
    % Default values
    if nargin < 1
        dlymean = 200e-9;
    end
    if nargin < 2
        aoaAzSpread = 60;
        aoaElSpread = 20;
    end
    
    % Create sub-path parameters
    npaths = 20;
    gain = zeros(npaths,1);      % path gain in dB
    dly = exprnd(dlymean, npaths,1);  % path delays in seconds
    aoaAz = unifrnd(-aoaAzSpread, aoaAzSpread, npaths,1);  % angles of arrival
    aoaEl = unifrnd(-aoaElSpread, aoaElSpread,npaths,1);
end

%