numFrames = 1; % 10 ms frames
snr = 20; % SNR in dB

% Create UE/carrier configuration
ue = nrCarrierConfig;
ue.NSizeGrid = 52;          % Bandwidth in number of resource blocks (52RBs at 15kHz SCS for 10MHz BW)
ue.SubcarrierSpacing = 15;  % 15, 30, 60, 120, 240 (kHz)
ue.CyclicPrefix = 'Normal'; % 'Normal' or 'Extended'

nTxAnts = 2;  % Number of transmit antennas (1,2,4)
nRxAnts = 2;  % Number of receive antennas
nLayers = min(nTxAnts,nRxAnts);

% Configure a periodic multi-port SRS and enable frequency hopping
srs = nrSRSConfig;
srs.NumSRSSymbols = 4;          % Number of OFDM symbols allocated per slot (1,2,4)
srs.SymbolStart = 8;            % Starting OFDM symbol within a slot
srs.NumSRSPorts = nTxAnts;      % Number of SRS antenna ports (1,2,4).
srs.FrequencyStart = 0;         % Frequency position of the SRS in BWP in RBs
srs.NRRC = 0;                   % Additional offset from FreqStart specified in blocks of 4 PRBs (0...67)
srs.CSRS = 14;                  % Bandwidth configuration C_SRS (0...63). It controls the allocated bandwidth to the SRS
srs.BSRS = 0;                   % Bandwidth configuration B_SRS (0...3). It controls the allocated bandwidth to the SRS
srs.BHop = 0;                   % Frequency hopping configuration (0...3). Set BHop < BSRS to enable frequency hopping
srs.KTC = 2;                    % Comb number (2,4). Frequency density in subcarriers
srs.Repetition = 2;             % Repetition (1,2,4). It disables frequency hopping in blocks of |Repetition| symbols
srs.SRSPeriod = [2 0];          % Periodicity and offset in slots. SRSPeriod(2) must be < SRSPeriod(1)
srs.ResourceType = 'periodic';  % Resource type ('periodic', 'semi-persistent','aperiodic'). Use 'aperiodic' to disable inter-slot frequency hopping

practicalSynchronization = true;

csiSubbandSize = 4; % Number of RBs per subband

channel = nrTDLChannel;
channel.DelayProfile = 'TDL-C';
channel.DelaySpread = 40e-9;
channel.MaximumDopplerShift = 30;
channel.NumTransmitAntennas = nTxAnts;
channel.NumReceiveAntennas = nRxAnts;
channel.Seed = 5;

% The sampling rate for the channel model is set using the value returned
% from hOFDMInfo.
ueInfo.NRB = ue.NSizeGrid;
ueInfo.SubcarrierSpacing = ue.SubcarrierSpacing;
ueInfo.CyclicPrefix = ue.CyclicPrefix;
ofdmInfo = hOFDMInfo(ueInfo);

channel.SampleRate = ofdmInfo.SamplingRate;

chInfo = info(channel);
maxChDelay = ceil(max(chInfo.PathDelays*channel.SampleRate));
maxChDelay = maxChDelay + chInfo.ChannelFilterDelay;

% Reset random generator for reproducibility
rng('default');

% Number of slots to simulate
numSlots = numFrames*ue.SlotsPerFrame;

% Total number of subcarriers and symbols per slot
K = ue.NSizeGrid * 12;
L = ue.SymbolsPerSlot;

% Initialize arrays storing channel estimates
allTxGrid = zeros([K L*numSlots nTxAnts]);

slotGridSize = [K L nRxAnts nTxAnts];
hEst = zeros(slotGridSize);
hestInterp = zeros(slotGridSize);
hEstUpdate = zeros(slotGridSize);

totalGridSize = [K L*numSlots nRxAnts nTxAnts];
allHest = zeros(totalGridSize);
allHestPerfect = zeros(totalGridSize);
allHestInterp = zeros(totalGridSize);

% Initialize noise power estimate
nvar = 0;

% Calculate the number of CSI subbands for the carrier
numCSISubbands = ceil(ue.NSizeGrid/csiSubbandSize);

% Initialize SINR per subband, slot, and PMI
maxPMI = hMaxPUSCHPrecodingMatrixIndicator(nLayers,nTxAnts);
sinrSubband = zeros([numCSISubbands numSlots maxPMI+1]);

% Initialize PMI matrix and SINR loss
pmi = NaN(numCSISubbands,numSlots);
pmiPerfect = pmi;
loss = zeros(size(pmi));

% Initialize timing estimation offset
offset = chInfo.ChannelFilterDelay;

% Calculate SRS CDM lengths
cdmLengths = hSRSCDMLengths(srs);

% OFDM symbols used for CSI acquisition
csiSelectSymbols = srs.SymbolStart + (1:srs.NumSRSSymbols);

for nSlot = 0:numSlots-1

    % Update slot and symbol counters
    ue.NSlot = nSlot;
    ueInfo.NSymbol = ue.NSlot * L;
    ueInfo.NSlot = ue.NSlot;

    % Generate SRS and map to slot grid
    [srsIndices,srsIndInfo] = nrSRSIndices(ue,srs);
    srsSymbols = nrSRS(ue,srs);

    % Create a slot-wise resource grid empty grid and map SRS symbols
    txGrid = zeros([K L nTxAnts]);
    txGrid(srsIndices) = srsSymbols;

    % Determine if the slot contains SRS
    isSRSSlot= ~isempty(srsSymbols);

    % OFDM Modulation
    [txWaveform,waveformInfo] = hOFDMModulate(ueInfo,txGrid);

    txWaveform = [txWaveform; zeros(maxChDelay, size(txWaveform,2))]; % required later to flush the channel filter to obtain the received signal

    % Transmission through channel
    [rxWaveform,pathGains] = channel(txWaveform);

    % Add AWGN to the received time domain waveform
    % Normalize noise power to take account of sampling rate, which is
    % a function of the IFFT size used in OFDM modulation. The SNR
    % is defined per RE for each receive antenna (TS 38.101-4).
    SNR = 10^(snr/20);    % Calculate linear noise gain
    N0 = 1/(sqrt(2.0*nRxAnts*double(waveformInfo.Nfft))*SNR);
    noise = N0*complex(randn(size(rxWaveform)),randn(size(rxWaveform)));

    rxWaveform = rxWaveform + noise;

    % Perform timing offset estimation
    pathFilters = getPathFilters(channel);

    % Timing estimation is only performed in the slots where the SRS is
    % transmitted to keep the perfect and practical channel estimation
    % synchronized.
    if isSRSSlot
        if practicalSynchronization
            % Practical synchronization. Correlate the received waveform
            % with the SRS to give timing offset estimate
            offset = nrTimingEstimate(rxWaveform,ueInfo.NRB,ueInfo.SubcarrierSpacing,nSlot,srsIndices,srsSymbols,'CyclicPrefix',ueInfo.CyclicPrefix);
        else
            offset = nrPerfectTimingEstimate(pathGains,pathFilters);
        end
    end

    % Perform OFDM demodulation
    rxGrid = hOFDMDemodulate(ueInfo,rxWaveform(1+offset:end,:));

    % Perform practical channel estimation
    % Update channel estimates only in the symbols and RBs containing SRS
    % in this slot and hold the estimates from the previous slot in other
    % locations. nvar is not updated when there is no SRS transmission. Use
    % a time-averaging window that covers all the SRS symbols transmitted.
    if isSRSSlot % this slot contains an SRS transmission
        [hEst,nvar] = nrChannelEstimate(rxGrid,srsIndices,srsSymbols,'AveragingWindow',[0 7],'CDMLengths',cdmLengths,'CyclicPrefix',ueInfo.CyclicPrefix);

        % Use channel estimate from previous slot for OFDM symbols before the first SRS symbol
        hestInterp = repmat(hestInterp(:,end,:,:),1,ue.SymbolsPerSlot);

        % Update channel estimate in OFDM symbols and RB where the SRS is
        % present and hold all channel estimates until the end of the slot
        firstSymbol = srs.SymbolStart+1;
        lastSymbol = srs.SymbolStart + srs.NumSRSSymbols;
        hEstUpdate(:,firstSymbol:lastSymbol,:,:) = hEst(:,firstSymbol:lastSymbol,:,:);
        hEstUpdate(:,lastSymbol:L,:,:) = repmat(hEst(:,lastSymbol,:,:),1,ue.SymbolsPerSlot-lastSymbol+1);

        idxHEstUpdate = hEstUpdate ~= 0; % Indices of updated channel estimates
        hestInterp(idxHEstUpdate) = hEstUpdate(idxHEstUpdate);
    else % Hold previous channel estimates if this slot does not contain SRS
        hestInterp = repmat(hestInterp(:,end,:,:),1,ue.SymbolsPerSlot);
    end

    % PMI Selection
    % Select the precoder matrix indicators for a number of layers using
    % the interpolated channel estimates in the OFDM symbols specified by
    % csiSelectSymbols. The PMIs are estimated per CSI subband
    [pmi(:,nSlot+1),sinrSubband(:,nSlot+1,:),subbandIndices] = hPMISelect(nLayers, hestInterp(:,csiSelectSymbols,:,:), nvar, csiSubbandSize);

    % PMI selection SINR loss
    % Calculate the performance loss as a ratio of the SINR after precoding
    % with PMIs selected using a practical channel estimate and the SINR
    % after precoding with PMIs selected using a perfect channel estimate.

    % Calculate perfect channel estimate for perfect PMI selection
    hEstPerfect = nrPerfectChannelEstimate(pathGains,pathFilters,ueInfo.NRB,ueInfo.SubcarrierSpacing,nSlot,offset,'CyclicPrefix',ueInfo.CyclicPrefix);

    % Perfect noise estimate from noise realization
    noiseGrid = hOFDMDemodulate(ueInfo,noise(1+offset:end,:));
    nvarPerfect = var(noiseGrid(:));

    [loss(:,nSlot+1),pmiPerfect(:,nSlot+1)] = hPMISelectionSINRLoss(pmi(:,nSlot+1), nLayers,  hEstPerfect(:,csiSelectSymbols,:,:), nvarPerfect);

    % Save a copy of all transmitted OFDM grids and channel estimates for
    % display purposes
    thisSlot = nSlot*L + (1:L); % Symbols of the current slot
    allTxGrid(:,thisSlot,:) = txGrid;
    allHest(:,thisSlot,:,:) = hEst;
    allHestInterp(:,thisSlot,:,:) = hestInterp;
    allHestPerfect(:,thisSlot,:,:) = hEstPerfect;

end

% Create x-axis and y-axis vectors
symbols = 0:(ue.NSlot+1)*ue.SymbolsPerSlot-1;
slots = 0:ue.NSlot;
subcarriers = 1:ue.NSizeGrid*12;
resourceBlocks = 1:ue.NSizeGrid;

figure
imagesc(symbols,subcarriers,abs(allTxGrid(:,:,1,1)));
xlabel('OFDM symbol'); ylabel('Subcarrier'); axis xy;
title('Transmitted SRS (port 1)');

% Remove first OFDM symbols not containing SRS to improve visualization of
% channel estimation error
hEstInterp = allHestInterp;
idx = 1:(srs.SRSPeriod(2)*ue.SymbolsPerSlot + srs.SymbolStart);
hEstInterp(:,idx,:,:) = NaN;
hEstInterp((srs.NRB*12+1):end,:,:,:) = NaN;

hChannelEstimationPlot(symbols,subcarriers,allHestPerfect,hEstInterp);

% First expand loss from subbands into RBs for display purposes
pmiRB = hExpandSubbandToRB(pmi, csiSubbandSize, ue.NSizeGrid);
pmiPerfectRB = hExpandSubbandToRB(pmiPerfect, csiSubbandSize, ue.NSizeGrid);
lossRB = hExpandSubbandToRB(loss, csiSubbandSize, ue.NSizeGrid);

hPMIPlot(slots,resourceBlocks,pmiRB,pmiPerfectRB,lossRB);

numLayers = min(size(allHestInterp,[3 4]));
if numLayers ~= 1
    pmiErr = sum( abs(pmi - pmiPerfect) > 0, [1 2])./ sum( ~isnan(pmi), [1 2]);

    TotPMIEst = sum(~isnan(pmi),[1 2]);
    fprintf('Number of estimated PMI: %d \n', TotPMIEst);
    fprintf('    Number of wrong PMI: %d \n', ceil(pmiErr*TotPMIEst));
    fprintf('         Relative error: %.1f (%%) \n', pmiErr*100);
else
    fprintf('For a single layer, PMI is always 0.\n');
end

hBestSINRPlot(slots,resourceBlocks,sinrSubband,pmi,csiSubbandSize);

% Displays perfect and practical channel estimates and the channel
% estimation error for the first transmit and receive ports. The channel
% estimation error is defined as the absolute value of the difference
% between the perfect and practical channel estimates.
function hChannelEstimationPlot(symbols,subcarriers,allHestPerfect,allHestInterp)

    figure
    subplot(311)
    imagesc(symbols, subcarriers, abs(allHestPerfect(:,:,1,1)));
    axis xy; xlabel('OFDM symbol'); ylabel('Subcarrier');
    colorbar;
    title('Perfect Channel Estimate (TxAnt=1, RxAnt=1)');

    subplot(312)
    imagesc(symbols, subcarriers, abs(allHestInterp(:,:,1,1)), ...
            'AlphaData',~isnan(allHestInterp(:,:,1,1)))
    axis xy; xlabel('OFDM symbol'); ylabel('Subcarrier');
    colorbar;
    title('SRS-based Practical Channel Estimate (TxAnt=1, RxAnt=1) ');

    % Display channel estimation error, defined as the difference between the
    % SRS-based and perfect channel estimates
    subplot(313)
    hestErr = abs(allHestInterp - allHestPerfect);
    imagesc(symbols, subcarriers, hestErr(:,:,1,1),...
            'AlphaData',~isnan(hestErr(:,:,1,1)));
    axis xy; xlabel('OFDM symbol'); ylabel('Subcarrier');
    colorbar;
    title('Channel Estimation Error (TxAnt=1, RxAnt=1)');

end

% Displays the PMI evolution and PMI estimation SINR loss over time and
% frequency. The SINR loss is defined as a ratio of the SINR after
% precoding with estimated and perfect PMIs. Estimated PMIs are
% obtained using a practical channel estimate and perfect PMIs are
% selected using a perfect channel estimate.
function hPMIPlot(slots,resourceBlocks,pmiRB,pmiPerfectRB,lossRB)

    figure
    subplot(311)
    imagesc(slots,resourceBlocks,pmiPerfectRB,'AlphaData',~isnan(pmiPerfectRB)); axis xy;
    c = caxis;
    cm = colormap;
    colormap( cm(1:floor(size(cm,1)/(c(2)-c(1)) -1):end,:) ); % Adjust colormap to PMI discrete values
    colorbar
    xlabel('Slot'); ylabel('Resource block'), title('PMI Selected using Perfect Channel Estimates')

    subplot(312)
    imagesc(slots,resourceBlocks,pmiRB,'AlphaData',~isnan(pmiRB)); axis xy;
    colorbar,
    xlabel('Slot'); ylabel('Resource block'), title('PMI Selected using SRS')

    subplot(313)
    imagesc(slots,resourceBlocks,lossRB,'AlphaData',~isnan(lossRB));
    colormap(gca,cm)
    xlabel('Slot'); ylabel('Resource block'); axis xy; colorbar;
    title('PMI Estimation SINR Loss (dB)')

end

% Displays the SINR per resource block obtained after precoding with the
% PMI that maximizes the SINR per subband.
function hBestSINRPlot(slots,resourceBlocks,sinrSubband,pmi,csiBandSize)

    % Display SINR after precoding with best PMI
    bestSINRPerSubband = nan(size(sinrSubband,[1 2]));

    % Get SINR per subband and slot using best PMI
    [sb,nslot] = find(~isnan(pmi));
    for i = 1:length(sb)
        bestSINRPerSubband(sb(i),nslot(i)) = sinrSubband(sb(i),nslot(i),pmi(sb(i),nslot(i))+1);
    end

    % First expand SINR from subbands into RBs for display purposes
    bestSINRPerRB = hExpandSubbandToRB(bestSINRPerSubband, csiBandSize, length(resourceBlocks));

    figure
    sinrdb = 10*log10(abs(bestSINRPerRB));
    imagesc(slots,resourceBlocks,sinrdb,'AlphaData',~isnan(sinrdb));
    axis xy; colorbar;
    xlabel('Slot');
    ylabel('Resource block')
    title('Average SINR Per Subband and Slot After Precoding with Best PMI (dB)')

end

% Expands a 2D matrix of values per subband in the first dimension into a
% matrix of values per resource block.
function rbValues = hExpandSubbandToRB(subbandValues, bandSize, NRB)

    lastBandSize = mod(NRB,bandSize);
    lastBandSize = lastBandSize + bandSize*(lastBandSize==0);

    rbValues = [kron(subbandValues(1:end-1,:),ones(bandSize,1));...
                subbandValues(end,:).*ones(lastBandSize,1)];
end