% Set the NR-TM parameters for the receiver
nrtm = "NR-FR1-TM3.1"; % Reference channel
bw   = "10MHz";        % Channel bandwidth
scs  = "30kHz";        % Subcarrier spacing
dm   = "FDD";          % Duplexing mode

% Set parameters for the spectrum analyzer
centerFrequency = 3.4e9; 
sampleRate = 15.36e6;
measurementTime = 20e-3;
mechanicalAttenuation = 0; %dB
startFrequency = 3.39e9; 
stopFrequency = 3.41e9;
resolutionBandwidth = 220e3; 
videoBandwidth = 220000;

foundVISA = visadevlist;
resourceID = foundVISA(foundVISA.Model == "N9030A",:).ResourceName;
resourceID = resourceID(contains(resourceID,"N9030A"));             % Extract resourceID which contains "N9030A"
sigAnalyzerObj = visadev(resourceID);
sigAnalyzerObj.ByteOrder = "big-endian";
configureCallback(sigAnalyzerObj,"byte",85e6,@callbackFcn) 
sigAnalyzerObj.Timeout = 20;

writeline(sigAnalyzerObj,"*RST");
instrumentInfo = writeread(sigAnalyzerObj,"*IDN?");
fprintf("Instrument identification information: %s",instrumentInfo);

% Set up signal analyzer mode to basic IQ mode
writeline(sigAnalyzerObj,":INSTrument:SELect BASIC");

% Set the center frequency
writeline(sigAnalyzerObj,strcat(":SENSe:FREQuency:CENTer ",num2str(centerFrequency)));

% Set the capture sample rate
writeline(sigAnalyzerObj,strcat(":SENSe:WAVeform:SRATe ",num2str(sampleRate)));

% Turn off averaging
writeline(sigAnalyzerObj,":SENSe:WAVeform:AVER OFF");

% Set the spectrum analyzer to take one single measurement after the trigger line goes high
writeline(sigAnalyzerObj,":INIT:CONT OFF"); 

% Set the trigger to external source 1 with positive slope triggering
writeline(sigAnalyzerObj,":TRIGger:WAVeform:SOURce IMMediate");
writeline(sigAnalyzerObj,":TRIGger:LINE:SLOPe POSitive");

% Set the time for which measurement needs to be made
writeline(sigAnalyzerObj,strcat(":WAVeform:SWE:TIME ",num2str(measurementTime)));

% Turn off electrical attenuation 
writeline(sigAnalyzerObj,":SENSe:POWer:RF:EATTenuation:STATe OFF");

% Set the mechanical attenuation level
writeline(sigAnalyzerObj,strcat(":SENSe:POWer:RF:ATTenuation ",num2str(mechanicalAttenuation)));

% Turn IQ signal ranging to auto
writeline(sigAnalyzerObj,":SENSe:VOLTage:IQ:RANGe:AUTO ON");

% Set the endianness of returned data
writeline(sigAnalyzerObj,":FORMat:BORDer NORMal");

% Set the format of the returned data
writeline(sigAnalyzerObj,":FORMat:DATA REAL,64");

% Trigger the instrument and initiate measurement
writeline(sigAnalyzerObj,"*TRG");
writeline(sigAnalyzerObj,":INITiate:WAVeform");

% Wait until measure operation is complete
measureComplete = writeread(sigAnalyzerObj,"*OPC?");

% Read the IQ data
writeline(sigAnalyzerObj,":READ:WAV0?");
data = readbinblock(sigAnalyzerObj,"double");

% Separate the data and build the complex IQ vector
inphase = data(1:2:end);
quadrature = data(2:2:end);
rxWaveform = inphase+1i*quadrature;

writeline(sigAnalyzerObj,":FETCH:WAV1?");
signalSpec = readbinblock(sigAnalyzerObj,"double");

% Display the measurement information
captureSampleRate = 1/signalSpec(1);
fprintf("Sample Rate (Hz) = %s",num2str(captureSampleRate));

fprintf("Number of points read = %s",num2str(signalSpec(4)));

fprintf("Max value of signal (dBm) = %s",num2str(signalSpec(6)));

fprintf("Min value of signal (dBm) = %s",num2str(signalSpec(7)));

% Ensure rxWaveform is a column vector
if ~iscolumn(rxWaveform)
    rxWaveform = rxWaveform.';
end

% Plot the power spectral density (PSD) of the acquired signal
spectrumPlotRx = dsp.SpectrumAnalyzer;
spectrumPlotRx.SampleRate =  captureSampleRate;
spectrumPlotRx.SpectrumType = "Power density";
spectrumPlotRx.PowerUnits =  "dBm";
spectrumPlotRx.Window = "Hamming";
spectrumPlotRx.SpectralAverages = 10;
spectrumPlotRx.YLimits = [-140 -90];
spectrumPlotRx.YLabel = "PSD";
spectrumPlotRx.ShowLegend = false;
spectrumPlotRx.Title = "Received Signal Spectrum: 10 MHz 5G NR-TM Waveform";
spectrumPlotRx(rxWaveform);

% Switch back to the spectrum analyzer view
writeline(sigAnalyzerObj,":INSTrument:SELect SA");

% Set the mechanical attenuation level
writeline(sigAnalyzerObj,strcat(":SENSe:POWer:RF:ATTenuation ",num2str(mechanicalAttenuation)));

% Set the center frequency, RBW, and VBW
writeline(sigAnalyzerObj,strcat(":SENSe:FREQuency:CENTer ",num2str(centerFrequency)));
writeline(sigAnalyzerObj,strcat(":SENSe:FREQuency:STARt ",num2str(startFrequency)));
writeline(sigAnalyzerObj,strcat(":SENSe:FREQuency:STOP ",num2str(stopFrequency)));
writeline(sigAnalyzerObj,strcat(":SENSe:BANDwidth:RESolution ",num2str(resolutionBandwidth)));
writeline(sigAnalyzerObj,strcat(":SENSe:BANDwidth:VIDeo ",num2str(videoBandwidth)));

% Enable continuous measurement on the spectrum analyzer
writeline(sigAnalyzerObj,":INIT:CONT ON"); 

% Begin receiving the over-the-air signal
writeline(sigAnalyzerObj,"*TRG");

clear sigAnalyzerObj;

tmwavegen = hNRReferenceWaveformGenerator(nrtm,bw,scs,dm);
[~,tmwaveinfo,resourcesInfo] = generateWaveform(tmwavegen);

frequencyCorrectionRange = -100e3:1e3:100e3;
[rxWaveform, coarseOffset] = DMRSFrequencyCorrection(rxWaveform,captureSampleRate,frequencyCorrectionRange,tmwavegen,resourcesInfo);
fprintf("Coarse frequency offset = %.0f Hz", coarseOffset)

frequencyCorrectionRange = -100:5:100;
[rxWaveform, fineOffset] = DMRSFrequencyCorrection(rxWaveform,captureSampleRate,frequencyCorrectionRange,tmwavegen,resourcesInfo);
fprintf("Fine frequency offset = %.1f Hz", fineOffset)

cfg = struct();
cfg.PlotEVM = true;                 % Plot EVM statistics 
cfg.DisplayEVM = true;              % Print EVM statistics
cfg.Label = nrtm;                   % Set to TM name of captured waveform
cfg.SampleRate = captureSampleRate; % Use sample rate during capture

[evmInfo,eqSym,refSym] = hNRPDSCHEVM(tmwavegen.Config,rxWaveform,cfg);

function [correctedWaveform,appliedFrequencyCorrection] = DMRSFrequencyCorrection(waveform,sampleRate,frequencyCorrectionRange,tmwavegen,resourcesInfo)
    % waveform - Waveform to be corrected. Needs to be a Nx1 column vector.
    % sampleRate - Sample rate of waveform
    % frequencyCorrectioRange - Range and granularity at which frequency
    % correction is inspected
    % tmwavegen and resourcesInfo - Outputs of generateWaveform method 
    [pdschArray,~,carrier] = hListTargetPDSCHs(tmwavegen.Config,resourcesInfo.WaveformResources);
    bwpCfg = tmwavegen.Config.BandwidthParts{1,1};
    nSlots = carrier.SlotsPerFrame;
    
    % Generate a reference grid spanning 10 ms (one frame). This grid
    % contains only the DM-RS and is used for synchronization.
    refGrid = referenceGrid(carrier,bwpCfg,pdschArray,nSlots);
    
    % Apply frequency offsets to the waveform as specified by
    % freuqnecyCorrectionRange.
    nSamples = (0:length(waveform)-1)';
    frequencyShift = (2*pi*frequencyCorrectionRange.*nSamples)./sampleRate;
    
    % Each column represents an offset waveform.
    offsetWaveforms = waveform.*exp(1j*frequencyShift);
    
    [~,mag] = nrTimingEstimate(offsetWaveforms,carrier.NSizeGrid,...
        carrier.SubcarrierSpacing,nSlots,refGrid, ...
        "SampleRate",sampleRate);
    
    % Find the frequency at which the DM-RS correlation is at a maximum.
    [~,index] = max(max(mag));    
    appliedFrequencyCorrection = frequencyCorrectionRange(index);
    correctedWaveform = offsetWaveforms(:,index);
end

function refGrid = referenceGrid(carrier,bwpCfg,pdschArray,nSlots)
    % Create a reference grid for the required number of slots. The grid
    % contains the DM-RS symbols specified in pdschArray. The function
    % returns REFGRID of dimensions K-by-S-by-L, where K is the number of
    % subcarriers of size carrier.NSizeGrid, S is the number of symbols
    % spanning nSlots, and L is the number of layers.

    nSubcarriers = carrier.NSizeGrid * 12;
    L = carrier.SymbolsPerSlot*nSlots;                           % Number of OFDM symbols in the reference grid
    nLayers = size(pdschArray(1).Resources(1).ChannelIndices,2);
    bwpStart = bwpCfg.NStartBWP;
    bwpLen = bwpCfg.NSizeBWP;
    refGrid = zeros(nSubcarriers,L,nLayers);                     % Empty grid
    bwpGrid = zeros(bwpLen*12,L,nLayers);
    rbsPerSlot = bwpLen*12*carrier.SymbolsPerSlot;

    % Populate the DM-RS symbols in the reference grid for all slots. Place
    % bwpGrid in a carrier grid (at an appropriate location) in case the
    % BWP size is not the same as the carrier grid
    for slotIdx = carrier.NSlot + (0:nSlots-1)
        [~,~,dmrsIndices,dmrsSymbols] = hSlotResources(pdschArray,slotIdx);
        if ~isempty(dmrsIndices)
            for layerIdx = 1:nLayers
                if layerIdx <= size(dmrsIndices,2)
                    dmrsIndices(:,layerIdx) = dmrsIndices(:,layerIdx) - rbsPerSlot*(layerIdx -1) + (L*bwpLen*12*(layerIdx-1));
                    bwpGrid(dmrsIndices(:,layerIdx)+(slotIdx-carrier.NSlot)*rbsPerSlot) = dmrsSymbols(:,layerIdx);
                end
            end
            refGrid(12*bwpStart+1:12*(bwpStart+bwpLen),:,:) = bwpGrid;
        end
    end
end