%hNBCellSearch Cell identity search using PSS and SSS
%   [NNCELLID,OFFSET,PEAK] = hNBCellSearch(...) returns the narrowband cell
%   identity NNCELLID carried by the NPSS and NSSS signals in the input
%   waveform, the timing offset OFFSET of the first frame head in that
%   waveform, and the peak correlation magnitude PEAK.
%
%   [NNCELLID,OFFSET,PEAK] = hNBCellSearch(WAVEFORM) returns the cell
%   identity NNCELLID, timing offset OFFSET, and peak correlation magnitude
%   PEAK given a waveform WAVEFORM.
%
%   The input WAVEFORM must be a T-by-P matrix where T is the number of
%   time domain samples and P is the number of receive antennas. The
%   waveform provided must contain the NPSS and NSSS signals for the cell
%   search to be successful.
%
%   NNCELLID is a scalar indicating the detected cell identity. OFFSET
%   indicates the number of samples from the start of the waveform WAVEFORM
%   to the position in that waveform where the first frame begins. PEAK is
%   a scalar indicating the peak magnitude of the correlation used to
%   detect a cell; a complete correlation output is available as the
%   output argument CORR from <a href="matlab:help('lteNBDLFrameOffset')
%   ">lteNBDLFrameOffset</a>.
%
%   [NNCELLID,OFFSET,PEAK] = hNBCellSearch(WAVEFORM,ALG) provides control
%   over the cell search algorithm as specified in the structure ALG, which
%   may include the following fields:
%   SSSDetection    - Optional. SSS detection method.
%                     ('PreFFT'(default),'PostFFT')
%   MaxCellCount    - Optional. The number of cell identities to detect
%                     (1...504) (default 1)
%   CellIDs         - Optional. A vector of cell identities for which to 
%                     constrain the cell search (default vector 0:503)
%
%   Unlike LTE, NPSS does not contain any information related to cell
%   identity. NPSS is purely used for timing synchronization and is
%   achieved by time-domain correlation using <a
%   href="matlab:help('lteNBDLFrameOffset')">lteNBDLFrameOffset</a>. Entire
%   cell identity is conveyed by NSSS and can be obtained using two methods
%   - "PreFFT" or "PostFFT". PreFFT NSSS detection operates in time domain
%   and cell identity is established by time-domain correlation by
%   searching for cell identities whose signal has high correlation with
%   the received waveform. PreFFT method is much slower than PostFFT.
%   'PostFFT' NSSS detection operates in the frequency domain. OFDM
%   demodulation is performed using the timing estimate from NPSS
%   detection. The demodulated NSSS Resource Elements (REs) are correlated
%   with all possible NSSS sequences to find the cell identity. For
%   'PostFFT' NSSS detection, the peak correlation magnitude PEAK is the
%   sum of the peak correlation magnitudes from time-domain NPSS detection
%   and frequency-domain NSSS detection. When ALG.MaxCellCount>1, NNCELLID,
%   OFFSET and PEAK are vectors, with each element corresponding to one
%   cell. If ALG.CellIDs is absent, the output vectors are sorted by
%   decreasing correlation peak magnitude i.e. decreasing PEAK value. If
%   ALG.CellIDs is present and ALG.MaxCellCount = numel(ALG.CellIDs), the
%   output vectors are in the same order as the cell identities in
%   ALG.CellIDs; this allows for monitoring of the PEAK output for a
%   pre-determined set of cells.
%
%   Example:
%   % Search for the cell identity (in this case 120) of an NB-IoT waveform.
%   
%   ngen = NBIoTDownlinkWaveformGenerator;
%   ngen.Config.NNCellID = 120;
%   ngen.Config.NBRefP = 2;
%   ngen.Config.TotSubframes = 650; %65 frames
%   [eNodeBOutput,~,ofdmInfo] = generateWaveform(ngen);
%   alg.SSSDetection = 'PostFFT';
%   [NNCellID, offset, peak] = hNBCellSearch(eNodeBOutput, alg)
% 
%   See also lteDLFrameOffset, lteFrequencyOffset, lteFrequencyCorrect,
%   lteOFDMDemodulate.

%   Copyright 2019-2020 The MathWorks, Inc.

function [cellIDs,offsets,peaks] = hNBCellSearch(waveform,varargin)

    narginchk(1,2);

    allcellids=0:503;
    strengthSort=true;
    defaultMaxCellCount=1;
    defaultSSSDetection='PreFFT';
    if (nargin==1)
        alg.CellIDs=allcellids;
        alg.MaxCellCount=defaultMaxCellCount;
        alg.SSSDetection=defaultSSSDetection;
    else
        alg=varargin{1};
        if (~isfield(alg,'MaxCellCount'))
            alg.MaxCellCount=defaultMaxCellCount;
        end
        if (~isfield(alg,'SSSDetection'))
            alg.SSSDetection=defaultSSSDetection;
        end
        if (~isfield(alg,'CellIDs'))
            alg.CellIDs=allcellids;
        else
            if(alg.MaxCellCount>=numel(alg.CellIDs))
                strengthSort=false;
            end
        end
        if (isempty(alg.CellIDs))
            error('lte:error','The function call resulted in an error: The cell identity constraint input must not be empty');
        end
    end
    
    enb = struct(); % Declaring an empty structure
    if (~isfield(enb,'CyclicPrefix')) % Though this is not needed for NB-IoT, keeping intact for enhancing lteCellSearch function in future
        enb.CyclicPrefix='Normal';
    end
    if (~isfield(enb,'DuplexMode')) % Though this is not needed for NB-IoT, keeping intact for enhancing lteCellSearch function in future
        enb.DuplexMode='FDD';
    end
    if (~isfield(enb,'OperationMode'))
        enb.OperationMode = 'Standalone';
    end
    
    if (~isfield(enb,'NFrame'))
        enb.NFrame = 0;
    end
    
    if (~any(strcmpi(alg.SSSDetection,{'PreFFT','PostFFT'})))        
        error('lte:error','The function call (hNBCellSearch) resulted in an error: For the parameter field SSSDetection, (%s) is not one of the set (PreFFT, PostFFT)',alg.SSSDetection);
    end
    postFFT=strcmpi(alg.SSSDetection,'PostFFT');
    
    if (alg.MaxCellCount<1 || alg.MaxCellCount>504)        
        error('lte:error','The function call (hNBCellSearch) resulted in an error: For the parameter field MaxCellCount, the value (%d) must be within the range [1,504]',alg.MaxCellCount);
    end
    
    % Form set of IDs and ID groups to search.
    alg.CellIDs=unique(alg.CellIDs,'stable');
    ids=0; % Though this is not needed for NB-IoT, keeping intact for enhancing lteCellSearch function in future
    idGroups=alg.CellIDs; %In NB-IoT case, entire cellIDs are on SSS
    
    % Primary cell search.
    enb.NDLRB = 6; % For setting cyclic prefix length from OFDM Info.
    ofdmInfo=lteOFDMInfo(enb); 
    subframeLength=ofdmInfo.SamplingRate*1e-3;
    frameLength=subframeLength*10;
    cpLength=double(ofdmInfo.CyclicPrefixLengths(2));
    corrcfg.NPSS='On';
    corrcfg.NSSS='Off';
    corrcfg.NRS='Off';
    % Determine the number of different NPSS peaks 'nPSSPeaks' to record for
    % each primary cell identity. For pre-FFT SSS detection, only record
    % one peak as SSS timing is independent; for post-FFT SSS detection,
    % record ALG.MaxCellCount peaks as SSS detection is performed after
    % OFDM demodulation using PSS timing.
    if (postFFT)
        nPSSPeaks=alg.MaxCellCount;
    else
        nPSSPeaks=1;
    end
    p_offsets=zeros(nPSSPeaks,1);
    p_peaks=zeros(nPSSPeaks,1);
    idx=1;
    enb.NCellID=0;
    [p_offsets(idx),corr]=lteNBDLFrameOffset(enb,waveform,corrcfg);
    if (postFFT)
        % lteNBDLFrameOffset returns the offset to the first NPSS
        % occurrence. The offset will be up to the length of a frame. For
        % post-FFT NSSS detection, take the offset modulo the length of a
        % frame to account for a potential negative offset.
        p_offsets(idx)=mod(p_offsets(idx),frameLength);
    end
    if (nPSSPeaks==1)
        % Record peak correlation value.
        p_peaks(idx)=max(corr(:));
    else
        % Perform frame averaging across half frames. In NB-IoT case,
        % averaging might be needed across entire frame.
        nFrames=floor(size(corr,1)/(frameLength));
        if (nFrames>0)
            corr=corr(1:nFrames*frameLength,:);
            corr=permute(mean(reshape(corr,frameLength,nFrames,size(corr,2)),2),[1 3 2]);
        end
        % Take the maximum across antennas.
        corr=max(corr,[],2);
        % Detect nPeaks peaks and associated timing offsets. For
        % multiple peaks closer together than the cyclic prefix length,
        % only the largest peak is detected.
        [foundpeaks,foundoffsets]=findpeaks(circshift(corr,floor(cpLength/2)),'NPeaks',nPSSPeaks,'SortStr','descend','MinPeakDistance',cpLength);
        foundoffsets=mod(foundoffsets-floor(cpLength/2)-1,frameLength);
        p_peaks(idx+(0:length(foundpeaks)-1))=foundpeaks;
        p_offsets(idx+(0:length(foundoffsets)-1))=foundoffsets;
    end
    
    % If only detecting one cell identity, only retain the single strongest
    % primary cell search result.
    if(alg.MaxCellCount==1)
        ids=ids(p_peaks==max(p_peaks));
        ids=ids(1);
        p_offsets=p_offsets(p_peaks==max(p_peaks));        
        p_peaks=max(p_peaks);
    else
        ids=zeros(nPSSPeaks,1);
        ids=ids(:).';        
    end        
    
    % Prepare cached SSS indices.
    if (postFFT)
        nsssInd=lteNSSSIndices(setfield(enb,'NSubframe',9)); %#ok<SFLD>
    end
    
    % Narrowband Secondary cell search.
    corrcfg.NPSS='Off';
    corrcfg.NSSS='On';
    s_peaks=zeros(numel(ids),numel(idGroups));
    for p=1:numel(ids)
        if (postFFT) 
            % Perform timing synchronization according to the PSS timing,
            % frequency offset estimation and correction, and OFDM
            % demodulation.
            timesynced=waveform(1+p_offsets(p):end,:);
            if(size(timesynced,1)<subframeLength)
                timesynced = [timesynced; zeros(subframeLength-size(timesynced,1),size(timesynced,2))]; %#ok<AGROW>
            end
            
            % Following commented code is not needed for the case currently
            % supported in NB-IoT.
            % if (strcmpi(enb.DuplexMode,'TDD'))
            %     enb.TDDConfig=0;
            %     enb.SSC=0;
            %     enb.NSubframe=0;
            % end
            
            % NB-IoT downlink OFDM modulation has a half subcarrier shift
            % just like the uplink SCFDMA modulation of LTE. Therefore call
            % lteFrequencyOffset having the field NULRB rather than NDLRB
            % so that frequency offset is properly computed. Also, use
            % lteSCFDMADemodulate function for OFDM demodulation in NB-IoT
            % for similar reason i.e., NB-IoT downlink OFDM modulation has
            % a half subcarrier shift just like the uplink SCFDMA
            % modulation of LTE.
            enb=rmfield(enb,'NDLRB');
            enb.NULRB=6; % For frequency offset computation
            enb.CyclicPrefixUL = 'Normal';
            foffset=lteFrequencyOffset(enb,timesynced,0);
            freqsynced=lteFrequencyCorrect(enb,timesynced,foffset);
            enb.NBULSubcarrierSpacing = '15kHz';
            rxgrid = lteSCFDMADemodulate(enb,freqsynced);% For OFDM demodulation in NB-IoT
            enb=rmfield(enb,'NULRB');
            enb.NDLRB=6; % Revert back the configuration after frequency offset compensation and OFDM demodulation
        end
        for s=1:numel(idGroups)
            enb.NNCellID=idGroups(s); % In NB-IoT, NSSS carries the entire Cell ID information
            if (postFFT)                
                s_peaks(p,s)=0;
                for nframe = 0:3
                    % NSSS is mapped onto every even numbered frames and
                    % NSSS values repeat for every 8 frames. So, 4
                    % contiguous frames will have different NSSS values.
                    % Current loop Searches in which of these 4 frames the
                    % input is in.
                    
                    % Perform NSSS detection using subframe 9 which carries NSSS.
                    tempenb = enb;
                    tempenb.NSubframe = 9;
                    tempenb.NFrame = 2*nframe;
                    sssSym9 = lteNSSS(tempenb); %For NB-IoT
                    % Perform channel estimation for the NSSS.
                    hestSSS=chestSSS(enb,rxgrid,nsssInd,sssSym9);
                    % Coherent combining of SSS REs in a given subframe and
                    % antenna.
                    combined=mean(hestSSS,3);
                    % Non-coherent combining across all subframes and
                    % antennas.
                    s_peaks(p,s)=max(s_peaks(p,s),sqrt(mean(abs(combined(:)).^2)));
                end
                % Adjust post-FFT NSSS correlation estimate to make it
                % comparable with peak correlation magnitudes obtained
                % in the time domain.
                s_peaks(p,s)=s_peaks(p,s)*double(ofdmInfo.Nfft+cpLength)*62/double(ofdmInfo.Nfft^2);
                if(size(hestSSS,1)>1)
                    s_peaks(p,s)=s_peaks(p,s)*2;
                end
            else
                [~,corr]=lteNBDLFrameOffset(enb,waveform,corrcfg);
                s_peaks(p,s)=max(corr(:));
            end
        end
    end
    
    % Establish overall peaks for the combinations of IDs and ID groups
    % configured for search.
    peaks=repmat(p_peaks,1,size(s_peaks,2))+s_peaks;
    if(numel(peaks)~=(numel(alg.CellIDs)*nPSSPeaks))
        peaksCellIDs=repmat(ids.',1,numel(idGroups))+repmat(idGroups,numel(ids),1)*3;
        validIdx=arrayfun(@(x)find(peaksCellIDs==x),alg.CellIDs,'UniformOutput',false);
        validIdx=cat(1,validIdx{:}).';
        peaks=peaks(validIdx);
    else
        peaks=peaks(:).';
    end
    
    % Establish the number of identities N to return, the minimum of the
    % configured cell identity constraint or ALG.MaxCellCount
    N=min(numel(alg.CellIDs),alg.MaxCellCount);
    
    % Sort results by strength unless ALG.CellIDs was configured, in which
    % case results are sorted by cell identity
    if (strengthSort)
        [~,idx]=sort(peaks,'descend');
        alg.CellIDs=repmat(alg.CellIDs,nPSSPeaks,1);
        alg.CellIDs=alg.CellIDs(:).';
    else
        peaks=max(reshape(peaks,nPSSPeaks,numel(peaks)/nPSSPeaks),[],1);
        idx=1:numel(alg.CellIDs);
    end
    
    % Return the N biggest peaks and the corresponding cell identities. 
    peaks=peaks(idx(1:N));
    cellIDs=alg.CellIDs(idx(1:N));
    
    % Return the corresponding timing offsets by rerunning timing
    % synchronization with the final cell identities; the NPSS and NSSS
    % correlation are both enabled.
    offsets=arrayfun(@(x)lteNBDLFrameOffset(setfield(enb,'NNCellID',x),waveform),cellIDs);
    
end

function ssshest=chestSSS(enb,rxgrid,sssInd,sym)
    
    dims=lteDLResourceGridSize(enb,1);
    L=dims(2);
    nSubframes=size(rxgrid,2)/L;
    N=floor(nSubframes/10);
    R=size(rxgrid,3);
    ssshest=zeros(N,R,numel(sssInd));
    % Run all correlations; correlations are indexed
    % into 'ssshest' with 'idx' and subframes are extracted from the
    % received grid using 'nsf' which is calculated from 'idx'.
    for idx=1:N
        nsf=(idx-1)*10+9;
        for r=1:R
            subframe=rxgrid(:,(nsf*L)+(1:L),r);
            ssshest(idx,r,:)=subframe(sssInd).*conj(sym);
        end
    end
end
