nFrames = 5; % Number of 10 ms frames
SNRdB = 0;   % SNR in dB

carrier = nrCarrierConfig;
carrier.SubcarrierSpacing = 15;
carrier.NSizeGrid = 52;

NStartBWP = 0;
NSizeBWP = 52;

csirs = nrCSIRSConfig;
csirs.CSIRSType = {'nzp','nzp','nzp'};
csirs.RowNumber = [4 4 4];
csirs.NumRB = 52;
csirs.RBOffset = 0;
csirs.CSIRSPeriod = [4 0];
csirs.SymbolLocations = {0, 0, 0};
csirs.SubcarrierLocations = {0, 4, 8};
csirs.Density = {'one','one','one'};

nTxAnts = 4;
nRxAnts = 4;

validateCSIRSConfig(carrier,csirs,nTxAnts);

numLayers = 4;

% Configure CSI reporting configuration parameters
reportConfig.NStartBWP = NStartBWP;
reportConfig.NSizeBWP = NSizeBWP;
reportConfig.PanelDimensions = [2 1];
reportConfig.CQIMode = 'Subband';
reportConfig.PMIMode = 'Subband';
reportConfig.SubbandSize = 4;
reportConfig.PRGSize = [];
reportConfig.CodebookMode = 1;
reportConfig.CodebookSubsetRestriction = [];
reportConfig.i2Restriction = [];

% Create channel object
channel = nrTDLChannel;
channel.NumTransmitAntennas = nTxAnts;
channel.NumReceiveAntennas = nRxAnts;
channel.DelayProfile = 'TDL-C';
channel.MaximumDopplerShift = 50;
channel.DelaySpread = 300e-9;

waveformInfo = nrOFDMInfo(carrier);
channel.SampleRate = waveformInfo.SampleRate;

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

% Initialize variables to store CQI, PMI, and subband SINR values for
% practical and perfect channel estimation scenarios
cqiPracticalPerSlot = [];
subbandCQIPractical = [];
pmiPracticalPerSlot = struct('i1',[],'i2',[]);
SINRPerSubbandPerCWPractical = [];
cqiPerfectPerSlot = [];
subbandCQIPerfect = [];
pmiPerfectPerSlot = struct('i1',[],'i2',[]);
SINRPerSubbandPerCWPerfect = [];

% Calculate total number of slots
totSlots = nFrames*carrier.SlotsPerFrame;

% Get number of CSI-RS ports
csirsPorts = csirs.NumCSIRSPorts(1);

% Get CDM lengths corresponding to configured CSI-RS resources
cdmLengths = getCDMLengths(csirs);

% Loop over all slots
for slotIdx = 0:totSlots - 1
    % Create carrier resource grid for one slot
    csirsSlotGrid = nrResourceGrid(carrier,csirsPorts);

    % Update slot number in carrier configuration object
    carrier.NSlot = slotIdx;

    % Generate CSI-RS indices and symbols
    csirsInd = nrCSIRSIndices(carrier,csirs);
    csirsSym = nrCSIRS(carrier,csirs);

    % Map CSI-RS to slot grid
    csirsSlotGrid(csirsInd) = csirsSym;

    % Map CSI-RS ports to transmit antennas
    wtx = eye(csirsPorts,nTxAnts);
    txGrid = reshape(reshape(csirsSlotGrid,[],csirsPorts)*wtx,size(csirsSlotGrid,1),size(csirsSlotGrid,2),nTxAnts);

    % Perform OFDM modulation to generate time-domain waveform
    txWaveform = nrOFDMModulate(carrier,txGrid);

    % Append zeros at the end of the transmitted waveform to flush channel
    % content. These zeros take into account any delay introduced in the
    % channel. The channel delay is a mix of multipath delay and
    % implementation delay. This value may change depending on the sampling
    % rate, delay profile, and delay spread.
    txWaveform = [txWaveform; zeros(maxChDelay,size(txWaveform,2))]; %#ok<AGROW> 

    % Transmit waveform through channel
    [rxWaveform,pathGains,sampleTimes] = channel(txWaveform);

    % Generate and add AWGN to received waveform
    SNR = 10^(SNRdB/20);                                                % Linear SNR value
    N0 = 1/(sqrt(2.0*nRxAnts*double(waveformInfo.Nfft))*SNR);           % Noise standard deviation
    rng('default');                                                     % Set RNG state for repeatability
    noise = N0*complex(randn(size(rxWaveform)),randn(size(rxWaveform)));
    rxWaveform = rxWaveform + noise;

    % Perform practical timing estimation
    offsetPractical = nrTimingEstimate(carrier,rxWaveform,csirsInd,csirsSym);
    if offsetPractical > maxChDelay
        offsetPractical = 0;
    end

    % Get path filters
    pathFilters = getPathFilters(channel);
    % Perform perfect timing estimation
    offsetPerfect = nrPerfectTimingEstimate(pathGains,pathFilters);

    % Perform time-domain offset correction for practical and
    % perfect timing estimation scenarios
    rxWaveformPractical = rxWaveform(1+offsetPractical:end,:);
    rxWaveformPerfect = rxWaveform(1+offsetPerfect:end,:);

    % Perform OFDM demodulation on previously synchronized waveforms
    rxGridPractical = nrOFDMDemodulate(carrier,rxWaveformPractical);
    rxGridPerfect = nrOFDMDemodulate(carrier,rxWaveformPerfect);

    % Consider only the NZP-CSI-RS symbols for channel estimation
    nzpCSIRSSym = csirsSym(csirsSym ~= 0);
    % Consider only the NZP-CSI-RS indices for channel estimation
    nzpCSIRSInd = csirsInd(csirsSym ~= 0);
    % Calculate practical channel estimate using 'AveragingWindow' as [5 1]
    [PracticalHest,nVarPractical] = nrChannelEstimate(carrier,rxGridPractical, ...
        nzpCSIRSInd,nzpCSIRSSym,'CDMLengths',cdmLengths,'AveragingWindow',[5 1]);

    % Perform perfect channel estimation
    PerfectHest = nrPerfectChannelEstimate(carrier,pathGains,pathFilters,offsetPerfect,sampleTimes);

    % Get perfect noise estimate value from noise realization
    noiseGrid = nrOFDMDemodulate(carrier,noise(1+offsetPerfect:end,:));
    nVarPerfect = var(noiseGrid(:));

    % Calculate CQI and PMI values using practical channel estimate
    [cqiPractical,pmiPractical,cqiInfoPractical] = hCQISelect(carrier,csirs, ...
        reportConfig,numLayers,PracticalHest,nVarPractical);

    % Store CQI, PMI, and subband SINR values of each slot for
    % practical channel estimation scenario
    cqiPracticalPerSlot(:,:,slotIdx+1) = cqiPractical; %#ok<SAGROW>
    pmiPracticalPerSlot(slotIdx+1) = pmiPractical;
    subbandCQIPractical(:,:,slotIdx+1) = cqiInfoPractical.SubbandCQI; %#ok<SAGROW>
    SINRPerSubbandPerCWPractical(:,:,slotIdx+1) = cqiInfoPractical.SINRPerSubbandPerCW; %#ok<SAGROW>

    % Calculate CQI and PMI values using perfect channel estimate
    [cqiPerfect,pmiPerfect,cqiInfoPerfect] = hCQISelect(carrier,csirs, ...
        reportConfig,numLayers,PerfectHest,nVarPerfect);

    % Store CQI, PMI, and subband SINR values of each slot for
    % perfect channel estimation scenario
    cqiPerfectPerSlot(:,:,slotIdx+1) = cqiPerfect; %#ok<SAGROW>
    subbandCQIPerfect(:,:,slotIdx+1) = cqiInfoPerfect.SubbandCQI; %#ok<SAGROW>
    pmiPerfectPerSlot(slotIdx+1) = pmiPerfect;
    SINRPerSubbandPerCWPerfect(:,:,slotIdx+1) = cqiInfoPerfect.SINRPerSubbandPerCW; %#ok<SAGROW>
end

plotWidebandCQIAndSINR(cqiPracticalPerSlot,cqiPerfectPerSlot,SINRPerSubbandPerCWPractical,SINRPerSubbandPerCWPerfect);

% Provide slot number for which subband CQI values and PMI i2 indices must be plotted
slotNum = 1; % Consider slot number as 1 (1-based) here, because CSI (CQI and PMI) is
             % reported in that slot for the configured CSI-RS resource(s)
plotSubbandCQIAndSINR(subbandCQIPractical,subbandCQIPerfect,SINRPerSubbandPerCWPractical,SINRPerSubbandPerCWPerfect,slotNum);

plotPMI(pmiPracticalPerSlot,pmiPerfectPerSlot,slotNum);

function validateCSIRSConfig(carrier,csirs,nTxAnts)
%   validateCSIRSPorts(CARRIER,CSIRS,NTXANTS) validates the CSI-RS
%   configuration, given the carrier specific configuration object CARRIER,
%   CSI-RS configuration object CSIRS, and the number of transmit antennas
%   NTXANTS.

    % Validate the number of CSI-RS ports
    if ~isscalar(unique(csirs.NumCSIRSPorts))
        error('nr5g:InvalidCSIRSPorts',...
            'All the CSI-RS resources must be configured to have the same number of CSI-RS ports.');
    end

    % Validate the CDM lengths
    if ~iscell(csirs.CDMType)
        cdmType = {csirs.CDMType};
    else
        cdmType = csirs.CDMType;
    end
    if (~all(strcmpi(cdmType,cdmType{1})))
        error('nr5g:InvalidCSIRSCDMTypes',...
            'All the CSI-RS resources must be configured to have the same CDM lengths.');
    end
    if nTxAnts < csirs.NumCSIRSPorts(1)
        error('nr5g:InvalidNumTxAnts',['Number of transmit antennas (' num2str(nTxAnts)...
            ') must be greater than or equal to the number of CSI-RS ports (' num2str(csirs.NumCSIRSPorts(1)) ').']);
    end

    % Check for the overlap between the CSI-RS indices
    csirsInd = nrCSIRSIndices(carrier,csirs,"OutputResourceFormat",'cell');
    numRes = numel(csirsInd);
    csirsIndAll = cell(1,numRes);
    ratioVal = csirs.NumCSIRSPorts(1)/prod(getCDMLengths(csirs));
    for resIdx = 1:numRes
        if ~isempty(csirsInd{resIdx})
            grid = nrResourceGrid(carrier,csirs.NumCSIRSPorts(1));
            [~,tempInd] = nrExtractResources(csirsInd{resIdx},grid);
            if numel(tempInd)/numel(csirsInd{resIdx}) ~= ratioVal
                error('nr5g:OverlappedCSIRSREsSingleResource',['CSI-RS indices of resource '...
                    num2str(resIdx) ' must be unique. Try changing the symbol or subcarrier locations.']);
            end
            csirsIndAll{resIdx} = tempInd(:);
            for idx = 1:resIdx-1
                overlappedInd = ismember(csirsIndAll{idx},csirsIndAll{resIdx});
                if any(overlappedInd)
                    error('nr5g:OverlappedCSIRSREsMultipleResources',['The resource elements of the '...
                        'configured CSI-RS resources must not overlap. Try changing the symbol or '...
                        'subcarrier locations of CSI-RS resource ' num2str(idx) ' and resource ' num2str(resIdx) '.']);
                end
            end
        end
    end
end

function cdmLengths = getCDMLengths(csirs)
%   CDMLENGTHS = getCDMLengths(CSIRS) returns the CDM lengths given the
%   CSI-RS configuration object CSIRS.

    CDMType = csirs.CDMType;
    if ~iscell(csirs.CDMType)
        CDMType = {csirs.CDMType};
    end
    CDMTypeOpts = {'noCDM','fd-CDM2','CDM4','CDM8'};
    CDMLengthOpts = {[1 1],[2 1],[2 2],[2 4]};
    cdmLengths = CDMLengthOpts{strcmpi(CDMTypeOpts,CDMType{1})};
end

function plotWidebandCQIAndSINR(cqiPracticalPerSlot,cqiPerfectPerSlot,SINRPerSubbandPerCWPractical,SINRPerSubbandPerCWPerfect)
%   plotWidebandCQIAndSINR(CQIPRACTICALPERSLOT,CQIPERFECTPERSLOT,SINRPERSUBBANDPERCWPRACTICAL,SINRPERSUBBANDPERCWPERFECT)
%   plots the wideband SINR and wideband CQI values for each codeword
%   across all active slots (in which the CQI is reported as other
%   than NaN) for practical and perfect channel estimation cases.

    numCW = size(cqiPracticalPerSlot,2);
    % Plot wideband CQI values
    for cwIdx = 1:numCW
        cqiPracticalPerCW = permute(cqiPracticalPerSlot(1,cwIdx,:),[1 3 2]);
        cqiPerfectPerCW = permute(cqiPerfectPerSlot(1,cwIdx,:),[1 3 2]);
        SINRPerCWPractical = permute(SINRPerSubbandPerCWPractical(1,cwIdx,:),[1 3 2]);
        SINRPerCWPerfect = permute(SINRPerSubbandPerCWPerfect(1,cwIdx,:),[1 3 2]);

        % Get logical array to represent active slots in which CQI is
        % reported
        activeSlotsLogAry = ~isnan(cqiPracticalPerCW);

        % Get slot numbers (0-based) in which CQI is reported
        activeSlotNum = find(activeSlotsLogAry)-1;

        % Extract wideband CQI indices for slots where CQI is reported
        % as other than NaN
        cqiPracticalPerCWActiveSlots = cqiPracticalPerCW(activeSlotsLogAry);
        cqiPerfectPerCWActiveSlots = cqiPerfectPerCW(activeSlotsLogAry);
        widebandSINRPractical = 10*log10(SINRPerCWPractical(1,activeSlotsLogAry));
        widebandSINRPerfect = 10*log10(SINRPerCWPerfect(1,activeSlotsLogAry));
        xTickLables = num2cell(activeSlotNum(:));

        if isempty(cqiPracticalPerCWActiveSlots)
            disp('No CQI data to plot, because all CQI values are NaNs.');
            return;
        end

        figure();
        subplot(211)
        plot(widebandSINRPerfect,'r^-')
        hold on;
        plot(widebandSINRPractical,'b^-')
        title(['Estimated Wideband SINR Values for Codeword ' num2str(cwIdx)]);
        xlabel('Slots');
        ylabel('Wideband SINR Values in dB');
        xticks(1:numel(cqiPracticalPerCWActiveSlots));
        xticklabels(xTickLables);
        [lowerBound,upperBound] = bounds([widebandSINRPractical(:);widebandSINRPerfect(:)]);
        ylim([lowerBound-1 upperBound+3.5]);
        legend({'Perfect channel est.','Practical channel est.'});

        subplot(212)
        plot(cqiPerfectPerCWActiveSlots,'r-o');
        hold on;
        plot(cqiPracticalPerCWActiveSlots,'b-*');
        xlabel('Slots');
        xticks(1:numel(cqiPracticalPerCWActiveSlots));
        xticklabels(xTickLables);
        ylabel('Wideband CQI');
        [lowerBound,upperBound] = bounds([cqiPracticalPerCWActiveSlots(:);cqiPerfectPerCWActiveSlots(:)]);
        ylim([lowerBound-1 upperBound+3.5]);
        title(['Wideband CQI Values for Codeword ' num2str(cwIdx)]);
        legend({'Perfect channel est.','Practical channel est.'});
    end
end

function plotSubbandCQIAndSINR(subbandCQIPractical,subbandCQIPerfect,SINRPerCWPractical,SINRPerCWPerfect,nslot)
%   plotSUBBANDCQIAndSINR(SUBBANDCQIPRACTICAL,SUBBANDCQIPERFECT,SINRPERCWPRACTICAL,SINRPERCWPERFECT,NSLOT)
%   plots the SINR and CQI values for each codeword across all subbands
%   for practical and perfect channel estimation cases for the given slot
%   NSLOT. The function does not plot the values if CQIMode is
%   'Wideband' or if the CQI and SINR values are all NaNs.

    numCW = size(subbandCQIPractical,2);
    numSubbands = size(subbandCQIPractical,1);
    if numSubbands > 1
        Mode = 'Subband';
    else
        Mode = 'Wideband';
    end

    % Plot subband CQI values
    if strcmpi(Mode,'subband')
        for cwIdx = 1:numCW
            subbandCQIPerCWPractical = subbandCQIPractical(2:end,cwIdx,nslot);
            subbandCQIPerCWPerfect = subbandCQIPerfect(2:end,cwIdx,nslot);
            subbandSINRPerCWPractical = 10*log10(SINRPerCWPractical(2:end,cwIdx,nslot));
            subbandSINRPerCWPerfect = 10*log10(SINRPerCWPerfect(2:end,cwIdx,nslot));

            xTickLables = num2cell(1:numel(subbandCQIPerfect(2:end,cwIdx,1)));
            if all(isnan(subbandCQIPerCWPractical))
                disp(['For the specified slot (' num2str(nslot) '), CQI values are not reported. Please choose another slot number.'])
            else
                figure();
                subplot(211)
                plot(subbandSINRPerCWPerfect,'r^-')
                hold on;
                plot(subbandSINRPerCWPractical,'b^-')
                title(['Estimated Subband SINR Values for Codeword ' num2str(cwIdx)]);
                xlabel('Subbands');
                ylabel('Subband SINR Values in dB');
                xticks(1:numSubbands);
                xticklabels(xTickLables);
                xlim([0 numSubbands+1]);
                [lowerBound,upperBound] = bounds([subbandSINRPerCWPerfect;subbandSINRPerCWPractical]);
                ylim([lowerBound-1 upperBound+3.5]);
                legend({'Perfect channel est.','Practical channel est.'});

                subplot(212)
                plot(subbandCQIPerCWPerfect,'r-o');
                hold on;
                plot(subbandCQIPerCWPractical,'b-*');
                title(['Subband CQI Values for Codeword ' num2str(cwIdx)]);
                xlabel('Subbands');
                ylabel('Subbands CQI');
                xticks(1:numSubbands);
                xticklabels(xTickLables);
                xlim([0 numSubbands+1]);
                [lowerBound,upperBound] = bounds([subbandCQIPerCWPerfect(:,1);subbandCQIPerCWPractical(:,1)]);
                ylim([lowerBound-1 upperBound+3.5]);
                legend({'Perfect channel est.','Practical channel est.'});
            end
        end
    end
end

function plotPMI(pmiPracticalPerSlot,pmiPerfectPerSlot,nslot)
%   plotPMI(PMIPRACTICALPERSLOT,PMIPERFECTPERSLOT,NSLOT) plots the PMI i1
%   indices across all active slots for practical and perfect
%   channel estimation scenarios. The function also plots the i2 indices of
%   practical and perfect channel estimation scenarios across all
%   active slots, when the PMI mode is 'Wideband' or plots i2 indices
%   across all subbands for the specified slot number NSLOT when the
%   PMI mode is 'Subband'.

    % Extract i1 values for slots where PMI is reported as other
    % than NaN
    i1PerfectVals = reshape([pmiPerfectPerSlot(:).i1],3,[])';
    i1PracticalVals = reshape([pmiPracticalPerSlot(:).i1],3,[])';

    % Get logical array to represent active slots in which PMI is
    % reported
    activeSlotsLogAry = ~isnan(i1PerfectVals(:,1));

    % Get slot numbers (0-based) in which PMI is reported
    activeSlotNum = find(activeSlotsLogAry)-1;

    % Get number of active slots
    numActiveSlots = numel(activeSlotNum);

    % Extract wideband PMI indices for slots where PMI is reported
    % as other than NaN
    i1PerfectValsActiveSlots = i1PerfectVals(activeSlotsLogAry,:);
    i1PracticalValsActiveSlots = i1PracticalVals(activeSlotsLogAry,:);

    if isempty(i1PerfectValsActiveSlots)
        disp('No PMI data to plot, because all PMI values are NaNs.');
        return;
    end
    % Get i11 indices
    i11PerfectVals = i1PerfectValsActiveSlots(:,1);
    i11PracticalVals = i1PracticalValsActiveSlots(:,1);

    % Plot i11 values
    figure;
    subplot(411)
    plot(i11PerfectVals,'r-o');
    hold on;
    plot(i11PracticalVals,'b-*');
    xlabel('Slots')
    ylabel('i11 Indices');
    xticks(1:numActiveSlots);
    xTickLables = num2cell(activeSlotNum(:));
    xticklabels(xTickLables);
    [lowerBound,upperBound] = bounds([i11PerfectVals; i11PracticalVals]);
    ylim([lowerBound-2 upperBound+2]);
    title('PMI: i11 Indices')
    legend({'Perfect channel est.','Practical channel est.'});

    % Get i12 indices
    i12PerfectVals = i1PerfectValsActiveSlots(:,2);
    i12PracticalVals = i1PracticalValsActiveSlots(:,2);

    % Plot i12 values
    subplot(412)
    plot(i12PerfectVals,'r-o');
    hold on;
    plot(i12PracticalVals,'b-*');
    xlabel('Slots')
    ylabel('i12 Indices');
    xticks(1:numActiveSlots);
    xticklabels(xTickLables);
    [lowerBound,upperBound] = bounds([i12PerfectVals; i12PracticalVals]);
    ylim([lowerBound-2 upperBound+2]);
    title('PMI: i12 Indices')
    legend({'Perfect channel est.','Practical channel est.'});

    % Get i13 indices
    i13PerfectVals = i1PerfectValsActiveSlots(:,3);
    i13PracticalVals = i1PracticalValsActiveSlots(:,3);

    % Plot i13 values
    subplot(413)
    plot(i13PerfectVals,'r-o');
    hold on;
    plot(i13PracticalVals,'b-*');
    xlabel('Slots')
    ylabel('i13 Indices');
    xticks(1:numActiveSlots);
    xticklabels(xTickLables);
    [lowerBound,upperBound] = bounds([i13PerfectVals; i13PracticalVals]);
    ylim([lowerBound-2 upperBound+2]);
    title('PMI: i13 Indices')
    legend({'Perfect channel est.','Practical channel est.'});

    % Extract i2 values
    numSlots = numel(pmiPracticalPerSlot);
    i2PerfectVals = reshape([pmiPerfectPerSlot(:).i2],[],numSlots)';     % Of size NumSlots-by-NumSubbands
    i2PracticalVals = reshape([pmiPracticalPerSlot(:).i2],[],numSlots)'; % Of size NumSlots-by-NumSubbands

    % Extract i2 values for all active slots where PMI is reported as other than
    % NaN
    i2PerfectValsActiveSlots = i2PerfectVals(activeSlotsLogAry,:);
    i2PracticalValsActiveSlots = i2PracticalVals(activeSlotsLogAry,:);

    numSubbands = size(i2PracticalVals,2);
    if numSubbands > 1
        mode = 'Subband';
    else
        mode = 'Wideband';
    end

    % Plot i2 values
    if strcmpi(mode,'Wideband')
        subplot(414);
        plot(i2PerfectValsActiveSlots,'r-o');
        hold on;
        plot(i2PracticalValsActiveSlots,'b-*');
        title('PMI: i2 Indices')
        xlabel('Slots')
        ylabel('i2 Indices');
        xticks(1:numActiveSlots);
        xticklabels(xTickLables);
        [lowerBound,upperBound] = bounds([i2PerfectValsActiveSlots; i2PracticalValsActiveSlots]);
        ylim([lowerBound-2 upperBound+2]);
        legend({'Perfect channel est.','Practical channel est.'});
    else
        if ~all(isnan(i2PerfectVals(nslot,:)))
            subplot(414);
            plot(i2PerfectVals(nslot,:),'r-o');
            hold on;
            plot(i2PracticalVals(nslot,:),'b-*');
            title(['PMI: i2 Indices for All Subbands in Slot ' num2str(nslot)]);
            xlabel('Subbands')
            ylabel('i2 Indices');
            xticks(1:numSubbands);
            xticklabels(num2cell(1:numSubbands));
            [lowerBound,upperBound] = bounds([i2PerfectVals(nslot,:) i2PracticalVals(nslot,:)]);
            yticks(lowerBound:upperBound);
            yticklabels(num2cell(lowerBound:upperBound));
            xlim([0 size(i2PerfectVals,2)+1])
            ylim([lowerBound-2 upperBound+2]);
            legend({'Perfect channel est.','Practical channel est.'});
        else
            disp(['For the specified slot (' num2str(nslot) '), PMI i2 indices are not reported. Please choose another slot number.'])
        end
    end
end