% Release 15 NR-TM test models for FR1 and FR2
% Select one of the NR-TMs among:
% "NR-FR1-TM1.1","NR-FR1-TM1.2","NR-FR1-TM2",...
% "NR-FR1-TM2a","NR-FR1-TM3.1","NR-FR1-TM3.1a",...
% "NR-FR1-TM3.2","NR-FR1-TM3.3","NR-FR2-TM1.1",...
% "NR-FR2-TM2","NR-FR2-TM3.1"
nrtm = "NR-FR1-TM3.2";

% Select the NR-TM waveform parameters
bw   = "10MHz"; % Channel bandwidth
scs  = "30kHz"; % Subcarrier spacing
dm   = "FDD";   % Duplexing mode

fprintf('TM = %s\n', nrtm);

tmwavegen = hNRReferenceWaveformGenerator(nrtm,bw,scs,dm);

[txWaveform,tmwaveinfo,resourcesinfo] = generateWaveform(tmwavegen);

% Visualize the associated PRB and subcarrier resource grids
displayResourceGrid(tmwavegen);

[pdschArray, targetRNTIs] = hListTargetPDSCHs(tmwavegen.Config,resourcesinfo.WaveformResources);

phaseNoiseOn = true;
nonLinearityModelOn = true;

txWaveform = txWaveform/max(abs(txWaveform));

txWaveform = repmat(txWaveform,2,1);

if phaseNoiseOn
    sr = tmwaveinfo.Info.SamplingRate;

    % Carrier frequency
    if tmwavegen.Config.FrequencyRange == "FR1" % carrier frequency for FR1
        fc = 4e9;
    else % carrier frequency for FR2
        fc = 28e9;
    end

    % Phase noise level
    foffsetLog = (3:0.1:log10(sr/2)); % model offset from 1e3Hz to sr/2
    foffset = 10.^foffsetLog;    % linear freq offset
    PN_dBc_Hz = PNmodelPoleZero(foffset,fc);
    figure; semilogx(foffset,PN_dBc_Hz); xlabel('Frequency offset (Hz)'); ylabel('dBc/Hz'); title('Phase noise model'); grid on

    % Apply phase noise to waveform
    pnoise = comm.PhaseNoise('FrequencyOffset',foffset,'Level',PN_dBc_Hz,'SampleRate',sr);
    rxWaveform = pnoise(txWaveform);
else
    rxWaveform = txWaveform;
end

if nonLinearityModelOn
    rapp = comm.MemorylessNonlinearity('Method','Rapp model');
    rapp.Smoothness = 1.55;
    rapp.OutputSaturationLevel = 1;

    % Plot non-linear characteristic
     plotNonLinearCharacteristic(rapp);

    % Apply nonlinearity
    rxWaveform = rapp(rxWaveform);
end

if dm == "FDD"
    nFrames = 1;
else % TDD
    nFrames = 2;
end
rxWaveform(1:nFrames*tmwaveinfo.Info.SamplesPerSubframe*10) = [];

gnb.NRB = tmwavegen.Config.Carriers.NRB;
gnb.CyclicPrefix = tmwavegen.Config.BWP.CyclicPrefix;
gnb.SubcarrierSpacing = tmwavegen.Config.Carriers.SubcarrierSpacing;

refGrid = referenceGrid(tmwaveinfo,pdschArray);
NSlot = 0;
[offset,mag] = nrTimingEstimate(rxWaveform,gnb.NRB,gnb.SubcarrierSpacing,NSlot,refGrid);
tmWaveformSync = rxWaveform(1+offset:end,:);

rxGrid = hOFDMDemodulate(gnb, tmWaveformSync);

% Get total number of slots and the number of subcarriers and symbols per
% slot
symbolsPerSlot = tmwaveinfo.Info.SymbolsPerSlot;
totalNrSlots = floor(size(rxGrid,2)/symbolsPerSlot);
noSubcarriers = size(rxGrid,1);

% Declare storage variables
constellationSymbols = [];  % equalized symbols for constellation plot
constellationRef = [];      % reference symbols for constellation plot
refSqGrid = []; % Grid with magnitude square of reference symbols
evSqGrid = [];  % Grid with magnitude square of error vector (symbols - reference)

for NSlot=0:totalNrSlots-1
    % Extract grid for current slot
    currentGrid = rxGrid(:,NSlot*symbolsPerSlot+(1:symbolsPerSlot),:);

    % Retrieve resources
    [pdschIndices,refSymbols,dmrsIndices,dmrsSymbols] = hSlotResources(pdschArray,NSlot);

    % Do not include first two slot symbols for PDSCH EVM (used for control
    % TS 38.141-1 table 4.9.2.2-2)
    idx = pdschIndices <= 2*noSubcarriers;
    pdschIndices(idx) = [];
    refSymbols(idx) = [];

    % Channel estimation
    [estChannelGrid,~] = nrChannelEstimate(currentGrid,dmrsIndices,dmrsSymbols);

    % Get PDSCH resource elements from the received grid
    [pdschRx,pdschHest] = nrExtractResources(pdschIndices,currentGrid,estChannelGrid);

    % Equalization: set noiseEst to 0 for zero-forcing equalization
    noiseEst = 0;
    [pdschEq,csi] = nrEqualizeMMSE(pdschRx,pdschHest,noiseEst);

    refSqGridSlot = zeros(size(currentGrid)); % slot grid magnitude square of reference symbols
    evSqGridSlot = zeros(size(currentGrid));  % slot grid magnitude square of error vector
    if ~isempty(refSymbols) % for slots with data

        % Error vector magnitude square
        evSq = abs(refSymbols-pdschEq).^2;
        % Reference symbols magnitude square
        refSymSq = abs(refSymbols).^2;

        % Store constellation symbols, reference symbols and square error vector
        constellationSymbols = [constellationSymbols; pdschEq];
        constellationRef = [constellationRef; refSymbols];
        evSqGridSlot(pdschIndices) = evSq;
        refSqGridSlot(pdschIndices) = refSymSq;

    end
    refSqGrid = [refSqGrid refSqGridSlot];
    evSqGrid = [evSqGrid evSqGridSlot];

end

figure
plot(constellationSymbols,'.'); hold on
plot(constellationRef,'+')
title('Equalized symbols constellation')
grid on; xlabel('In-Phase'); ylabel('Quadrature');

% EVM per symbol
[evmPerSymbol, evmMaxSymbol] = evm(evSqGrid,refSqGrid);

figure; subplot(3,1,1)
plot(0:length(evmPerSymbol)-1,evmPerSymbol,'.-',0:length(evmPerSymbol)-1,evmMaxSymbol,'.:');
title('EVM vs OFDM symbol')
grid on; xlabel('Symbol number'); ylabel('EVM (%)'); legend('rms EVM','peak EVM','Location','bestoutside')

% EVM per slot
% reshape grids to one column per slot
evSqGridSlot = reshape(evSqGrid,size(evSqGrid).*symbolsPerSlot.^[1 -1]);
refSqGridSlot = reshape(refSqGrid,size(refSqGrid).*symbolsPerSlot.^[1 -1]);

[evmPerSlot, evmMaxSlot] = evm(evSqGridSlot,refSqGridSlot);

subplot(3,1,2)
slotNo = 0:length(evmPerSlot)-1;
plot(slotNo,evmPerSlot,'.-',slotNo,evmMaxSlot,'.:');
title('EVM vs slot')
grid on; xlabel('Slot number'); ylabel('EVM (%)'); legend('rms EVM','peak EVM','Location','bestoutside')

% EVM per subcarrier
[evmPerSubcarrier, evmMaxSubcarrier] = evm(evSqGrid.',refSqGrid.');

subplot(3,1,3)
subcarrierNo = 0:length(evmPerSubcarrier)-1;
plot(subcarrierNo,evmPerSubcarrier,'.-',subcarrierNo,evmMaxSubcarrier,'.:')
title('EVM vs subcarrier')
grid on; xlabel('Subcarrier number'); ylabel('EVM (%)'); legend('rms EVM','max EVM','Location','bestoutside')

[evmRMS, evmMax] = evm(evSqGrid(:),refSqGrid(:));
disp("PDSCH RNTIs: " + mat2str(targetRNTIs))
disp("RMS EVM = " + string(evmRMS) + "%");
disp("Max EVM = " + string(evmMax) + "%");

function [evmRMS, evmMax] = evm(evSqGrid,refSqGrid)
% Calculates the evm per column of the input matrix
% inputs (both of the same size):
%   - evSqGrid = matrix of error vector magnitude square |refSymb-rxSymb|.^2
%   - refSqGrid = matrix of ref symbols magnitude square |refSymb|.^2

    if size(evSqGrid)~=size(refSqGrid)
        error('Input matrices must have the same size');
    end

    % RMS EVM (%)
    evmRMS = rmsEVM(sum(evSqGrid,1),sum(refSqGrid,1));

    % max EVM (%)
    evmMax = zeros(1,size(refSqGrid,2));
    % find non zero REs in reference grid
    nonZeroREs = (refSqGrid~=0);
    for ncol = 1:length(evmMax)
        % remove reference symbols with value zero
        nzEvSqREsSlot = evSqGrid(nonZeroREs(:,ncol),ncol);
        nzRefSqREsSlot = refSqGrid(nonZeroREs(:,ncol),ncol);
        evmMax(ncol) = maxEVM(nzEvSqREsSlot,nzRefSqREsSlot);
    end
end

function evmRMS = rmsEVM(evSq,refSq)
% Calcuates RMS EVM. Inputs
%   - evSq = matrix of error vector magnitude square
%   - refSq = matrix of ref symbols magnitude square
% They have the same size or refSq can be a scalar.
    evmRMS = 100*sqrt(evSq./refSq);
end

function evmMax = maxEVM(evSq,refSq)
% inputs (both of the same size):
%   - evSq = array of error vector magnitude square
%   - refSq = array of ref symbols magnitude square
    if ~isempty(evSq) && ~isempty(refSq)
        evmMax = max(100*sqrt(evSq./mean(refSq)));
    else
        evmMax = NaN;
    end
end

function refGrid = referenceGrid(tmwaveinfo,pdschArray)
% Create a reference grid for synchronization. This should span a whole
% frame, starting in slot 0. It contains the DM-RSs specified in
% pdschArray.

    L = tmwaveinfo.Info.SymbolsPerSubframe*10; % symbols in a frame
    refGrid = zeros(tmwaveinfo.Info.NSubcarriers,L); % empty grid
    rbsPerSlot = tmwaveinfo.Info.NSubcarriers*tmwaveinfo.Info.SymbolsPerSlot;
    % Polpulate the DM-RSs in the reference grid for all slots
    for NSlot=0:(10*tmwaveinfo.Info.SlotsPerSubframe)-1 % 1 frame worth of subframes
        [~,~,dmrsIndices,dmrsSymbols] = hSlotResources(pdschArray,NSlot);
        refGrid(dmrsIndices+NSlot*rbsPerSlot) = dmrsSymbols;
    end

end

function PN_dBC_Hz = PNmodelPoleZero(f,fc)
% Generate the phase noise characteristic in dBc/Hz for the frequency
% offset values specified by vector f for the carrier frequency fc.
%
% The model used here is the multi-pole/zero model as proposed in:
% Yinan Qi et al. "On the Phase Tracking Reference Signal (PT-RS) Design
% for 5G New Radio (NR)". Vehicular Technology Conference (VTC-Fall), Aug
% 2018.

% Pole/zeros and PSD0 as specified in the paper mentioned above and in 3GPP
% R1-163984, "Discussion on phase noise modeling". This corresponds to
% "parameter set A" for a base frequency of 30GHz.
fcBase = 30e9;
fz= [1.8 2.2 40]*1e6;
fp = [0.1 0.2 8]*1e6;
PSD0 = -79.4;

% Compute numerator
num = ones(size(f));
for ii=1:numel(fz)
    num = num.* ( 1 +(f./fz(ii)).^2);
end

% Compute denominator
den = ones(size(f));
for ii=1:numel(fz)
    den = den.* ( 1 +(f./fp(ii)).^2);
end

% Compute phase noise and apply a shift for carrier frequencies different
% from the base frequency.
PN_dBC_Hz = 10*log10(num./den) + PSD0 + 20*log10(fc/fcBase);

end

function plotNonLinearCharacteristic(memoryLessNonlinearity)
% Plot the non linear characteristic of the PA impairment represented by
% memoryLessNonlinearity system object. This input parameter is a
% comm.MemorylessNonlinearity system object.

% input samples
x = complex((1/sqrt(2))*(-1+2*rand(1000,1)),(1/sqrt(2))*(-1+2*rand(1000,1)));

% Non linearity
yRapp = memoryLessNonlinearity(x);
% release object so we can feed it different number of samples
release(memoryLessNonlinearity);

% Plot characteristic
figure; plot(10*log10(abs(x).^2),10*log10(abs(x).^2)); hold on; grid on
plot(10*log10(abs(x).^2),10*log10(abs(yRapp).^2),'.');
xlabel('Input power (dBW)'); ylabel('Output power (dBW)'); title('nonlinearity impairment')
legend('Linear characteristic', 'Rapp nonlinearity','Location','Northwest');

end