NFrames = 2;        % Total number of frames to simulate
SNRdB = 5;          % SNR in dB
fc = 6e9;           % Carrier freq. (Hz)

bsAntSize = [4 4];  % number of rows and columns in rectangular array (base station)
ueAntSize = [2 2];  % number of rows and columns in rectangular array (UE). Total number of antennas must be 1, 2 or 4

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

noDLSlotsPerULSlot = 4; % number of DL slots per UL slot

% Cyclic uplink/downlink pattern. One uplink slot followed by one or some
% downlink slots
ULDLpattern = [false true(1,noDLSlotsPerULSlot)]; % false: UL, true: DL
ULDLstring(ULDLpattern) = "DL";
ULDLstring(~ULDLpattern) = "UL";
disp("Uplink/Downlink cyclic pattern:")

disp("Slot " + string((0:length(ULDLstring)-1)') + ": " + ULDLstring(:))

% Numerology
SCS = 30; % subcarrier spacing: 15, 30, 60, 120, 240 (kHz)
NRB = 24; % bandwidth in number of resource blocks (24RBs at 30kHz SCS for 10MHz BW)

carrier = nrCarrierConfig;
carrier.NSizeGrid = NRB;
carrier.SubcarrierSpacing = SCS;
carrier.CyclicPrefix = 'Normal';  % 'Normal' or 'Extended'

% Get OFDM related information (sampling rate, symbol lengths, symbols
% per slot, slots per subframe/frame, FFT size)
ofdmInfo = nrOFDMInfo(carrier);

modulation = "16QAM";       % PDSCH modulation scheme
numLayers = 1;

pdsch = nrPDSCHConfig;
pdsch.NumLayers = numLayers;
pdsch.Modulation = modulation;
pdsch.PRBSet = 0:NRB-1;     % assume full band allocation
pdsch.DMRS.DMRSAdditionalPosition = 1; % one additional DMRS position to deal better with changing channels

verifyNumLayers(numLayers,bsAntSize,ueAntSize);

% DL-SCH parameters
codeRate = 490/1024;
RV = 0;

% Create DL-SCH encoder object
encodeDLSCH = nrDLSCH;
encodeDLSCH.TargetCodeRate = codeRate;

% Create DL-SCH decoder object
decodeDLSCH = nrDLSCHDecoder;
decodeDLSCH.TargetCodeRate = codeRate;
decodeDLSCH.LDPCDecodingAlgorithm = 'Normalized min-sum';
decodeDLSCH.MaximumLDPCIterationCount = 6;

channel = nrCDLChannel;
channel.DelayProfile = 'CDL-A';
channel.DelaySpread = 300e-9;
channel.CarrierFrequency = fc;
channel.MaximumDopplerShift = 100;
channel.SampleRate = ofdmInfo.SampleRate;

channel.TransmitAntennaArray.Size = [bsAntSize 1 1 1]; % Assume only 1 polarization and 1 panel of arrays
channel.ReceiveAntennaArray.Size = [ueAntSize 1 1 1];  % Assume only 1 polarization and 1 panel of arrays

chInfo = channel.info();
nBSAnts = chInfo.NumTransmitAntennas;   % Number of BS antennas
nUEAnts = chInfo.NumReceiveAntennas;    % Number of UE antennas

channel.displayChannel();

% SRS configuration, configure a multi-port SRS
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 = nUEAnts;      % Number of SRS antenna ports (1,2,4).
srs.FrequencyStart = 0;         % Frequency position of the SRS in RBs
srs.NRRC = 0;                   % Additional offset from FreqStart specified in blocks of 4 PRBs (0...67)
srs.CSRS = 7;                   % 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.SRSPeriod = [1 0];          % Periodicity and offset in slots. SRSPeriod(2) must be < SRSPeriod(1). Set the period to 1 so there is an SRS in all UL slots
srs.ResourceType = 'periodic';  % Resource type ('periodic','semi-persistent','aperiodic'). Use 'aperiodic' to disable inter-slot frequency hopping

% Generate uplink SRS waveform.
[srsWaveform,srsGrid,srsWaveformInfo,srsIndices,srsSymbols] = generateULSRSWaveform(carrier,srs);
% Display the transmitted OFDM grid containing SRS
% Create x-axis and y-axis vectors
symbols = 0:(carrier.NSlot+1)*carrier.SymbolsPerSlot-1;
subcarriers = 1:carrier.NSizeGrid*12;
figure;imagesc(symbols,subcarriers,abs(srsGrid(:,:,1)));
xlabel('OFDM symbol'); ylabel('Subcarrier'); axis xy;
title('Transmitted SRS (port 1)');

% Initialize result arrays
bitTput = [];           % Number of successfully received bits per transmission
txedTrBlkSizes = [];    % Number of transmitted info bits per transmission
evmPerSlot = zeros(ofdmInfo.SlotsPerFrame*NFrames,pdsch.NumLayers);
calculateEVM = comm.EVM();

diagnosticsOn = true;

for nSlot = 0:ofdmInfo.SlotsPerFrame*NFrames-1 
    
    if ~ULDLpattern(mod(nSlot,length(ULDLpattern))+1)
        
        if diagnosticsOn
            disp("Slot "+string(nSlot)+", UL slot");
        end
        
        % Reverse channel direction to swap the roles of the transmit and
        % receive arrays. Initially the channel is configured for the DL.
        % Assume there is always a DL slot before any UL slot.
        if ~channel.TransmitAndReceiveSwapped
            channel.swapTransmitAndReceive();
        end
        
        % New slot
        carrier.NSlot = nSlot;
        
        % Generate srsWaveform
        [srsWaveform,srsGrid,srsWaveformInfo,srsIndices,srsSymbols] = generateULSRSWaveform(carrier,srs);
        
        % Send uplink signal through the channel and add AWGN.        
        % Pad UL waveforms with zeros to flush the channel filter
        chInfo = info(channel);
        maxChDelay = ceil(max(chInfo.PathDelays*channel.SampleRate));
        maxChDelay = maxChDelay + chInfo.ChannelFilterDelay;
        srsWaveform = [srsWaveform; zeros(maxChDelay, size(srsWaveform,2))]; % required later to flush the channel filter to obtain the received signal
        
        % Transmission through channel
        rxWaveform = channel(srsWaveform);
        
        % Add AWGN to the received time domain waveform
        noise = generateAWGN(SNRdB,ofdmInfo.Nfft,size(rxWaveform));
        rxWaveform = rxWaveform + noise;
        % Use SRS for timing estimation
        offset = nrTimingEstimate(carrier,rxWaveform,srsIndices,srsSymbols);
        
        % Perform OFDM demodulation
        rxGrid = nrOFDMDemodulate(carrier,rxWaveform(1+offset:end,:));
        
        % Perform practical channel estimation
        cdmLengths = hSRSCDMLengths(srs); % Calculate SRS CDM lengths
        hEst = nrChannelEstimate(carrier,rxGrid,srsIndices,srsSymbols,'AveragingWindow',[0 7],'CDMLengths',cdmLengths);
        
        % Get Beamforming Weights
        % Calculate these weights for the maximum number of layers the
        % system can theoretically support, i.e. min(nBSAnts,nUEAnts).
        % Later we can select a smaller number of layers.
        nUplinkLayers = min(nBSAnts,nUEAnts); % number of layers used in the uplink transmission
        [wue,wbs,~] = getBeamformingWeights(hEst,nUplinkLayers);
        wbs = wbs(1:pdsch.NumLayers,:); % for precoding only use weights corresponding to the number of PDSCH layers
        wbs = wbs /sqrt(pdsch.NumLayers);
        
        % The SRS channel estimate was obtained with an uplink transmission. Switch
        % the channel direction to model the downlink. This switches the roles of the
        % transmit and receive arrays, i.e. the transmit array becomes the receive array
        % and vice versa.        
        if channel.TransmitAndReceiveSwapped
            channel.swapTransmitAndReceive();
        end
        
    else % DL slot        
        % New slot
        carrier.NSlot = nSlot;
        
        % Generate PDSCH waveform and precode using the beamforming weights
        % obtained from SRS channel estimation
        [txWaveform,waveformInfo,pdschResources,dmrsResources,trBlkSizes] = ...
            generatePDSCHWaveform(carrier,pdsch,encodeDLSCH,RV,wbs);
        
        pdschSymbolsTx = pdschResources.pdschSymbolsTx;
        pdschIndices = pdschResources.pdschIndices;
        
        dmrsSymbols = dmrsResources.dmrsSymbols;
        dmrsIndices = dmrsResources.dmrsIndices;
        
        % Send signal through the CDL channel and add noise
        chInfo = info(channel);
        % Pad with zeros to flush the channel filter and obtain the received signal
        maxChDelay = ceil(max(chInfo.PathDelays*channel.SampleRate)) + chInfo.ChannelFilterDelay;
        txWaveform = [txWaveform; zeros(maxChDelay, size(txWaveform,2))];
        
        rxWaveform = channel(txWaveform);
        
        % Add AWGN to the received time-domain waveform
        noise = generateAWGN(SNRdB,ofdmInfo.Nfft,size(rxWaveform));
        rxWaveform = rxWaveform + noise;
        
        % Synchronize to received signal using the DM-RS
        [offset,mag] = nrTimingEstimate(carrier,rxWaveform,dmrsIndices,dmrsSymbols);
        rxWaveform = rxWaveform(1+offset:end,:);
        
        % OFDM demodulation
        rxGrid = nrOFDMDemodulate(carrier,rxWaveform);
        
        % Channel estimation
        [estChannelGrid,noiseEst] = nrChannelEstimate(carrier,rxGrid,dmrsIndices,dmrsSymbols,'CDMLengths',pdsch.DMRS.CDMLengths);
        
        % Equalization and PDSCH decoding
        [pdschRx,pdschHest] = nrExtractResources(pdschIndices,rxGrid,estChannelGrid);
        [pdschEq,csi] = nrEqualizeMMSE(pdschRx,pdschHest,noiseEst);        
        [dlschLLRs,rxSymbols] = nrPDSCHDecode(carrier,pdsch,pdschEq,noiseEst);
        
        % Scale LLRs by CSI
        csi = nrLayerDemap(csi); % CSI layer demapping
        for cwIdx = 1:pdsch.NumCodewords
            Qm = length(dlschLLRs{cwIdx})/length(rxSymbols{cwIdx}); % bits per symbol
            csi{cwIdx} = repmat(csi{cwIdx}.',Qm,1);   % expand by each bit per symbol
            dlschLLRs{cwIdx} = dlschLLRs{cwIdx} .* csi{cwIdx}(:);   % scale
        end

        % Decode the DL-SCH transport channel
        decodeDLSCH.TransportBlockLength = trBlkSizes;
        [decbits,blkerr] = decodeDLSCH(dlschLLRs,pdsch.Modulation,pdsch.NumLayers,RV);
        reset(decodeDLSCH); % reset internal soft buffers, assume no retransmissions
        
        %  Store values to calculate throughput (only for active PDSCH instances)
        if(any(trBlkSizes ~= 0))
            bitTput = [bitTput trBlkSizes.*(1-blkerr)];
            txedTrBlkSizes = [txedTrBlkSizes trBlkSizes];                         
        end
        
        if diagnosticsOn
            if blkerr
                disp("Error in DL slot "+string(nSlot));
            else
                disp("Slot "+string(nSlot)+", successful DL reception");
            end
            evmPerSlot(nSlot+1,:)=calculateEVM(pdschSymbolsTx,pdschEq);
        end
    end
end

if diagnosticsOn
    tegendTxt =cell(pdsch.NumLayers,1);
    for n=1:pdsch.NumLayers
        tegendTxt{n} = "layer "+string(n);
    end
    figure;bar(0:ofdmInfo.SlotsPerFrame*NFrames-1,evmPerSlot)
    xlabel("Slot Number");ylabel("EVM (%)")
    title("PDSCH Symbols EVM")
    legend(tegendTxt);
end

% Calculate maximum and simulated throughput
maxThroughput = sum(txedTrBlkSizes); % Max possible throughput
simThroughput = sum(bitTput,2);      % Simulated throughput

% Display the results dynamically in the command window
fprintf([['\n\nThroughput(Mbps) for ', num2str(NFrames) ' frame(s) '],...
    '= %.4f Mbps\n'], 1e-6*simThroughput/(NFrames*10e-3));

fprintf(['Throughput(%%) for ', num2str(NFrames) ' frame(s) = %.4f %%\n'],...
    simThroughput*100/maxThroughput);

function noise = generateAWGN(SNRdB,Nfft,sizeRxWaveform)
% Generate AWGN for a given value of SNR in dB (SNRDB). This is the
% receiver SNR per resource element and antenna assuming the channel does
% not affect the power of the signal. NFFT is the FFT size used in OFDM
% demodulation. SIZERXWAVEFORM is the size of the receive waveform used to
% calculate the size of the noise matrix.

    % 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).
    nRxAnts = sizeRxWaveform(2);
    SNR = 10^(SNRdB/20); % Calculate linear noise gain
    N0 = 1/(sqrt(2.0*nRxAnts*double(Nfft))*SNR);
    noise = N0*complex(randn(sizeRxWaveform),randn(sizeRxWaveform));
end

function [wtx,wrx,D] = getBeamformingWeights(hEst,nLayers)
% Get beamforming weights given a channel matrix HEST and the number of
% layers NLAYERS. One set of weights is provided for the whole bandwidth.
% The beamforming weights are calculated using singular value (SVD)
% decomposition. 
%
% This function returns the beamforming weights for the transmitter and the
% receiver (WTX and WRX) and the singular value matrix (D).
    
    % Average channel estimate
    [~,~,R,P] = size(hEst);
    H = permute(mean(reshape(hEst,[],R,P)),[2 3 1]);
    
    % SVD decomposition
    [U,D,V] = svd(H);
    wtx = V(:,1:nLayers).';
    wrx = U(:,1:nLayers)';
end

function [srsWaveform,srsGrid,srsWaveformInfo,srsIndices,srsSymbols] = generateULSRSWaveform(ue,srs)
% Generate the SRS symbols, map to slot grid and OFDM modulate. This
% function returns the time and frequency domain waveforms (SRSWAVEFORM and SRSGRID),
% associated information (SRSWAVEFORMINFO), and the SRS symbols and its
% mapping indices (SRSSYMBOLS and SRSINDICES).
%
% The function inputs are the UE carrier and the SRS configuration object.
    
    [srsIndices,~] = nrSRSIndices(ue,srs);
    srsSymbols = nrSRS(ue,srs);
    
    % Create a slot-wise resource grid empty grid and map SRS symbols
    srsGrid = nrResourceGrid(ue,srs.NumSRSPorts);
    srsGrid(srsIndices) = srsSymbols;
    
    % Determine if the slot contains SRS
    isSRSSlot= ~isempty(srsSymbols);
    if ~isSRSSlot
        error("nr5g:TDDReciprocityExample:NoSRS","Slot "+ string(ue.NSlot) + ": No SRS present");
    end
    
    % OFDM Modulation
    [srsWaveform,srsWaveformInfo] = nrOFDMModulate(ue,srsGrid);
end

function [txWaveform,waveformInfo,pdschResources,dmrsResources,trBlkSizes] = ...
    generatePDSCHWaveform(carrier,pdsch,encodeDLSCH,RV,wbs)
% For a given slot, generate transport block, apply DL-SCH encoding and
% PDSCH modulate. Generate the PDSCH DM-RS. Precode resulting signal and
% OFDM modulate. This function returns the time-domain signal TXWAVEFORM and
% associated information WAVEFORMINFO. The resource structures
% PDSCHRESOURCES and DMRSRESOURCES include the PDSCH and DM-RS symbols and
% their respective mapping indices. The function also returns the transport
% block sizes (TRBLKSIZES).
%
% The function inputs include the CARRIER and PDSCH configuration objects,
% the DL-SCH encoder object (ENCODEDLSCH), the redundancy version (RV), and
% the precoding weights (WBS).

    % Generate PDSCH indices
    [pdschIndices,pdschInfo] = nrPDSCHIndices(carrier,pdsch);
    
    % Calculate transport block sizes
    Xoh_PDSCH = 0;
    codeRate = encodeDLSCH.TargetCodeRate;
    trBlkSizes = nrTBS(pdsch.Modulation,pdsch.NumLayers,numel(pdsch.PRBSet),pdschInfo.NREPerPRB,codeRate,Xoh_PDSCH);
    
    for cwIdx = 1:pdsch.NumCodewords
        trBlk = randi([0 1],trBlkSizes(cwIdx),1);
        setTransportBlock(encodeDLSCH,trBlk,cwIdx-1);
    end
    
    % Encode the DL-SCH transport blocks
    codedTrBlock = encodeDLSCH(pdsch.Modulation,pdsch.NumLayers,pdschInfo.G,RV);
    
    % Generate PDSCH symbols and precode
    pdschSymbolsTx = nrPDSCH(carrier,pdsch,codedTrBlock);
    pdschSymbolsTxPrecoded = pdschSymbolsTx*wbs;
    
    pdschResources.pdschSymbolsTx = pdschSymbolsTx;
    pdschResources.pdschIndices = pdschIndices;
    
    % Create an empty resource grid and map the PDSCH symbols. The PDSCH
    % indices are calculated for a single antenna port. After precoding,
    % map the resulting signal to multiple antennas. The function
    % nrExtractResources() helps to convert the single-port PDSCH indices
    % to multiple antennas
    nBSAnts = size(wbs,2);
    pdschGrid = nrResourceGrid(carrier,nBSAnts);
    [~,pdschAntIndices] = nrExtractResources(pdschIndices,pdschGrid);
    pdschGrid(pdschAntIndices) = pdschSymbolsTxPrecoded;
    
    % Generate DM-RS symbols and indices
    dmrsSymbols = nrPDSCHDMRS(carrier,pdsch);
    dmrsIndices = nrPDSCHDMRSIndices(carrier,pdsch);
    
    dmrsResources.dmrsSymbols = dmrsSymbols;
    dmrsResources.dmrsIndices = dmrsIndices;
    
    % PDSCH DM-RS precoding and mapping. Similarly to the PDSCH, the DM-RS
    % indices are calculated for a single port, the function
    % nrExtractResources() helps to convert these to multiantenna indices
    for p = 1:size(dmrsSymbols,2)
        [~,dmrsAntIndices] = nrExtractResources(dmrsIndices(:,p),pdschGrid);
        pdschGrid(dmrsAntIndices) = pdschGrid(dmrsAntIndices) + dmrsSymbols(:,p)*wbs(p,:);
    end
    
    % Perform OFDM modulation of the resource grid
    [txWaveform,waveformInfo] = nrOFDMModulate(carrier,pdschGrid);
end

function cdmLengths = hSRSCDMLengths(srs)
%   Get CDM lengths for the SRS configuration SRS.
%
%   SRS is an SRS-specific configuration object with properties:
%   NumSRSPorts     - Number of SRS antenna ports (1,2,4)
%   KTC             - Transmission comb number (2,4)
%   CyclicShift     - Cyclic shift number offset (0...NCSmax-1).
%                     NCSmax = 12 if KTC = 4 and NCSmax = 8 if KTC = 2
    if srs.NumSRSPorts == 1
        cdmLengths = [1 1];
    elseif srs.NumSRSPorts == 2
        cdmLengths = [2 1];
    elseif (srs.KTC == 2 && srs.CyclicShift >= 4) || (srs.KTC == 4 && srs.CyclicShift >= 6)
        cdmLengths = [2 1];
    else
        cdmLengths = [4 1];
    end
end

function verifyNumLayers(numLayers,bsAntSize,ueAntSize)
    if numLayers>min(prod(bsAntSize),prod(ueAntSize))
        error("nr5g:TDDReciprocityExample:InvalidNumLayers",...
            'The number of layers (%d) must satisfy NLayers <= min(num BS ants,num UE ants) = min(%d,%d) = (%d)',...
            numLayers,prod(bsAntSize),prod(ueAntSize),min(prod(bsAntSize),prod(ueAntSize)));
    end
end