rng(211);                           % Set RNG state for repeatability

prm.NCellID = 1;                    % Cell ID
prm.FreqRange = 'FR1';              % Frequency range: 'FR1' or 'FR2'
prm.CenterFreq = 3.5e9;             % Hz
prm.SSBlockPattern = 'Case B';      % Case A/B/C/D/E
prm.SSBTransmitted = [ones(1,8) zeros(1,0)];   % 4/8 or 64 in length

prm.TxArraySize = [8 8];            % Transmit array size, [rows cols]
prm.TxAZlim = [-60 60];             % Transmit azimuthal sweep limits
prm.TxELlim = [-90 0];              % Transmit elevation sweep limits

prm.RxArraySize = [2 2];            % Receive array size, [rows cols]
prm.RxAZlim = [-180 180];           % Receive azimuthal sweep limits
prm.RxELlim = [0 90];               % Receive elevation sweep limits

prm.ElevationSweep = false;         % Enable/disable elevation sweep
prm.SNRdB = 30;                     % SNR, dB
prm.RSRPMode = 'SSSwDMRS';          % {'SSSwDMRS', 'SSSonly'}

prm = validateParams(prm);

txBurst.BlockPattern = prm.SSBlockPattern;
txBurst.SSBTransmitted = prm.SSBTransmitted;
txBurst.NCellID = prm.NCellID;
txBurst.SSBPeriodicity = 20;
txBurst.NFrame = 0;
txBurst.Windowing = 0;
txBurst.DisplayBurst = true;

% Assume same subcarrier spacing for carrier as the burst
carrier = nrCarrierConfig('NCellID',prm.NCellID,'NFrame',txBurst.NFrame);
carrier.SubcarrierSpacing = prm.SCS;
carrierInfo = nrOFDMInfo(carrier);
txBurst.SampleRate = carrierInfo.SampleRate;

c = physconst('LightSpeed');   % Propagation speed
lambda = c/prm.CenterFreq;     % Wavelength

prm.posTx = [0;0;0];           % Transmit array position, [x;y;z], meters
prm.posRx = [100;50;0];        % Receive array position, [x;y;z], meters

toRxRange = rangeangle(prm.posTx,prm.posRx);
spLoss = fspl(toRxRange,lambda);    % Free space path loss

% Transmit array
if prm.IsTxURA
    % Uniform rectangular array
    arrayTx = phased.URA(prm.TxArraySize,0.5*lambda, ...
        'Element',phased.IsotropicAntennaElement('BackBaffled',true));
else
    % Uniform linear array
    arrayTx = phased.ULA(prm.NumTx, ...
        'ElementSpacing',0.5*lambda, ...
        'Element',phased.IsotropicAntennaElement('BackBaffled',true));
end

% Receive array
if prm.IsRxURA
    % Uniform rectangular array
    arrayRx = phased.URA(prm.RxArraySize,0.5*lambda, ...
        'Element',phased.IsotropicAntennaElement);
else
    % Uniform linear array
    arrayRx = phased.ULA(prm.NumRx, ...
        'ElementSpacing',0.5*lambda, ...
        'Element',phased.IsotropicAntennaElement);
end

% Scatterer locations
prm.FixedScatMode = true;
if prm.FixedScatMode
    % Fixed single scatterer location
    prm.ScatPos = [50; 80; 0];
else
    % Generate scatterers at random positions
    Nscat = 10;        % Number of scatterers
    azRange = -180:180;
    elRange = -90:90;
    randAzOrder = randperm(length(azRange));
    randElOrder = randperm(length(elRange));
    azAngInSph = azRange(randAzOrder(1:Nscat));
    elAngInSph = elRange(randElOrder(1:Nscat));
    r = 20;            % radius
    [x,y,z] = sph2cart(deg2rad(azAngInSph),deg2rad(elAngInSph),r);
    prm.ScatPos = [x;y;z] + (prm.posTx + prm.posRx)/2;
end

% Configure channel
channel = phased.ScatteringMIMOChannel;
channel.PropagationSpeed = c;
channel.CarrierFrequency = prm.CenterFreq;
channel.SampleRate = txBurst.SampleRate;
channel.SimulateDirectPath = false;
channel.ChannelResponseOutputPort = true;
channel.Polarization = 'None';
channel.TransmitArray = arrayTx;
channel.TransmitArrayPosition = prm.posTx;
channel.ReceiveArray = arrayRx;
channel.ReceiveArrayPosition = prm.posRx;
channel.ScattererSpecificationSource = 'Property';
channel.ScattererPosition = prm.ScatPos;
channel.ScattererCoefficient = ones(1,size(prm.ScatPos,2));

% Get maximum channel delay
[~,~,tau] = channel(complex(randn(txBurst.SampleRate*1e-3,prm.NumTx), ...
    randn(txBurst.SampleRate*1e-3,prm.NumTx)));
maxChDelay = ceil(max(tau)*txBurst.SampleRate);

% Create and display burst information
txBurstInfo = hSSBurstInfo(txBurst);
disp(txBurstInfo);

% Generate burst waveform and grid
[burstWaveform,txBurstGrid] = hSSBurst(txBurst);

% Number of beams at both transmit and receive ends
numBeams = sum(txBurst.SSBTransmitted);

% Transmit beam angles in azimuth and elevation, equi-spaced
azBW = beamwidth(arrayTx,prm.CenterFreq,'Cut','Azimuth');
elBW = beamwidth(arrayTx,prm.CenterFreq,'Cut','Elevation');
txBeamAng = hGetBeamSweepAngles(numBeams,prm.TxAZlim,prm.TxELlim, ...
    azBW,elBW,prm.ElevationSweep);

% For evaluating transmit-side steering weights
SteerVecTx = phased.SteeringVector('SensorArray',arrayTx, ...
    'PropagationSpeed',c);

% Apply steering per OFDM symbol for each SSB
carrier.NSizeGrid = txBurstInfo.NRB;
ofdmInfo = nrOFDMInfo(carrier);
gridSymLengths = repmat(ofdmInfo.SymbolLengths,1, ...
    size(txBurstGrid,2)/length(ofdmInfo.SymbolLengths));
%   repeat burst over numTx to prepare for steering
strTxWaveform = repmat(burstWaveform,1,prm.NumTx)./sqrt(prm.NumTx);
for ssb = 1:length(txBurstInfo.SSBIndex)

    % Extract SSB waveform from burst
    blockSymbols = txBurstInfo.OccupiedSymbols(ssb,:);
    startSSBInd = sum(gridSymLengths(1:blockSymbols(1)-1))+1;
    endSSBInd = sum(gridSymLengths(1:blockSymbols(4)));
    ssbWaveform = strTxWaveform(startSSBInd:endSSBInd,1);

    % Generate weights for steered direction
    wT = SteerVecTx(prm.CenterFreq,txBeamAng(:,ssb));

    % Apply weights per transmit element to SSB
    strTxWaveform(startSSBInd:endSSBInd,:) = ssbWaveform.*(wT');

end

% Receive beam angles in azimuth and elevation, equi-spaced
azBW = beamwidth(arrayRx,prm.CenterFreq,'Cut','Azimuth');
elBW = beamwidth(arrayRx,prm.CenterFreq,'Cut','Elevation');
rxBeamAng = hGetBeamSweepAngles(numBeams,prm.RxAZlim,prm.RxELlim, ...
    azBW,elBW,prm.ElevationSweep);

% For evaluating receive-side steering weights
SteerVecRx = phased.SteeringVector('SensorArray',arrayRx, ...
    'PropagationSpeed',c);

% AWGN level
SNR = 10^(prm.SNRdB/20);                        % Convert to linear gain
N0 = 1/(sqrt(2.0*prm.NumRx*double(ofdmInfo.Nfft))*SNR); % Noise Std. Dev.

% Generate a reference grid for timing correction
%   assumes an SSB in first slot
pssRef = nrPSS(carrier.NCellID);
pssInd = nrPSSIndices;
pbchdmrsRef = nrPBCHDMRS(carrier.NCellID,txBurstInfo.ibar_SSB(1));
pbchDMRSInd = nrPBCHDMRSIndices(carrier.NCellID);
pssGrid = zeros([240 4]);
pssGrid(pssInd) = pssRef;
pssGrid(pbchDMRSInd) = pbchdmrsRef;
refGrid = zeros([12*carrier.NSizeGrid ofdmInfo.SymbolsPerSlot]);
refGrid(txBurstInfo.OccupiedSubcarriers, ...
    txBurstInfo.OccupiedSymbols(1,:)) = pssGrid;

% Loop over all receive beams
rsrp = zeros(numBeams,numBeams);
for rIdx = 1:numBeams

    % Fading channel
    txWave = [strTxWaveform; zeros(maxChDelay,size(strTxWaveform,2))];
    fadWave = channel(txWave);

    % AWGN, after accounting for path loss
    noise = N0*complex(randn(size(fadWave)),randn(size(fadWave)));
    rxWaveform = fadWave*(10^(spLoss/20)) + noise;

    % Generate weights for steered direction
    wR = SteerVecRx(prm.CenterFreq,rxBeamAng(:,rIdx));

    % Apply weights per receive element
    if strcmp(prm.FreqRange, 'FR1')
        strRxWaveform = rxWaveform.*(wR');
    else  % for FR2, combine signal from antenna elements
        strRxWaveform = rxWaveform*conj(wR);
    end

    % Correct timing
    offset = nrTimingEstimate(carrier, ...
        strRxWaveform(1:ofdmInfo.SampleRate*1e-3,:),refGrid*wR(1)');
    if offset > maxChDelay
        offset = 0;
    end
    strRxWaveformS = strRxWaveform(1+offset:end,:);

    % OFDM Demodulate
    rxGrid = nrOFDMDemodulate(carrier,strRxWaveformS);

    % Loop over all SSBs in rxGrid (transmit end)
    for tIdx = 1:numBeams
        % Get each SSB grid
        rxSSBGrid = rxGrid(txBurstInfo.OccupiedSubcarriers, ...
            txBurstInfo.OccupiedSymbols(tIdx,:),:);

        % Make measurements, store per receive, transmit beam
        rsrp(rIdx,tIdx) = measureSSB(rxSSBGrid,prm.RSRPMode,txBurst.NCellID);
    end
end

[m,i] = max(rsrp,[],'all','linear');    % First occurence is output
% i is column-down first (for receive), then across columns (for transmit)
[rxBeamID,txBeamID] = ind2sub([numBeams numBeams],i(1));

% Display the selected beam pair
disp(['Selected Beam pair with RSRP: ' num2str(10*log10(rsrp(rxBeamID, ...
    txBeamID))+30) ' dBm', 13 '  Transmit #' num2str(txBeamID) ...
    ' (Azimuth: ' num2str(txBeamAng(1,txBeamID)) ', Elevation: ' ...
    num2str(txBeamAng(2,txBeamID)) ')' 13 '  Receive #' num2str(rxBeamID) ...
    ' (Azimuth: ' num2str(rxBeamAng(1,rxBeamID)) ', Elevation: ' ...
    num2str(rxBeamAng(2,rxBeamID)) ')' ]);

% Display final beam pair patterns
h = figure('Position',figposition([32 55 32 40]),'MenuBar','none');
h.Name = 'Selected Transmit Array Response Pattern';
wT = SteerVecTx(prm.CenterFreq,txBeamAng(:,txBeamID));
pattern(arrayTx,prm.CenterFreq,'PropagationSpeed',c,'Weights',wT);

h = figure('Position',figposition([32 55 32 40]),'MenuBar','none');
h.Name = 'Selected Receive Array Response Pattern';
wR = SteerVecRx(prm.CenterFreq,rxBeamAng(:,rxBeamID));
pattern(arrayRx,prm.CenterFreq,'PropagationSpeed',c,'Weights',wR);

% Plot MIMO scenario with tx, rx, scatterers, and determined beams
prmScene = struct();
prmScene.txArraySize = prm.TxArraySize;
prmScene.rxArraySize = prm.RxArraySize;
prmScene.txElemPos = getElementPosition(arrayTx); % meters
prmScene.rxElemPos = getElementPosition(arrayRx); % meters
prmScene.txArrayPos = prm.posTx;
prmScene.rxArrayPos = prm.posRx;
prmScene.txAzAngles = -90:90;
prmScene.rxAzAngles = [90:180 -179:-90];
prmScene.scatPos = prm.ScatPos;
prmScene.lambda = lambda;
prmScene.arrayScaling = 1;
hPlotSpatialMIMOScene(prmScene,wT,wR);
if ~prm.ElevationSweep
    view(2);
end

function prm = validateParams(prm)
% Validate user specified parameters and return updated parameters
%
% Only cross-dependent checks are made for parameter consistency.

    if strcmpi(prm.FreqRange,'FR1')
        if prm.CenterFreq > 7.125e9 || prm.CenterFreq < 410e6
            error(['Specified center frequency is outside the FR1 ', ...
                   'frequency range (410 MHz - 7.125 GHz).']);
        end
        if strcmpi(prm.SSBlockPattern,'Case D') ||  ...
           strcmpi(prm.SSBlockPattern,'Case E')
            error(['Invalid SSBlockPattern for selected FR1 frequency ' ...
                'range. SSBlockPattern must be one of ''Case A'' or ' ...
                '''Case B'' or ''Case C'' for FR1.']);
        end
        if ~((length(prm.SSBTransmitted)==4) || ...
             (length(prm.SSBTransmitted)==8))
            error(['SSBTransmitted must be a vector of length 4 or 8', ...
                   'for FR1 frequency range.']);
        end
        if (prm.CenterFreq <= 3e9) && (length(prm.SSBTransmitted)~=4)
            error(['SSBTransmitted must be a vector of length 4 for ' ...
                   'center frequency less than or equal to 3GHz.']);
        end
        if (prm.CenterFreq > 3e9) && (length(prm.SSBTransmitted)~=8)
            error(['SSBTransmitted must be a vector of length 8 for ', ...
                   'center frequency greater than 3GHz and less than ', ...
                   'or equal to 7.125GHz.']);
        end
    else % 'FR2'
        if prm.CenterFreq > 52.6e9 || prm.CenterFreq < 24.25e9
            error(['Specified center frequency is outside the FR2 ', ...
                   'frequency range (24.25 GHz - 52.6 GHz).']);
        end
        if ~(strcmpi(prm.SSBlockPattern,'Case D') || ...
                strcmpi(prm.SSBlockPattern,'Case E'))
            error(['Invalid SSBlockPattern for selected FR2 frequency ' ...
                'range. SSBlockPattern must be either ''Case D'' or ' ...
                '''Case E'' for FR2.']);
        end
        if length(prm.SSBTransmitted)~=64
            error(['SSBTransmitted must be a vector of length 64 for ', ...
                   'FR2 frequency range.']);
        end
    end

    prm.NumTx = prod(prm.TxArraySize);
    prm.NumRx = prod(prm.RxArraySize);
    if prm.NumTx==1 || prm.NumRx==1
        error(['Number of transmit or receive antenna elements must be', ...
               ' greater than 1.']);
    end
    prm.IsTxURA = (prm.TxArraySize(1)>1) && (prm.TxArraySize(2)>1);
    prm.IsRxURA = (prm.RxArraySize(1)>1) && (prm.RxArraySize(2)>1);

    if ~( strcmpi(prm.RSRPMode,'SSSonly') || ...
          strcmpi(prm.RSRPMode,'SSSwDMRS') )
        error(['Invalid RSRP measuring mode. Specify either ', ...
               '''SSSonly'' or ''SSSwDMRS'' as the mode.']);
    end

    % Select SCS based on SSBlockPattern
    switch lower(prm.SSBlockPattern)
        case 'case a'
            scs = 15;
        case {'case b', 'case c'}
            scs = 30;
        case 'case d'
            scs = 120;
        case 'case e'
            scs = 240;
    end
    prm.SCS = scs;

end

function rsrp = measureSSB(rxSSBGrid,mode,NCellID)
% Compute the reference signal received power (RSRP) based on SSS, and if
% selected, also PBCH DM-RS.

    sssInd = nrSSSIndices;                       % SSS indices

    numRx = size(rxSSBGrid,3);
    rsrpSSS = zeros(numRx,1);
    for rxIdx = 1:numRx
        % Extract signals per rx element
        rxSSBGridperRx = rxSSBGrid(:,:,rxIdx);
        rxSSS = rxSSBGridperRx(sssInd);

        % Average power contributions over all REs for RS
        rsrpSSS(rxIdx) = mean(rxSSS.*conj(rxSSS));
    end

    if strcmpi(mode,'SSSwDMRS')
        pbchDMRSInd = nrPBCHDMRSIndices(NCellID);    % PBCH DM-RS indices
        rsrpDMRS = zeros(numRx,1);
        for rxIdx = 1:numRx
            % Extract signals per rx element
            rxSSBGridperRx = rxSSBGrid(:,:,rxIdx);
            rxPBCHDMRS = rxSSBGridperRx(pbchDMRSInd);

            % Average power contributions over all REs for RS
            rsrpDMRS(rxIdx) = mean(rxPBCHDMRS.*conj(rxPBCHDMRS));
        end
    end

    switch lower(mode)
        case 'sssonly'  % Only SSS
           rsrp = max(rsrpSSS);     % max over receive elements
        case 'ssswdmrs' % Both SSS and PBCH-DMRS
           rsrp = max(rsrpSSS+rsrpDMRS)/2; % max over receive elements
    end
end