equalizationMethod = 'mmse-sic';

allocationIndex = 195; % Four uplink users in a MU-MIMO configuration
cfgSys = heTBSystemConfig(allocationIndex);

% These parameters are same for all users in the MU-MIMO system
cfgSys.HELTFType = 4;          % HE-LTF compression mode
cfgSys.GuardInterval = 3.2;    % Guard interval type
cfgSys.SingleStreamPilots = 1; % Single stream pilot transmission of HE-LTF
numRx = 8;                     % Number of receive(AP) antennas

% The individual parameters for each user are specified below
allocInfo = ruInfo(cfgSys);
numUsers = allocInfo.NumUsers; % Number of uplink users

for userIdx = 1:numUsers
    cfgSys.User{userIdx}.NumTransmitAntennas = 1;
    cfgSys.User{userIdx}.NumSpaceTimeStreams = 1;
    cfgSys.User{userIdx}.SpatialMapping = 'Direct';
    cfgSys.User{userIdx}.MCS = 7;
    cfgSys.User{userIdx}.APEPLength = 1e3;
    cfgSys.User{userIdx}.ChannelCoding = 'LDPC';
end

cfgTB = getUserConfig(cfgSys);

snr = 20:2:24;

% The sample rate and field indices for the HE TB packet is same for all
% users. Here the trigger configuration of the first user is used to get
% the sample rate and field indices of the HE TB PPDU.
fs = wlanSampleRate(cfgTB{1});    % Same for all users
ind = wlanFieldIndices(cfgTB{1}); % Same for all users

tgaxBase = wlanTGaxChannel;
tgaxBase.SampleRate = fs;
tgaxBase.TransmissionDirection = 'Uplink';
tgaxBase.TransmitReceiveDistance = 10;
chanBW = cfgSys.ChannelBandwidth;
tgaxBase.ChannelBandwidth = chanBW;
tgaxBase.NumReceiveAntennas = numRx;

% A cell array stores the channel objects, one per user
tgax = cell(1,numUsers);
for userIdx = 1:numUsers
    tgax{userIdx} = clone(tgaxBase);
    tgax{userIdx}.NumTransmitAntennas = cfgSys.User{userIdx}.NumTransmitAntennas;
    tgax{userIdx}.UserIndex = userIdx;
end

ofdmInfo = wlanHEOFDMInfo('HE-Data',cfgSys.ChannelBandwidth,cfgSys.GuardInterval);
numSNR = numel(snr); % Number of SNR points
numPackets = 50; % Number of packets to simulate
packetErrorRate = zeros(numUsers,numSNR);
txPSDU = cell(numUsers);

% parfor isnr = 1:numSNR % Use 'parfor' to speed up the simulation
for isnr = 1:numSNR
    % Set random substream index per iteration to ensure that each
    % iteration uses a repeatable set of random numbers
    stream = RandStream('combRecursive','Seed',0);
    stream.Substream = isnr;
    RandStream.setGlobalStream(stream);

    % Create an instance of the AWGN channel per SNR point simulated
    awgn = comm.AWGNChannel;
    awgn.NoiseMethod = 'Signal to noise ratio (SNR)';
    awgn.SignalPower = 1/numRx;
    sysInfo = ruInfo(cfgSys);

    % Simulate multiple packets
    numPacketErrors = zeros(numUsers,1);
    for pktIdx = 1:numPackets

        % Transmit processing
        rxWaveform = 0;
        packetError = zeros(numUsers,1);
        txPSDU = cell(1,numUsers);

        % Generate random channel realization for each packet by varying
        % the UserIndex property of the channel. This assumes all users
        % have the same number of transmit antennas.
        chPermutations = randperm(numUsers);
        for userIdx = 1:numUsers
            % HE TB config object for each user
            cfgUser = cfgTB{userIdx};

            % Generate a packet with random PSDU
            txPSDU{userIdx} = randi([0 1],getPSDULength(cfgUser)*8,1,'int8');

            % Generate HE TB waveform, containing payload for single user
            txTrig = wlanWaveformGenerator(txPSDU{userIdx},cfgUser);

            % Pass waveform through a random TGax Channel
            channelIdx = chPermutations(userIdx);
            reset(tgax{channelIdx}); % New channel realization
            rxTrig = tgax{channelIdx}([txTrig; zeros(15,size(txTrig,2))]);

            % Scale the transmit power of the user within an RU. This is to
            % ensure same SNR for each user after the addition of noise.
            ruNum = cfgSys.User{userIdx}.RUNumber;
            SF = sqrt(1/sysInfo.NumUsersPerRU(ruNum))*sqrt(cfgUser.RUSize/(sum(sysInfo.RUSizes)));

            % Combine uplink users into one waveform
            rxWaveform = rxWaveform+SF*rxTrig;
        end

        % Pass the waveform through AWGN channel. Account for noise
        % energy in nulls so the SNR is defined per active subcarriers.
        awgn.SNR = snr(isnr)-10*log10(ofdmInfo.FFTLength/sum(sysInfo.RUSizes));
        rxWaveform = awgn(rxWaveform);

        % Receive processing
        % Packet detect and determine coarse packet offset
        coarsePktOffset = wlanPacketDetect(rxWaveform,chanBW,0,0.05);
        if isempty(coarsePktOffset) % If empty no L-STF detected; packet error
            numPacketErrors = numPacketErrors+1;
            continue; % Go to next loop iteration
        end

        % Extract the non-HT fields and determine fine packet offset
        nonhtfields = rxWaveform(coarsePktOffset+(ind.LSTF(1):ind.LSIG(2)),:);
        finePktOffset = wlanSymbolTimingEstimate(nonhtfields,chanBW);

        % Determine final packet offset
        pktOffset = coarsePktOffset+finePktOffset;

        % If packet detected out with the range of expected delays from
        % the channel modeling; packet error
        if pktOffset>50
            numPacketErrors = numPacketErrors+1;
            continue; % Go to next loop iteration
        end

        % Extract HE-LTF and HE-Data fields for all RUs
        rxLTF = rxWaveform(pktOffset+(ind.HELTF(1):ind.HELTF(2)),:);
        rxData = rxWaveform(pktOffset+(ind.HEData(1):ind.HEData(2)),:);

        for ruIdx = 1:allocInfo.NumRUs
            % Demodulate HE-LTF and HE-Data field for the RU of interest
            ru = [allocInfo.RUSizes(ruIdx) allocInfo.RUIndices(ruIdx)];
            demodHELTFRU = wlanHEDemodulate(rxLTF,'HE-LTF',chanBW,cfgSys.GuardInterval,cfgSys.HELTFType,ru);
            demodHEDataRU = wlanHEDemodulate(rxData,'HE-Data',chanBW,cfgSys.GuardInterval,ru);

            % Channel estimate
            [chanEst,ssPilotEst] = heLTFChannelEstimate(demodHELTFRU,cfgSys,ruIdx);

            % Get indices of data and pilots within RU (without nulls)
            ruOFDMInfo = wlanHEOFDMInfo('HE-Data',cfgSys.ChannelBandwidth,cfgSys.GuardInterval, ...
                [allocInfo.RUSizes(ruIdx) allocInfo.RUIndices(ruIdx)]);

            % Estimate noise power in HE fields of each user
            nVarEst = heNoiseEstimate(demodHEDataRU(ruOFDMInfo.PilotIndices,:,:),ssPilotEst,cfgSys,ruIdx);

            % Discard pilot subcarriers
            demodDataSym = demodHEDataRU(ruOFDMInfo.DataIndices,:,:);
            chanEstData = chanEst(ruOFDMInfo.DataIndices,:,:);

            % Equalize
            if strcmpi(equalizationMethod,'mmse-sic')
                [eqSym,csi] = heSuccessiveEqualize(demodDataSym,chanEstData,nVarEst,cfgSys,ruIdx);
            else
                [eqSym,csi] = heEqualizeCombine(demodDataSym,chanEstData,nVarEst,cfgSys);
            end

            for userIdx = 1:allocInfo.NumUsersPerRU(ruIdx)
                % Get TB config object for each user
                userNum = cfgSys.RU{ruIdx}.UserNumbers(userIdx);
                cfgUser = cfgTB{userNum};

                % Get space-time stream indices for the current user
                stsIdx = cfgUser.StartingSpaceTimeStream-1+(1:cfgUser.NumSpaceTimeStreams);

                % Demap and decode bits
                rxPSDU = wlanHEDataBitRecover(eqSym(:,:,stsIdx),nVarEst,csi(:,stsIdx),cfgUser,'LDPCDecodingMethod','layered-bp');

                % PER calculation
                packetError(userNum) = any(biterr(txPSDU{userNum},rxPSDU));
            end
        end
        numPacketErrors = numPacketErrors+packetError;
    end

    % Calculate packet error rate (PER) at SNR point
    packetErrorRate(:,isnr)= numPacketErrors/numPackets;
    disp(['SNR ' num2str(snr(isnr)) ...
          ' completed for ' num2str(numUsers) ' users']);

end

markers = 'ox*sd^v><ph+ox*sd^v';
color = 'bmcrgbrkymcrgbrkymc';
figure;

for nSTA = 1:numUsers
    semilogy(snr,packetErrorRate(nSTA,:).',['-' markers(nSTA) color(nSTA)]);
    hold on;
end

grid on;
xlabel('SNR (dB)');
ylabel('PER');
dataStr = arrayfun(@(x)sprintf('STA- %d',x),1:numUsers,'UniformOutput',false);
legend(dataStr);
title('PER for uplink 802.11ax link');