% A mixed OFDMA and MU-MIMO configuration is defined for an HE-MU packet.
% The allocation index 17 defines two 52-tone RUs, with one user in each
% RU, and one 106-tone RU. The 106-tone RU has two users in a MU-MIMO
% configuration.
cfgMU = wlanHEMUConfig(17);
cfgMU.NumTransmitAntennas = 2;

% Configure RU 1 and user 1
cfgMU.RU{1}.SpatialMapping = 'Direct';
cfgMU.User{1}.STAID = 1;
cfgMU.User{1}.APEPLength = 1e3;
cfgMU.User{1}.MCS = 5;
cfgMU.User{1}.NumSpaceTimeStreams = 2;
cfgMU.User{1}.ChannelCoding = 'LDPC';

% Configure RU 2 and user 2
cfgMU.RU{2}.SpatialMapping = 'Fourier';
cfgMU.User{2}.STAID = 2;
cfgMU.User{2}.APEPLength = 500;
cfgMU.User{2}.MCS = 4;
cfgMU.User{2}.NumSpaceTimeStreams = 1;
cfgMU.User{2}.ChannelCoding = 'BCC';

% Configure RU 3 and user 1
cfgMU.RU{3}.SpatialMapping = 'Fourier';
cfgMU.User{3}.STAID = 3;
cfgMU.User{3}.APEPLength = 100;
cfgMU.User{3}.MCS = 2;
cfgMU.User{3}.NumSpaceTimeStreams = 1;
cfgMU.User{3}.ChannelCoding = 'BCC';

% Configure RU 3 and user 2
cfgMU.User{4}.STAID = 4;
cfgMU.User{4}.APEPLength = 500;
cfgMU.User{4}.MCS = 3;
cfgMU.User{4}.NumSpaceTimeStreams = 1;
cfgMU.User{4}.ChannelCoding = 'LDPC';

% Specify propagation channel
numRx = 2; % Number of receive antennas
delayProfile = 'Model-D'; % TGax channel delay profile

% Specify impairments
noisePower = -40; % Noise power to apply in dBW
cfo = 62e3; % Carrier frequency offset Hz

% Generate waveform
rx = heSigRecGenerateWaveform(cfgMU,numRx,delayProfile,noisePower,cfo);

chanBW = cfgMU.ChannelBandwidth; % Assume channel bandwidth is known
sr = wlanSampleRate(cfgMU); % Sample rate

% Specify pilot tracking method for recovering the data field. This can be:
%  'Joint' - use joint common phase error and sample rate offset tracking
%  'CPE' - use only common phase error tracking
% When recovering 26-tone RUs only CPE tracking is used as the joint
% tracking algorithm is susceptible to noise.
pilotTracking = 'Joint';

% Create an HE recovery configuration object and set the channel bandwidth
cfgRx = wlanHERecoveryConfig;
cfgRx.ChannelBandwidth = chanBW;

% The recovery configuration object is used to get the start and end
% indices of the pre-HE-SIG-B field.
ind = wlanFieldIndices(cfgRx);

% Setup plots for the example
[spectrumAnalyzer,timeScope,ConstellationDiagram,EVMPerSubcarrier,EVMPerSymbol] = heSigRecSetupPlots(sr);

% Minimum packet length is 10 OFDM symbols
lstfLength = double(ind.LSTF(2));
minPktLen = lstfLength*5; % Number of samples in L-STF

rxWaveLen = size(rx,1);

searchOffset = 0; % Offset from start of waveform in samples
while (searchOffset + minPktLen) <= rxWaveLen
    % Packet detection
    pktOffset = wlanPacketDetect(rx,chanBW,searchOffset);

    % Adjust packet offset
    pktOffset = searchOffset + pktOffset;
    if isempty(pktOffset) || (pktOffset + ind.LSIG(2) > rxWaveLen)
        error('** No packet detected **');
    end

    % Coarse frequency offset estimation and correction using L-STF
    rxLSTF = rx(pktOffset+(ind.LSTF(1):ind.LSTF(2)), :);
    coarseFreqOffset = wlanCoarseCFOEstimate(rxLSTF,chanBW);
    rx = helperFrequencyOffset(rx,sr,-coarseFreqOffset);

    % Symbol timing synchronization
    searchBufferLLTF = rx(pktOffset+(ind.LSTF(1):ind.LSIG(2)),:);
    pktOffset = pktOffset+wlanSymbolTimingEstimate(searchBufferLLTF,chanBW);

    % Fine frequency offset estimation and correction using L-STF
    rxLLTF = rx(pktOffset+(ind.LLTF(1):ind.LLTF(2)),:);
    fineFreqOffset = wlanFineCFOEstimate(rxLLTF,chanBW);
    rx = helperFrequencyOffset(rx,sr,-fineFreqOffset);

    % Timing synchronization complete: packet detected
    fprintf('Packet detected at index %d\n',pktOffset + 1);

    % Display estimated carrier frequency offset
    cfoCorrection = coarseFreqOffset + fineFreqOffset; % Total CFO
    fprintf('Estimated CFO: %5.1f Hz\n\n',cfoCorrection);

    break; % Front-end processing complete, stop searching for a packet
end

% Scale the waveform based on L-STF power (AGC)
gain = 1./(sqrt(mean(rxLSTF.*conj(rxLSTF))));
rx = rx.*gain;

rxLLTF = rx(pktOffset+(ind.LLTF(1):ind.LLTF(2)),:);
lltfDemod = wlanLLTFDemodulate(rxLLTF,chanBW);
lltfChanEst = wlanLLTFChannelEstimate(lltfDemod,chanBW);
noiseVar = helperNoiseEstimate(lltfDemod);

disp('Detect packet format...');
rxSIGA = rx(pktOffset+(ind.LSIG(1):ind.HESIGA(2)),:);
pktFormat = wlanFormatDetect(rxSIGA,lltfChanEst,noiseVar,chanBW);
fprintf('  %s packet detected\n\n',pktFormat);

% Set the packet format in the recovery object and update the field indices
cfgRx.PacketFormat = pktFormat;
ind = wlanFieldIndices(cfgRx);

lltfDemod = wlanHEDemodulate(rxLLTF,'L-LTF',chanBW);
lltfChanEst = wlanLLTFChannelEstimate(lltfDemod,chanBW);

disp('Decoding L-SIG... ');
% Extract L-SIG and RL-SIG fields
rxLSIG = rx(pktOffset+(ind.LSIG(1):ind.RLSIG(2)),:);

% OFDM demodulate
helsigDemod = wlanHEDemodulate(rxLSIG,'L-SIG',chanBW);

% Estimate CPE and phase correct symbols
helsigDemod = preHECommonPhaseErrorTracking(helsigDemod,lltfChanEst,'L-SIG',chanBW);

% Estimate channel on extra 4 subcarriers per subchannel and create full
% channel estimate
preheInfo = wlanHEOFDMInfo('L-SIG',chanBW);
preHEChanEst = preHEChannelEstimate(helsigDemod,lltfChanEst,preheInfo.NumSubchannels);

% Average L-SIG and RL-SIG before equalization
helsigDemod = mean(helsigDemod,2);

% Equalize data carrying subcarriers, merging 20 MHz subchannels
[eqLSIGSym,csi] = preHESymbolEqualize(helsigDemod(preheInfo.DataIndices,:,:), ...
    preHEChanEst(preheInfo.DataIndices,:,:),noiseVar,preheInfo.NumSubchannels);

% Decode L-SIG field
[~,failCheck,lsigInfo] = wlanLSIGBitRecover(eqLSIGSym,noiseVar,csi);

if failCheck
    disp(' ** L-SIG check fail **');
else
    disp(' L-SIG check pass');
end
% Get the length information from the recovered L-SIG bits and update the
% L-SIG length property of the recovery configuration object
lsigLength = lsigInfo.Length;
cfgRx.LSIGLength = lsigLength;

% Measure EVM of L-SIG symbols
EVM = comm.EVM;
EVM.ReferenceSignalSource = 'Estimated from reference constellation';
EVM.Normalization = 'Average constellation power';
EVM.ReferenceConstellation = wlanReferenceSymbols('BPSK');
rmsEVM = EVM(eqLSIGSym);
fprintf(' L-SIG EVM: %2.2fdB\n\n',20*log10(rmsEVM/100));

% Calculate the receive time and corresponding number of samples in the
% packet
RXTime = ceil((lsigLength + 3)/3) * 4 + 20; % In microseconds
numRxSamples = round(RXTime * 1e-6 * sr);   % Number of samples in time

fprintf(' RXTIME: %dus\n',RXTime);
fprintf(' Number of samples in the packet: %d\n\n',numRxSamples);

sampleOffset = max((-lstfLength + pktOffset),1); % First index to plot
sampleSpan = numRxSamples + 2*lstfLength; % Number samples to plot
% Plot as much of the packet (and extra samples) as we can
plotIdx = sampleOffset:min(sampleOffset + sampleSpan,rxWaveLen);

% Configure timeScope to display the packet
timeScope.TimeSpan = sampleSpan/sr;
timeScope.TimeDisplayOffset = sampleOffset/sr;
timeScope.YLimits = [0 max(abs(rx(:)))];
timeScope(abs(rx(plotIdx,:)));
release(timeScope);

% Display the spectrum of the detected packet
spectrumAnalyzer(rx(pktOffset + (1:numRxSamples),:));
release(spectrumAnalyzer);

disp('Decoding HE-SIG-A...')
rxSIGA = rx(pktOffset+(ind.HESIGA(1):ind.HESIGA(2)),:);
sigaDemod = wlanHEDemodulate(rxSIGA,'HE-SIG-A',chanBW);
hesigaDemod = preHECommonPhaseErrorTracking(sigaDemod,preHEChanEst,'HE-SIG-A',chanBW);

% Equalize data carrying subcarriers, merging 20 MHz subchannels
preheInfo = wlanHEOFDMInfo('HE-SIG-A',chanBW);
[eqSIGASym,csi] = preHESymbolEqualize(hesigaDemod(preheInfo.DataIndices,:,:), ...
                                      preHEChanEst(preheInfo.DataIndices,:,:), ...
                                      noiseVar,preheInfo.NumSubchannels);
% Recover HE-SIG-A bits
[sigaBits,failCRC] = wlanHESIGABitRecover(eqSIGASym,noiseVar,csi);

% Perform the CRC on HE-SIG-A bits
if failCRC
    disp(' ** HE-SIG-A CRC fail **');
else
    disp(' HE-SIG-A CRC pass');
end

% Measure EVM of HE-SIG-A symbols
release(EVM);
if strcmp(pktFormat,'HE-EXT-SU')
    % The second symbol of an HE-SIG-A field for an HE-EXT-SU packet is
    % QBPSK.
    EVM.ReferenceConstellation = wlanReferenceSymbols('BPSK',[0 pi/2 0 0]);
    % Account for scaling of L-LTF for an HE-EXT-SU packet
    rmsEVM = EVM(eqSIGASym*sqrt(2));
else
    EVM.ReferenceConstellation = wlanReferenceSymbols('BPSK');
    rmsEVM = EVM(eqSIGASym);
end
fprintf(' HE-SIG-A EVM: %2.2fdB\n\n',20*log10(mean(rmsEVM)/100));

cfgRx = interpretHESIGABits(cfgRx,sigaBits);
ind = wlanFieldIndices(cfgRx); % Update field indices

disp(cfgRx)

if strcmp(pktFormat,'HE-MU')
    fprintf('Decoding HE-SIG-B...\n');
    if ~cfgRx.SIGBCompression
        fprintf(' Decoding HE-SIG-B common field...\n');
        s = getSIGBLength(cfgRx);
        % Get common field symbols. The start of HE-SIG-B field is known
        rxSym = rx(pktOffset+(ind.HESIGA(2)+(1:s.NumSIGBCommonFieldSamples)),:);
        % Decode HE-SIG-B common field
        [status,cfgRx] = heSIGBCommonFieldDecode(rxSym,preHEChanEst,noiseVar,cfgRx);

        % CRC on HE-SIG-B content channels
        if strcmp(status,'Success')
            fprintf('  HE-SIG-B (common field) CRC pass\n');
        elseif strcmp(status,'ContentChannel1CRCFail')
            fprintf('  ** HE-SIG-B CRC fail for content channel-1\n **');
        elseif strcmp(status,'ContentChannel2CRCFail')
            fprintf('  ** HE-SIG-B CRC fail for content channel-2\n **');
        elseif any(strcmp(status,{'UnknownNumUsersContentChannel1','UnknownNumUsersContentChannel2'}))
            error('  ** Unknown packet length, discard packet\n **');
        else
            % Discard the packet if all HE-SIG-B content channels fail
            error('  ** HE-SIG-B CRC fail **');
        end
        % Update field indices as the number of HE-SIG-B symbols are
        % updated
        ind = wlanFieldIndices(cfgRx);
    end

    % Get complete HE-SIG-B field samples
    rxSIGB = rx(pktOffset+(ind.HESIGB(1):ind.HESIGB(2)),:);
    fprintf(' Decoding HE-SIG-B user field... \n');
    % Decode HE-SIG-B user field
    [failCRC,cfgUsers] = heSIGBUserFieldDecode(rxSIGB,preHEChanEst,noiseVar,cfgRx);

    % CRC on HE-SIG-B users
    if ~all(failCRC)
        fprintf('  HE-SIG-B (user field) CRC pass\n\n');
        numUsers = numel(cfgUsers);
    elseif all(failCRC)
        % Discard the packet if all users fail the CRC
        error('  ** HE-SIG-B CRC fail for all users **');
    else
        fprintf('  ** HE-SIG-B CRC fail for at least one user\n **');
        % Only process users with valid CRC
        numUsers = numel(cfgUsers);
    end

else % HE-SU, HE-EXT-SU
    cfgUsers = {cfgRx};
    numUsers = 1;
end

cfgDataRec = trackingRecoveryConfig;
cfgDataRec.PilotTracking = pilotTracking;

fprintf('Decoding HE-Data...\n');
for iu = 1:numUsers
    % Get recovery configuration object for each user
    user = cfgUsers{iu};
    if strcmp(pktFormat,'HE-MU')
        fprintf(' Decoding User:%d, STAID:%d, RUSize:%d\n',iu,user.STAID,user.RUSize);
    else
        fprintf(' Decoding RUSize:%d\n',user.RUSize);
    end

    heInfo = wlanHEOFDMInfo('HE-Data',chanBW,user.GuardInterval,[user.RUSize user.RUIndex]);

    % HE-LTF demodulation and channel estimation
    rxHELTF = rx(pktOffset+(ind.HELTF(1):ind.HELTF(2)),:);
    heltfDemod = wlanHEDemodulate(rxHELTF,'HE-LTF',chanBW,user.GuardInterval, ...
        user.HELTFType,[user.RUSize user.RUIndex]);
    [chanEst,pilotEst] = heLTFChannelEstimate(heltfDemod,user);

    % Number of expected data OFDM symbols
    symLen = heInfo.FFTLength+heInfo.CPLength;
    numOFDMSym = (ind.HEData(2)-ind.HEData(1)+1)/symLen;

    % HE-Data demodulation with pilot phase and timing tracking
    % Account for extra samples when extracting data field from the packet
    % for sample rate offset tracking. Extra samples may be required if the
    % receiver clock is significantly faster than the transmitter.
    maxSRO = 120; % Parts per million
    Ne = ceil(numRxSamples*maxSRO*1e-6); % Number of extra samples
    Ne = min(Ne,rxWaveLen-numRxSamples); % Limited to length of waveform
    numRxSamplesProcess = numRxSamples+Ne;
    rxData = rx(pktOffset+(ind.HEData(1):numRxSamplesProcess),:);
    if user.RUSize==26
        % Force CPE only tracking for 26-tone RU as algorithm susceptible
        % to noise
        cfgDataRec.PilotTracking = 'CPE';
    else
        cfgDataRec.PilotTracking = pilotTracking;
    end
        [demodSym,cpe,peg] = heTrackingOFDMDemodulate(rxData,chanEst,numOFDMSym,user,cfgDataRec);

    % Estimate noise power in HE fields
    demodPilotSym = demodSym(heInfo.PilotIndices,:,:);
    nVarEst = heNoiseEstimate(demodPilotSym,pilotEst,user);

    % Equalize
    [eqSym,csi] = heEqualizeCombine(demodSym,chanEst,nVarEst,user);

    % Discard pilot subcarriers
    eqSymUser = eqSym(heInfo.DataIndices,:,:);
    csiData = csi(heInfo.DataIndices,:);

    % Demap and decode bits
    rxPSDU = wlanHEDataBitRecover(eqSymUser,nVarEst,csiData,user,'LDPCDecodingMethod','layered-bp');

    % Deaggregate the A-MPDU
    [mpduList,~,status] = wlanAMPDUDeaggregate(rxPSDU,wlanHESUConfig);
    if strcmp(status,'Success')
        fprintf('  A-MPDU deaggregation successful \n');
    else
        fprintf('  A-MPDU deaggregation unsuccessful \n');
    end

    % Decode the list of MPDUs and check the FCS for each MPDU
    for i = 1:numel(mpduList)
        [~,~,status] = wlanMPDUDecode(mpduList{i},wlanHESUConfig,'DataFormat','octets');
        if strcmp(status,'Success')
            fprintf('  FCS pass for MPDU:%d\n',i);
        else
            fprintf('  FCS fail for MPDU:%d\n',i);
        end
    end

    % Plot equalized constellation of the recovered HE data symbols for all
    % spatial streams per user
    hePlotEQConstellation(eqSymUser,user,ConstellationDiagram,iu,numUsers);

    % Measure EVM of HE-Data symbols
    release(EVM);
    EVM.ReferenceConstellation = wlanReferenceSymbols(user);
    rmsEVM = EVM(eqSymUser(:));
    fprintf('  HE-Data EVM:%2.2fdB\n\n',20*log10(rmsEVM/100));

    % Plot EVM per symbol of the recovered HE data symbols
    hePlotEVMPerSymbol(eqSymUser,user,EVMPerSymbol,iu,numUsers);

    % Plot EVM per subcarrier of the recovered HE data symbols
    hePlotEVMPerSubcarrier(eqSymUser,user,EVMPerSubcarrier,iu,numUsers);
end