function [bits, crcBits, eqDataSym, varargout] = trackingVHTDataRecover( ...
    rxVHTData, chanEst, noiseVarEst, cfgVHT, varargin)
%trackingVHTDataRecover Recover bits from VHT Data field signal with pilot tracking
% 
%   [BITS, CRCBITS] = trackingVHTDataRecover(RXVHTDATA, CHANEST, NOISEVAREST,
%   CFGVHTSU) recovers the bits in the VHT-Data field for a VHT format
%   single-user transmission with joint sample rate offset and residual
%   carrier frequency offset tracking. LDPC coding is not supported.
%
%   BITS is an int8 column vector of length 8*CFGVHT.PSDULength containing
%   the recovered information bits.
%
%   CRCBITS is an int8 column vector of length 8 containing the VHT-Data
%   field checksum bits.
%
%   RXVHTDATA is the received time-domain VHT Data field signal, specified
%   as an Ns-by-Nr matrix of real or complex values. Ns represents the
%   number of time-domain samples in the VHT Data field and Nr represents
%   the number of receive antennas. Ns can be greater than the VHT Data
%   field length; in this case additional samples at the end of RXVHTDATA,
%   if not required, are not used. When sample rate offset tracking is
%   enabled using the optional CFGREC argument, additional samples may be
%   required in RXVHTDATA. This is to allow for the receiver running at a
%   higher sample rate than the transmitter and therefore more samples
%   being required.
% 
%   CHANEST is the estimated channel at data and pilot subcarriers based on
%   the VHT-LTF. It is an array of size Nst-by-Nsts-by-Nr, where Nst
%   represents the total number of occupied subcarriers, Nsts represents
%   the total number of space-time streams used for the transmission and Nr
%   is the number of receive antennas.
%
%   NOISEVAREST is the noise variance estimate. It is a nonnegative scalar.
%
%   CFGVHTSU is the format configuration object of type <a 
%   href="matlab:help('wlanVHTConfig')">wlanVHTConfig</a>, which
%   specifies the parameters for the single-user VHT format.
%
%   [BITS, CRCBITS] = trackingVHTDataRecover(RXVHTDATA, CHANEST, NOISEVAREST,
%   CFGVHTMU, USERNUMBER) recovers the bits in the VHT-Data field of a VHT
%   format multi-user transmission for an individual user of interest.
%
%   CFGVHTMU is the VHT format configuration for a multi-user transmission,
%   specified as a <a href="matlab:help('wlanVHTConfig')">wlanVHTConfig</a> object.
%
%   USERNUMBER is the user of interest, specified as an integer between 1
%   and NumUsers, where NumUsers is the number of users in the
%   transmission.
%
%   [BITS, CRCBITS] = trackingVHTDataRecover(RXVHTDATA, CHANEST, NOISEVAREST,
%   CFGVHTSU, USERNUMBER, NUMSTS) recovers the bits in the VHT-Data field
%   of a VHT format multi-user transmission for an individual user of
%   interest.
%   
%   CFGVHTSU is the VHT format configuration for the user of interest,
%   specified as a <a href="matlab:help('wlanVHTConfig')">wlanVHTConfig</a> object.
%
%   NUMSTS is the number of space-time streams, specified as a
%   1-by-NumUsers vector. Element values specify the number of space-time
%   streams per user.
%
%   [BITS, CRCBITS] = trackingVHTDataRecover(..., CFGREC) allows different
%   algorithm options for data recovery via the input CFGREC, which is a
%   <a href="matlab:help('trackingRecoveryConfig')">trackingRecoveryConfig</a> configuration object. When the CFGREC input is not
%   specified, the default property values of the <a href="matlab:help('trackingRecoveryConfig')">trackingRecoveryConfig</a> object
%   are adopted in the recovery. Joint sample rate offset and residual
%   carrier frequency offset tracking is enabled by default.
%
%   [..., EQDATASYM, CPE, PEG] = trackingVHTDataRecover(...) also returns the
%   equalized subcarriers, common phase error, and phase error gradient.
%
%   EQDATASYM is a complex Nsd-by-Nsym-by-Nss array containing the
%   equalized symbols at data carrying subcarriers. Nsd represents the
%   number of data subcarriers, Nsym represents the number of OFDM symbols
%   in the VHT-Data field, and Nss represents the number of spatial
%   streams assigned to the user.
%
%   CPE is a column vector of length Nsym containing the common phase error
%   between each received and expected OFDM symbol.
%
%   PEG is a column vector of length Nsym containing the phase error
%   gradient per OFDM symbol in degrees per subcarrier. This error is
%   caused by a sample rate offset between transmitter and receiver.
%
%   Example: 
%   %  Recover bits in VHT Data field via channel estimation on VHT-LTF 
%   %  over a 2 x 2 quasi-static fading channel
%
%     % Configure a VHT configuration object 
%     chanBW = 'CBW160';
%     cfgVHT = wlanVHTConfig('ChannelBandwidth',    chanBW, ...
%         'NumTransmitAntennas', 2, 'NumSpaceTimeStreams', 2, ...
%         'APEPLength',          512); 
%  
%     % Generate VHT-LTF and VHT Data field signals
%     txDataBits = randi([0 1], 8*cfgVHT.PSDULength, 1);
%     txVHTLTF  = wlanVHTLTF(cfgVHT); 
%     txVHTData = wlanVHTData(txDataBits, cfgVHT);
% 
%     % Pass through a 2 x 2 quasi-static fading channel with AWGN 
%     H = 1/sqrt(2)*complex(randn(2, 2), randn(2, 2));
%     rxVHTLTF  = awgn(txVHTLTF  * H, 10);
%     rxVHTData = awgn(txVHTData * H, 10);
% 
%     % Perform channel estimation based on VHT-LTF
%     demodVHTLTF = wlanVHTLTFDemodulate(rxVHTLTF, cfgVHT, 1);
%     chanEst = wlanVHTLTFChannelEstimate(demodVHTLTF, cfgVHT);
% 
%     % Configure a recovery object using ZF equalization
%     cfgRec = trackingRecoveryConfig('EqualizationMethod', 'ZF'); 
% 
%     % Recover information bits in VHT Data
%     rxDataBits = trackingVHTDataRecover(rxVHTData, chanEst, 0.1, ...
%         cfgVHT, cfgRec);
%
%     % Compare against original information bits
%     disp(isequal(txDataBits, rxDataBits));
%
%   See also trackingRecoveryConfig.

%   Copyright 2016-2019 The MathWorks, Inc.

%#codegen
%#ok<*EMCA>

narginchk(4,7);
nargoutchk(0,5);

% Calculate CPE and/or PEG if requested
calculateCPE = false;
calculatePEG = false;
if nargout>3
    calculateCPE = true;
    if nargout>4
        calculatePEG = true;
    end
end

% VHT configuration input self-validation
validateattributes(cfgVHT, {'wlanVHTConfig'}, {'scalar'}, mfilename, 'VHT format configuration object');

% Optional parameter and case determination - with cfgRec last
if nargin == 7      % (..., cfgVHTSU, userNum, numSTS, cfgRec)
    cfgRecSpec = true;
    cfgRec = varargin{3};

    muSpec = 2;    
    userNum = varargin{1};
    numSTSVec = varargin{2};
elseif nargin == 6  
    if isa(varargin{2}, 'trackingRecoveryConfig') % (..., cfgVHTMU, userNum, cfgRec)
        cfgRecSpec = true;
        cfgRec = varargin{2};
        
        muSpec = 1;
        userNum = varargin{1};    
    else            % (..., cfgVHTSU, userNum, numSTS) 
        cfgRecSpec = false;

        muSpec = 2;
        userNum = varargin{1};
        numSTSVec = varargin{2};        
    end        
elseif nargin == 5  
    if isa(varargin{1}, 'trackingRecoveryConfig') % (..., cfgVHTSU, cfgRec)
        cfgRecSpec = true;
        cfgRec = varargin{1};

        muSpec = 0;
    else            % (..., cfgVHTMU, userNum)
        cfgRecSpec = false;
        muSpec = 1;
        userNum = varargin{1};
    end
else % 4 (..., cfgVHTSU)
    cfgRecSpec = false;
    muSpec = 0;
end

% Validate optional inputs
if muSpec==2 % SU CFGVHT
    validateattributes(userNum, {'numeric'}, {'real','integer','scalar','>=',1,'<=',4}, mfilename, 'USERNUMBER');

    wlan.internal.validateParam('NUMSTS', numSTSVec, mfilename);

    % If UserNum>1, numSTSVec must be a vector
    coder.internal.errorIf(userNum > length(numSTSVec), 'wlan:wlanVHTDataRecover:InvalidUserNum', length(numSTSVec));

    propIdx = 1;     % Have a SU cfgVHT object
    numSTSu = numSTSVec(userNum);
elseif muSpec==1 % MU CFGVHT
    validateattributes(userNum, {'numeric'}, {'real','integer','scalar','>=',1,'<=',cfgVHT.NumUsers}, mfilename, 'USERNUMBER');

    % Have a MU cfgVHT object as input
    numSTSVec = cfgVHT.NumSpaceTimeStreams;
    propIdx = userNum;
    numSTSu = numSTSVec(propIdx);
else % not specified, set defaults
    % Single-user case
    userNum     = 1;
    propIdx     = 1;
    numSTSVec   = cfgVHT.NumSpaceTimeStreams;
    numSTSu     = numSTSVec(propIdx);
end

if cfgRecSpec
    validateattributes(cfgRec, {'trackingRecoveryConfig'}, {'scalar'}, mfilename, 'recovery configuration object');

    symOffset = cfgRec.OFDMSymbolOffset;
    pilotTracking = cfgRec.PilotTracking;
    eqMethod = cfgRec.EqualizationMethod;
    pilotTrackingWindow = cfgRec.PilotTrackingWindow;
else    % set defaults
    symOffset = 0.75;
    eqMethod = 'MMSE';
    pilotTracking = 'Joint';
    pilotTrackingWindow = 9;
end
cfgTrack = struct('calculateCPE', calculateCPE, 'pilotTracking', pilotTracking, 'pilotTrackingWindow', pilotTrackingWindow, 'pilotGainTracking', false);

cfgInfo = validateConfig(cfgVHT, 'MCS');
mcsTable = wlan.internal.getRateTable(cfgVHT);
chanBW = cfgVHT.ChannelBandwidth;

% NDP only for SU, so idx is (1)
if cfgVHT.APEPLength(1) == 0 
    bits     = zeros(0, 1, 'int8');
    crcBits  = zeros(0, 1, 'int8');
    eqDataSym = zeros(mcsTable.NSD(1), 0, mcsTable.Nss(1));
    if calculateCPE==true
        varargout{1} = []; % CPE
    end
    if calculatePEG==true
        varargout{2} = []; % PEG
    end
    return;
end

% All optional params: parsed and validated
numSTSTotal = sum(numSTSVec);

% Signal input self-validation
validateattributes(rxVHTData, {'double'}, {'2d','finite'}, 'rxVHTData', 'VHT-Data field signal'); 
validateattributes(chanEst, {'double'}, {'3d','finite'}, 'chanEst', 'channel estimation'); 
validateattributes(noiseVarEst, {'double'}, {'real','scalar','nonnegative','finite'}, 'noiseVarEst', 'noise variance estimation'); 

% Set up some implicit configuration parameters
numBPSCS   = mcsTable.NBPSCS(propIdx);    % Number of coded bits per single carrier
numCBPS    = mcsTable.NCBPS(propIdx);     % Number of coded bits per OFDM symbol
numES      = mcsTable.NES(propIdx);       % Number of encoded streams
numSS      = mcsTable.Nss(propIdx);       % Number of spatial streams
numSeg     = strcmp(chanBW, 'CBW160') + 1;
% Number of coded bits per OFDM symbol, per spatial stream, per segment
numCBPSSI  = numCBPS/numSS/numSeg;  
numRx      = size(rxVHTData, 2);

% Get OFDM configuration
[cfgOFDM, dataInd, pilotInd] = wlan.internal.wlanGetOFDMConfig(chanBW, cfgVHT.GuardInterval, 'VHT', numSTSTotal);
  
% Set channel coding
coder.varsize('channelCoding',[1,4]);
channelCoding = getChannelCoding(cfgVHT);

% Validate only BCC channel coding supported 
coder.internal.errorIf(strcmp(channelCoding{propIdx},'LDPC'),'wlan:trackingVHTDataRecover:LDPCNotSupported');

% Cross-validation between inputs
if muSpec==2
    coder.internal.errorIf(cfgVHT.NumSpaceTimeStreams(1) ~= numSTSu, ...
    'wlan:wlanVHTDataRecover:InvalidNumSTS', numSTSu, cfgVHT.NumSpaceTimeStreams(1));
end
coder.internal.errorIf(cfgVHT.STBC && muSpec > 0, 'wlan:wlanVHTDataRecover:InvalidSTBCMU');

numST = numel([dataInd; pilotInd]); % Total number of occupied subcarriers
coder.internal.errorIf(size(chanEst, 1) ~= numST, 'wlan:wlanVHTDataRecover:InvalidChanEst1D', numST);
coder.internal.errorIf(size(chanEst, 2) ~= numSTSTotal, 'wlan:wlanVHTDataRecover:InvalidChanEst2D', numSTSTotal);
coder.internal.errorIf(size(chanEst, 3) ~= numRx, 'wlan:wlanVHTDataRecover:InvalidChanEst3D');

% Extract data and pilot subcarriers from channel estimate
chanEstData = chanEst(dataInd,:,:);
chanEstPilots = chanEst(pilotInd,:,:);

% Cross-validation between inputs
numOFDMSym = cfgInfo.NumDataSymbols;
minInputLen = numOFDMSym*(cfgOFDM.FFTLength+cfgOFDM.CyclicPrefixLength);
coder.internal.errorIf(size(rxVHTData, 1) < minInputLen, 'wlan:wlanVHTDataRecover:ShortDataInput', minInputLen);

% Index into streams for the user of interest
stsIdx = sum(numSTSVec(1:(userNum-1)))+(1:numSTSu); 

% OFDM demodulation and optional pilot tracking
fnRefPilots = @()getRefPilots(chanBW, numOFDMSym, numSTSVec, stsIdx);
[ofdmDemod, cpe, peg] = trackingOFDMDemodulate(rxVHTData, chanEstPilots(:,stsIdx,:), fnRefPilots, numOFDMSym, symOffset, cfgOFDM, cfgTrack);
if calculateCPE
    varargout{1} = cpe; 
end
if calculatePEG
    varargout{2} = peg;
end

% Equalization
if cfgVHT.STBC  % Only SU
    [eqDataSym, dataCSI] = wlan.internal.wlanSTBCCombine(ofdmDemod(dataInd,:,:), chanEstData, numSS, eqMethod, noiseVarEst);
else    % Both SU and MU
    [eqDataSym, dataCSI] = wlan.internal.wlanEqualize(ofdmDemod(dataInd,:,:), chanEstData(:,stsIdx,:), eqMethod, noiseVarEst);
end

% Segment parsing of symbols
parserOut = wlanSegmentParseSymbols(eqDataSym, chanBW);  % [Nsd/Nseg Nsym Nss Nseg]
csiParserOut = wlanSegmentParseSymbols(reshape(dataCSI, [], 1, numSS), chanBW); % [Nsd/Nseg 1 Nss Nseg]
 
% Constellation demapping
qamDemodOut = wlanConstellationDemap(parserOut, noiseVarEst, numBPSCS); % [Ncbpssi,Nsym,Nss,Nseg]

% Apply bit-wise CSI and concatenate OFDM symbols in the first dimension
qamDemodOut = bsxfun(@times, ...
        reshape(qamDemodOut, numBPSCS, [], numOFDMSym, numSS, numSeg), ...
        reshape(csiParserOut, 1, [], 1, numSS, numSeg));
qamDemodOut = reshape(qamDemodOut, [], numSS, numSeg); % [(Ncbpssi*Nsym),Nss,Nseg]

% BCC Deinterleaving
deintlvrOut = wlanBCCDeinterleave(qamDemodOut, 'VHT', numCBPSSI, chanBW); % [(Ncbpssi*Nsym),Nss,Nseg]

% Segment deparsing of bits
segDeparserOut = wlanSegmentDeparseBits(deintlvrOut, chanBW, numES, numCBPS, numBPSCS); % [(Ncbpss*Nsym),Nss]

% Stream deparsing
streamDeparserOut = wlanStreamDeparse(segDeparserOut(:,:), numES, numCBPS, numBPSCS); % [(Ncbps*Nsym/Nes),Nes]
% Indexing for codegen

% Channel decoding for BCC
numTailBits = 6;
chanDecOutPreDeparse = wlanBCCDecode(streamDeparserOut, mcsTable.Rate(propIdx));
% BCC decoder deparser
chanDecOut = reshape(chanDecOutPreDeparse(1:end-numTailBits,:)', [], 1);

% Derive initial state of the scrambler 
scramInit = wlan.internal.scramblerInitialState(chanDecOut(1:7));

% Remove pad and tail bits, and descramble
if all(scramInit==0)
    % Scrambler initialization invalid (0), therefore do not descramble
    descramBits = chanDecOut(1:16+8*cfgVHT.PSDULength(propIdx));
else
    descramBits = wlanScramble(chanDecOut(1:16+8*cfgVHT.PSDULength(propIdx)), scramInit);
end

% Outputs
crcBits = descramBits(9:16);
bits = descramBits(17:end);

end

function ref = getRefPilots(chanBW,numOFDMSym,numSTSVec,stsIdx)
    % Get reference pilots, from Eqn 22-95, IEEE Std 802.11ac-2013
    % Offset by 4 to allow for L-SIG, VHT-SIG-A, VHT-SIG-B pilot symbols
    n = (0:numOFDMSym-1).';
    z = 4;
    refPilots = wlan.internal.vhtPilots(n, z, chanBW, sum(numSTSVec));
    ref = refPilots(:,:,stsIdx); % Extract for MU
end