NumTxAnts = 4; % Number of transmit antennas
NumSTS = 2;    % Number of space-time streams
NumRxAnts = 2; % Number of receive antennas
cfgHEBase = wlanHESUConfig;
cfgHEBase.ChannelBandwidth = 'CBW20';      % Channel bandwidth
cfgHEBase.NumSpaceTimeStreams = NumSTS;    % Number of space-time streams
cfgHEBase.NumTransmitAntennas = NumTxAnts; % Number of transmit antennas
cfgHEBase.APEPLength = 1e3;                % Payload length in bytes
cfgHEBase.ExtendedRange = false;           % Do not use extended range format
cfgHEBase.Upper106ToneRU = false;          % Do not use upper 106 tone RU
cfgHEBase.PreHESpatialMapping = false;     % Spatial mapping of pre-HE fields
cfgHEBase.GuardInterval = 0.8;             % Guard interval duration
cfgHEBase.HELTFType = 4;                   % HE-LTF compression mode
cfgHEBase.ChannelCoding = 'LDPC';          % Channel coding
cfgHEBase.MCS = 3;                         % Modulation and coding scheme
cfgHEBase.SpatialMapping = 'Custom';       % Custom for beamforming

cfgNDP = cfgHEBase;
cfgNDP.APEPLength = 0;                  % NDP has no data
cfgNDP.NumSpaceTimeStreams = NumTxAnts; % For feedback matrix calculation
cfgNDP.SpatialMapping = 'Direct';       % Each TxAnt carries a STS

% Create and configure the TGax channel
chanBW = cfgHEBase.ChannelBandwidth;
tgaxChannel = wlanTGaxChannel;
tgaxChannel.DelayProfile = 'Model-B';
tgaxChannel.NumTransmitAntennas = NumTxAnts;
tgaxChannel.NumReceiveAntennas = NumRxAnts;
tgaxChannel.TransmitReceiveDistance = 5; % Distance in meters for NLOS
tgaxChannel.ChannelBandwidth = chanBW;
tgaxChannel.LargeScaleFadingEffect = 'None';
fs = wlanSampleRate(cfgHEBase);
tgaxChannel.SampleRate = fs;

codeBookSize = [0 1 Inf];

snr = 10:2:18;

maxNumErrors = 10;   % The maximum number of packet errors at an SNR point
maxNumPackets = 100; % The maximum number of packets at an SNR point

numQuant = numel(codeBookSize);
numSNR = numel(snr); % Number of SNR points
packetErrorRate = zeros(numQuant,numSNR);

% Get occupied subcarrier indices and OFDM parameters
ofdmInfo = wlanHEOFDMInfo('HE-Data',cfgHEBase);

% Indices to extract fields from the PPDU
ind = wlanFieldIndices(cfgHEBase);
indSound = wlanFieldIndices(cfgNDP);

for ibf = 1:numQuant
    switch codeBookSize(ibf) % See P802.11ax/D4.1 Section 9.4.1.64
        case 0
            NumBitsPsi = 2; % Number of bits for psi
            NumBitsPhi = 4; % Number of bits for phi
            disp('End-to-End simulation with compressed beamforming quantization with');
            disp(['Number of Bits for phi = ' num2str(NumBitsPhi) ...
                ' and Number of Bits for psi = ' num2str(NumBitsPsi)]);
        case 1
            NumBitsPsi = 4; % Number of bits for psi
            NumBitsPhi = 6; % Number of bits for phi
            disp('End-to-End simulation with compressed beamforming quantization with');
            disp(['Number of Bits for phi = ' num2str(NumBitsPhi) ...
                ' and Number of Bits for psi = ' num2str(NumBitsPsi)]);
        otherwise
            disp('End-to-End simulation with non-compressed beamforming');
    end

    %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',100);
        stream.Substream = isnr;
        RandStream.setGlobalStream(stream);

        % Create an instance of the AWGN channel per SNR point simulated
        awgnChannel = comm.AWGNChannel;
        awgnChannel.NoiseMethod = 'Signal to noise ratio (SNR)';
        awgnChannel.SignalPower = 1/tgaxChannel.NumReceiveAntennas;
        % Account for noise energy in nulls so the SNR is defined per
        % active subcarrier
        awgnChannel.SNR = snr(isnr)-10*log10(ofdmInfo.FFTLength/ofdmInfo.NumTones);

        % Create an instance of the HE configuration object per SNR point
        % simulated. This will enable to use parfor
        cfgHE = cfgHEBase;

        % Loop to simulate multiple packets
        numPacketErrors = 0;
        numPkt = 1; % Index of packet transmitted
        while numPacketErrors<=maxNumErrors && numPkt<=maxNumPackets
            % Null data packet transmission
            tx = wlanWaveformGenerator([],cfgNDP);

            % Add trailing zeros to allow for channel delay
            txPad = [tx; zeros(50,cfgNDP.NumTransmitAntennas)];

            % Pass through a fading indoor TGax channel
            reset(tgaxChannel); % Reset channel for different realization
            rx = tgaxChannel(txPad);

            % Pass the waveform through AWGN channel
            rx = awgnChannel(rx);

            % Calculate the steering matrix at the beamformee
            V = heUserBeamformingFeedback(rx,cfgNDP,true);

            if isempty(V)
                % User feedback failed, packet error
                numPacketErrors = numPacketErrors+1;
                numPkt = numPkt+1;
                continue; % Go to next loop iteration
            end

            if ~isinf(codeBookSize(ibf))
                % Find quantized angles of the beamforming feedback matrix
                angidx = bfCompressQuantize(V(:,1:NumSTS,:),NumBitsPhi,NumBitsPsi);

                % Calculate steering matrix from the quantized angles at
                % beamformer:
                % Assuming zero delay in transmitting the quantized angles
                % from beamformee to beamformer, the steering matrix is
                % calculated from the quantized angles and is used in the
                % data transmission of beamformer.

                [~,Nc,Nr] = size(V(1,1:NumSTS,:));
                V = bfDecompress(angidx,Nr,Nc,NumBitsPhi,NumBitsPsi);
            end

            steeringMat = V(:,1:NumSTS,:);

            % Beamformed data transmission
            psduLength = getPSDULength(cfgHE); % PSDU length in bytes
            txPSDU = randi([0 1],psduLength*8,1); % Generate random PSDU
            cfgHE.SpatialMappingMatrix = steeringMat;
            tx = wlanWaveformGenerator(txPSDU,cfgHE);

            % Add trailing zeros to allow for channel delay
            txPad = [tx; zeros(50,cfgHE.NumTransmitAntennas)];

            % Pass through a fading indoor TGax channel
            rx = tgaxChannel(txPad);

            % Pass the waveform through AWGN channel
            rx = awgnChannel(rx);

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

            % Extract L-STF and perform coarse frequency offset correction
            lstf = rx(coarsePktOffset+(ind.LSTF(1):ind.LSTF(2)),:);
            coarseFreqOff = wlanCoarseCFOEstimate(lstf,chanBW);
            rx = helperFrequencyOffset(rx,fs,-coarseFreqOff);

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

            % Determine final packet offset
            pktOffset = coarsePktOffset+finePktOffset;

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

            % Extract L-LTF and perform fine frequency offset correction
            rxLLTF = rx(pktOffset+(ind.LLTF(1):ind.LLTF(2)),:);
            fineFreqOff = wlanFineCFOEstimate(rxLLTF,chanBW);
            rx = helperFrequencyOffset(rx,fs,-fineFreqOff);

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

            % Data demodulate
            rxData = rx(pktOffset+(ind.HEData(1):ind.HEData(2)),:);
            demodSym = wlanHEDemodulate(rxData,'HE-Data',cfgHE);

            % Pilot phase tracking
            % Average single-stream pilot estimates over symbols (2nd dimension)
            pilotEstTrack = mean(pilotEst,2);
            demodSym = heCommonPhaseErrorTracking(demodSym,pilotEstTrack,cfgHE);

            % Estimate noise power in HE fields
            nVarEst = heNoiseEstimate(demodSym(ofdmInfo.PilotIndices,:,:),pilotEstTrack,cfgHE);

            % Extract data subcarriers from demodulated symbols and channel
            % estimate
            demodDataSym = demodSym(ofdmInfo.DataIndices,:,:);
            chanEstData = chanEst(ofdmInfo.DataIndices,:,:);

            % Equalization and STBC combining
            [eqDataSym,csi] = heEqualizeCombine(demodDataSym,chanEstData,nVarEst,cfgHE);

            % Recover data
            rxPSDU = wlanHEDataBitRecover(eqDataSym,nVarEst,csi,cfgHE,'LDPCDecodingMethod','layered-bp');

            % Determine if any bits are in error, i.e. a packet error
            packetError = ~isequal(txPSDU,rxPSDU);
            numPacketErrors = numPacketErrors+packetError;
            numPkt = numPkt+1;
        end

        % Calculate packet error rate (PER) at SNR point
        packetErrorRate(ibf,isnr) = numPacketErrors/(numPkt-1);
        disp(['MCS ' num2str(cfgHE.MCS) ','...
            ' SNR ' num2str(snr(isnr)) ...
            ' completed after ' num2str(numPkt-1) ' packets,'...
            ' PER:' num2str(packetErrorRate(ibf,isnr))]);
    end
    disp(newline);
end

figure;
lineTypes = ["k-o" "b-s" "r-*"];
semilogy(snr,packetErrorRate(1,:),lineTypes(1));
hold on;
grid on;
xlabel('SNR (dB)');
ylabel('PER');
for ibf = 2:numQuant
    semilogy(snr,packetErrorRate(ibf,:),lineTypes(ibf));
end
dataStr = [string(['Compressed Beamforming, ' newline ...
                   'NumBitsPhi = 4, NumBitsPsi = 2' newline])...
           string(['Compressed Beamforming, ' newline ...
                   'NumBitsPhi = 6, NumBitsPsi = 4' newline]) ...
                   "Non-Compressed Beamforming"];
legend(dataStr);
title(sprintf('802.11ax Beamforming PER for Channel %s, %s, %s',tgaxChannel.DelayProfile,cfgHEBase.ChannelBandwidth,cfgHEBase.ChannelCoding));