s = rng(10);                             % Set RNG seed for repeatability

% Transmission parameters
chanBW      = 'CBW80';                   % Channel bandwidth
numUsers    = 3;                         % Number of users
numSTSVec   = [3 1 4];                   % Number of streams per user
userPos     = [0 1 2];                   % User positions
mcsVec      = [4 6 8];                   % MCS per user: 16QAM, 64QAM, 256QAM
apepVec     = [520 192 856];             % Payload per user, in bytes
chCodingVec = {'BCC', 'LDPC', 'LDPC'};   % Channel coding per user

% Precoding and equalization parameters
precodingType = 'ZF';                    % Precoding type; ZF or MMSE
snr           = 47;                      % SNR in dB
eqMethod      = 'ZF';                    % Equalization method

% Create the multi-user VHT format configuration object
numTx = sum(numSTSVec);
cfgVHTMU = wlanVHTConfig('ChannelBandwidth', chanBW,...
    'NumUsers', numUsers, ...
    'NumTransmitAntennas', numTx, ...
    'GroupID', 2, ...
    'NumSpaceTimeStreams', numSTSVec,...
    'UserPositions', userPos, ...
    'MCS', mcsVec, ...
    'APEPLength', apepVec, ...
    'ChannelCoding', chCodingVec);

% VHT sounding (NDP) configuration, for same number of streams
cfgVHTNDP = wlanVHTConfig('ChannelBandwidth', chanBW,...
    'NumUsers', 1, ...
    'NumTransmitAntennas', numTx, ...
    'GroupID', 0, ...
    'NumSpaceTimeStreams', sum(numSTSVec),...
    'MCS', 0, ...
    'APEPLength', 0);

% Generate the null data packet, with no data
txNDPSig = wlanWaveformGenerator([], cfgVHTNDP);
NPDSigLen = size(txNDPSig, 1);

% Set up layout parameters for WINNER II channel
AA = winner2.AntennaArray('UCA', numTx, 0.2);
for i = 1:numUsers
    AA(i+1) = winner2.AntennaArray('ULA', numSTSVec(i), 0.05);
end
STAIdx   = 2:(numUsers+1);
APIdx   = {1};
rndSeed = 12;
cfgLayout = winner2.layoutparset(STAIdx, APIdx, numUsers, AA, [], rndSeed);
cfgLayout.Pairing = [ones(1,numUsers);2:(numUsers+1)]; % One access point to all users
cfgLayout.ScenarioVector = ones(1,numUsers);           % A1 scenario for all links
cfgLayout.PropagConditionVector = zeros(1, numUsers);  % NLOS
for i = 1:numUsers % Randomly set velocity for each user
    v = rand(3,1) - 0.5;
    cfgLayout.Stations(i+1).Velocity = v/norm(v, 'fro');
end

% Set up model parameters for WINNER II channel
cfgModel = winner2.wimparset;
cfgModel.FixedPdpUsed       = 'yes';
cfgModel.FixedAnglesUsed    = 'yes';
cfgModel.IntraClusterDsUsed = 'no';
cfgModel.RandomSeed         = 111;    % Repeatability

% The maximum velocity for the 3 users is 1m/s. Set up the SampleDensity
% field to ensure that the sample rate matches the channel bandwidth.
maxMSVelocity = max(cell2mat(cellfun(@(x) norm(x, 'fro'), ...
    {cfgLayout.Stations.Velocity}, 'UniformOutput', false)));
cfgModel.UniformTimeSampling = 'yes';
cfgModel.SampleDensity = round(physconst('LightSpeed')/ ...
    cfgModel.CenterFrequency/2/(maxMSVelocity/wlanSampleRate(cfgVHTMU)));

% Create the WINNER II channel System object
WINNERChan = comm.WINNER2Channel(cfgModel, cfgLayout);

% Call the info method to check some derived channel parameters
chanInfo = info(WINNERChan) %#ok<NOPTS>

chanDelay   = chanInfo.ChannelFilterDelay;
numPadZeros = 10;

% Set ModelConfig.NumTimeSamples to match the length of the input signal to
% avoid warning
WINNERChan.ModelConfig.NumTimeSamples = NPDSigLen + numPadZeros;

% Sound the WINNER II channel for all users
chanOutNDP = WINNERChan([txNDPSig;zeros(numPadZeros,numTx)]);

% Add AWGN
rxNDPSig = cellfun(@awgn, chanOutNDP, ...
    num2cell(snr*ones(numUsers, 1)), 'UniformOutput', false);

mat = cell(numUsers,1);
for uIdx = 1:numUsers
    % Compute the feedback matrix based on received signal per user
    mat{uIdx} = vhtCSIFeedback(rxNDPSig{uIdx}(chanDelay(uIdx)+1:end,:), ...
        cfgVHTNDP, uIdx, numSTSVec);
end

% Pack the per user CSI into a matrix
numST = length(mat{1});         % Number of subcarriers
steeringMatrix = zeros(numST, sum(numSTSVec), sum(numSTSVec));
%   Nst-by-Nt-by-Nsts
for uIdx = 1:numUsers
    stsIdx = sum(numSTSVec(1:uIdx-1))+(1:numSTSVec(uIdx));
    steeringMatrix(:,:,stsIdx) = mat{uIdx};     % Nst-by-Nt-by-Nsts
end

% Zero-forcing or MMSE precoding solution
if strcmp(precodingType, 'ZF')
    delta = 0; % Zero-forcing
else
    delta = (numTx/(10^(snr/10))) * eye(numTx); % MMSE
end
for i = 1:numST
    % Channel inversion precoding
    h = squeeze(steeringMatrix(i,:,:));
    steeringMatrix(i,:,:) = h/(h'*h + delta);
end

% Set the spatial mapping based on the steering matrix
cfgVHTMU.SpatialMapping = 'Custom';
cfgVHTMU.SpatialMappingMatrix = permute(steeringMatrix,[1 3 2]);

% Create data sequences, one for each user
txDataBits = cell(numUsers, 1);
psduDataBits = cell(numUsers, 1);
for uIdx = 1:numUsers
    % Generate payload for each user
    txDataBits{uIdx} = randi([0 1], cfgVHTMU.APEPLength(uIdx)*8, 1, 'int8');

    % Pad payload with zeros to form a PSDU
    psduDataBits{uIdx} = [txDataBits{uIdx}; ...
        zeros((cfgVHTMU.PSDULength(uIdx)-cfgVHTMU.APEPLength(uIdx))*8, 1, 'int8')];
end

txSig = wlanWaveformGenerator(psduDataBits, cfgVHTMU);

release(WINNERChan);

% Set ModelConfig.NumTimeSamples to match the total length of NDP plus
% waveform and padded zeros
WINNERChan.ModelConfig.NumTimeSamples = ...
    WINNERChan.ModelConfig.NumTimeSamples + length(txSig) + numPadZeros;

% Transmit through the WINNER II channel for all users, with 10 all-zero
% samples appended to account for channel filter delay
chanOut = WINNERChan([txNDPSig; zeros(numPadZeros, numTx); ...
	txSig; zeros(numPadZeros, numTx)]);

% Extract the waveform output for each user
chanOut = cellfun(@(x) x(NPDSigLen+numPadZeros+1:end,:), chanOut, 'UniformOutput', false);

% Add AWGN
rxSig = cellfun(@awgn, chanOut, ...
    num2cell(snr*ones(numUsers, 1)), 'UniformOutput', false);

% Get field indices from configuration, assumed known at receiver
ind = wlanFieldIndices(cfgVHTMU);

% Single-user receivers recover payload bits
rxDataBits = cell(numUsers, 1);
scaler = zeros(numUsers, 1);
spAxes = gobjects(sum(numSTSVec), 1);
hfig = figure('Name','Per-stream equalized symbol constellation');
for uIdx = 1:numUsers
    rxNSig = rxSig{uIdx}(chanDelay(uIdx)+1:end, :);

    % User space-time streams
    stsU = numSTSVec(uIdx);

    % Estimate noise power in VHT fields
    lltf = rxNSig(ind.LLTF(1):ind.LLTF(2),:);
    demodLLTF = wlanLLTFDemodulate(lltf, chanBW);
    nVar = helperNoiseEstimate(demodLLTF, chanBW, sum(numSTSVec));

    % Perform channel estimation based on VHT-LTF
    rxVHTLTF  = rxNSig(ind.VHTLTF(1):ind.VHTLTF(2),:);
    demodVHTLTF = wlanVHTLTFDemodulate(rxVHTLTF, chanBW, numSTSVec);
    chanEst = wlanVHTLTFChannelEstimate(demodVHTLTF, chanBW, numSTSVec);

    % Recover information bits in VHT Data field
    rxVHTData = rxNSig(ind.VHTData(1):ind.VHTData(2),:);
    [rxDataBits{uIdx}, ~, eqsym] = wlanVHTDataRecover(rxVHTData, ...
        chanEst, nVar, cfgVHTMU, uIdx, ...
        'EqualizationMethod', eqMethod, 'PilotPhaseTracking', 'None', ...
        'LDPCDecodingMethod', 'layered-bp', 'MaximumLDPCIterationCount', 6);

    % Plot equalized symbols for all streams per user
    scaler(uIdx) = ceil(max(abs([real(eqsym(:)); imag(eqsym(:))])));
    for i = 1:stsU
        subplot(numUsers, max(numSTSVec), (uIdx-1)*max(numSTSVec)+i);
        plot(reshape(eqsym(:,:,i), [], 1), '.');
        axis square
        spAxes(sum([0 numSTSVec(1:(uIdx-1))])+i) = gca; % Store axes handle
        title(['User ' num2str(uIdx) ', Stream ' num2str(i)]);
        grid on;
    end
end

% Scale axes for all subplots and scale figure
for i = 1:numel(spAxes)
    xlim(spAxes(i),[-max(scaler) max(scaler)]);
    ylim(spAxes(i),[-max(scaler) max(scaler)]);
end
pos = get(hfig, 'Position');
set(hfig, 'Position', [pos(1)*0.7 pos(2)*0.7 1.3*pos(3) 1.3*pos(4)]);

% Compare recovered bits against per-user APEPLength information bits
ber = inf(1, numUsers);
for uIdx = 1:numUsers
    idx = (1:cfgVHTMU.APEPLength(uIdx)*8).';
    [~, ber(uIdx)] = biterr(txDataBits{uIdx}(idx), rxDataBits{uIdx}(idx));
    disp(['Bit Error Rate for User ' num2str(uIdx) ': ' num2str(ber(uIdx))]);
end

rng(s); % Restore RNG state