function [commonParams,varargout] = ehtCodingParameters(cfg,varargin)
%ehtCodingParameters EHT coding parameters
% 
%   [COMMONPARAMS,USERPARAMS] = ehtCodingParameters(CFG) returns a
%   structure COMMONPARAMS, containing the coding parameters common to all
%   users, and an array of structures, USERPARAMS, containing the coding
%   parameters for each user, and an array as define in IEEE
%   P802.11be/D1.0.
%
%   CFG is a format configuration object of type <a href="matlab:help('ehtMUConfig')">ehtMUConfig</a>.
%
%   [COMMONPARAMS,USERPARAMS] = ehtCodingParameters(CFG,RUIDX,USERIDX)
%   returns the user parameters given the specific RU (resource unit)
%   index, RUIDX, and user index, USERIDX

%   Copyright 2020-2021 The MathWorks, Inc.

% Form a vector of parameters, each element is the parameter for a user
numUsers = numel(cfg.User);

% Determine which user objects are active - the STAID is not 2046
userActive = true(1,numUsers); % Vector indicating if a user object is active
for i = 1:numUsers
    if cfg.User{i}.STAID==2046
        % If STAID is 2046, then RU carries no data, and user is not
        % active. Therefore do not take user into account when
        % calculating coding parameters.
        userActive(i) = false;
    end
end

% Get a vector of the user object numbers which are active (STAID is
% not 2046)
numActiveUsers = sum(userActive==true);
activeUserNumbers = zeros(1,numActiveUsers);
userIdx = 1;
for i = 1:numel(userActive)
    if userActive(i)
        % If user active then store active user number
        activeUserNumbers(userIdx) = i;
        userIdx = userIdx+1;
    end
end

ruSize = cell(numUsers,1);
mcs = zeros(numUsers,1);
numSTS = zeros(numUsers,1);
apepLength = zeros(numUsers,1);
channelCoding = repmat({'LDPC'},numUsers,1); % Initialize
dcm = false(numUsers,1); % DCM is only applicable for MCS 14, 15

for userIdx = 1:numUsers
    ruSize{userIdx} = cfg.RU{cfg.User{userIdx}.RUNumber}.Size;
    mcs(userIdx) = cfg.User{userIdx}.MCS;
    numSTS(userIdx) = cfg.User{userIdx}.NumSpaceTimeStreams;
    if cfg.User{userIdx}.STAID==2046
        % If STAID is 2046, then RU carries no data, therefore do not
        % include user in coding calculations by setting the APEPLength
        % to 0
        apepLength(userIdx) = 0;
    else
        apepLength(userIdx) = cfg.User{userIdx}.APEPLength;
    end
    channelCoding{userIdx} = cfg.User{userIdx}.ChannelCoding;
    if any(mcs(userIdx)==[14 15])
        dcm(userIdx) = true;
    end
end

nss = numSTS;
Nservice = 16; % Table 36-18

% First calculation: the initial common number of symbols and pre-FEC
% padding factor. Calculate values for all users and take the maximum. Also
% calculate NCBPSSHORT, NDPBSSHORT, and the rate dependent parameters per
% user as part of this.

NSYMinit = zeros(numUsers,1);   % Number of symbols (initial)
ainit = zeros(numUsers,1);      % Pre-FEC padding factor (initial)
NCBPSSHORT = zeros(numUsers,1); % Number of coded bits per symbol (short)
NDBPSSHORT = zeros(numUsers,1); % Number of data bits per symbol (short)
Ntail = zeros(numUsers,1);      % Number of tail bits
R = zeros(numUsers,1);          % Rate
NSS = zeros(numUsers,1);        % Number of spatial streams
NBPSCS = zeros(numUsers,1);     % Number of bits per subcarrier
NDBPS = zeros(numUsers,1);      % Number of data bits per symbol
NCBPS = zeros(numUsers,1);      % Number of coded bits per symbols
NSD = zeros(numUsers,1);        % Number of data carrying subcarriers

for u = 1:numUsers
    switch channelCoding{u}
        case 'BCC'
            Ntail(u) = 6;
        case 'LDPC'
            Ntail(u) = 0;
    end

    % Get rate dependent parameters for all users
    params = wlan.internal.heRateDependentParameters(sum(ruSize{u}),mcs(u),nss(u),dcm(u));
    R(u) = params.Rate;
    NSS(u) = params.NSS;
    NBPSCS(u) = params.NBPSCS;
    NDBPS(u) = params.NDBPS;
    NCBPS(u) = params.NCBPS;
    NSD(u) = params.NSD;

    % Number of excess bits in last OFDM symbol
    NEXCESS = mod(8*apepLength(u)+Ntail(u)+Nservice,NDBPS(u)); % Equation 36-47

    % NSD, SHORT values. Table 36-46
    NSDSHORT = wlan.internal.heNSDShort(sum(ruSize{u}),dcm(u),mcs(u));

    % Initial number of symbol segments in the last OFDM symbol(s).
    % Equation 36-48
    NCBPSSHORT(u) = NSDSHORT*NSS(u)*NBPSCS(u);
    NDBPSSHORT(u) = NCBPSSHORT(u)*R(u);
    if NEXCESS==0 % Equation 36-48
        ainit(u) = 4;
    else
        ainit(u) = min(ceil(NEXCESS/(NDBPSSHORT(u))),4);
    end

    NSYMinit(u) = ceil((8*apepLength(u)+Ntail(u)+Nservice)/(NDBPS(u))); % Equation 36-49
end

% Derive user index for active users with longest encoded packet duration,
% Equation 36-48
[~,umax] = max(NSYMinit(userActive)-1+ainit(userActive)/4);

% Use values from max for all users. Equation 36-50
NSYMinitCommon = NSYMinit(activeUserNumbers(umax));
ainitCommon = ainit(activeUserNumbers(umax));

% Now we know the common pre-FEC padding factor and number of symbols,
% update each users number of coded bits in the last symbol

NDBPSLASTinit = zeros(numUsers,1);
NCBPSLASTinit = zeros(numUsers,1);
for u = 1:numUsers
    % Update each user's initial number of coded bits in its last
    % symbol, Equation 36-52, 36-53
    if ainitCommon<4
        NDBPSLASTinit(u) = ainitCommon*NDBPSSHORT(u);
        NCBPSLASTinit(u) = ainitCommon*NCBPSSHORT(u);
    else
        NDBPSLASTinit(u) = NDBPS(u);
        NCBPSLASTinit(u) = NCBPS(u);
    end
end

%%
% For each user which uses LDPC calculate the number of pre FEC padding
% bits and if an LDPC extra symbol is required.

NPADPreFEC = zeros(numUsers,1);
ldpcExtraSymbol = false(numUsers,1);
for u = 1:numUsers
    if strcmp(channelCoding{u},'LDPC')
        % Equation 36-54
        NPADPreFEC(u) = (NSYMinitCommon-1)*NDBPS(u)+NDBPSLASTinit(u)-8*apepLength(u)-Nservice; 

        % Equation 36-55, 36-56, 35-57, 36-58
        mSTBC = 1; % No STBC in EHT
        ldpcParms = wlan.internal.heCommonLDPCParameters(NSYMinitCommon,mSTBC,NDBPS(u),NCBPS(u),NDBPSLASTinit(u),NCBPSLASTinit(u),R(u));
        ldpcExtraSymbol(u) = ldpcParms.LDPCExtraSymbol;
    end
end

%% 
% Update NSYM, the pre-FEC padding factor, NDBPSLast, and NCBPSLast for all
% users now we know if an LDPC extra symbol is required. We can also
% calculate the Pre FEC padding factor for BCC users.

commonLDPCExtraSymbol = any(ldpcExtraSymbol(userActive));
if commonLDPCExtraSymbol
    % Equation 36-59
    if ainitCommon==4
        NSYM = NSYMinitCommon+1;
        a = 1;
    else
        NSYM = NSYMinitCommon;
        a = ainitCommon+1;
    end
else
    % Equation 36-60
    NSYM = NSYMinitCommon;
    a = ainitCommon;
end

NDBPSLAST = zeros(numUsers,1);
NCBPSLAST = zeros(numUsers,1);
NPADPreFECMAC = zeros(numUsers,1);
NPADPreFECPHY = zeros(numUsers,1);
NPADPostFEC = zeros(numUsers,1);
for u = 1:numUsers
    % Equation 36-62
    if a<4
        NCBPSLAST(u) = a*NCBPSSHORT(u);
    else
        NCBPSLAST(u) = NCBPS(u);
    end

    switch channelCoding{u}
        case 'LDPC'
            % Equation 36-61
            NDBPSLAST(u) = NDBPSLASTinit(u);
        case 'BCC'
            % Equation 36-62
            if a<4
                NDBPSLAST(u) = a*NDBPSSHORT(u);
            else
                NDBPSLAST(u) = NDBPS(u);
            end

            % Equation 36-64
            NPADPreFEC(u) = (NSYM-1)*NDBPS(u)+NDBPSLAST(u)-8*apepLength(u)-Ntail(u)-Nservice;
    end

    NPADPostFEC(u) = NCBPS(u)-NCBPSLAST(u); % Equation 36-65
    
    NPADPreFECMAC(u) = floor(NPADPreFEC(u)/8)*8; % Equation 36-66
    NPADPreFECPHY(u) = mod(NPADPreFEC(u),8); % Equation 36-67
end

if all(apepLength==0)
    % For NDP set all parameters to 0 so no data symbols transmitted
    NSYM = 0;
    NSYMinitCommon = 0;
    NCBPSSHORT = zeros(numUsers,1);
    NDBPSSHORT = zeros(numUsers,1);
    NCBPSLAST = zeros(numUsers,1);
    NCBPSLASTinit = zeros(numUsers,1);
    NDBPSLAST = zeros(numUsers,1);
    NDBPSLASTinit = zeros(numUsers,1);
    NPADPreFECMAC = zeros(numUsers,1);
    NPADPreFECPHY = zeros(numUsers,1);
    NPADPostFEC = zeros(numUsers,1);
    a = 4;
    ainitCommon = 4;
    commonLDPCExtraSymbol = false;
end

% Parameters common to all users
commonParams = struct;
commonParams.NSYM = NSYM;
commonParams.NSYMInit = NSYMinitCommon;
commonParams.PreFECPaddingFactor = a;
commonParams.PreFECPaddingFactorInit = ainitCommon;
commonParams.LDPCExtraSymbol = commonLDPCExtraSymbol;

if nargout>1
    if nargin>1
        u = varargin{1}; % User of interest

        % Return structure for user of interest
        userParams = struct;
        userParams.NSYM = NSYM;
        userParams.NSYMInit = NSYMinitCommon; % Use the common MU one as per 36-49
        userParams.mSTBC = 1; % No STBC in EHT (leaving this field, needed by wlan.internal.heCommonLDPCParameters)
        userParams.Rate = R(u);
        userParams.NBPSCS = NBPSCS(u);
        userParams.NSD = NSD(u);
        userParams.NCBPS = NCBPS(u);
        userParams.NDBPS = NDBPS(u);
        userParams.DCM = dcm(u);
        userParams.NSS = NSS(u);
        userParams.ChannelCoding = channelCoding{u};

        userParams.NCBPSSHORT = NCBPSSHORT(u);
        userParams.NDBPSSHORT = NDBPSSHORT(u);
        userParams.NCBPSLAST = NCBPSLAST(u);
        userParams.NCBPSLASTInit = NCBPSLASTinit(u);
        userParams.NDBPSLAST = NDBPSLAST(u);
        userParams.NDBPSLASTInit = NDBPSLASTinit(u);
        userParams.NPADPreFECMAC = NPADPreFECMAC(u);
        userParams.NPADPreFECPHY = NPADPreFECPHY(u);
        userParams.NPADPostFEC = NPADPostFEC(u);
        userParams.PreFECPaddingFactor = a;
        userParams.PreFECPaddingFactorInit = ainitCommon;
        userParams.LDPCExtraSymbol = commonLDPCExtraSymbol;
        userParams.PSDULength = getPSDUlength(userParams);
    else
        % Initialize structure
        p = struct;
        p.NSYM = 0;
        p.NSYMInit = 0; % Use the common MU one as per 36-49
        p.mSTBC = 0; % No STBC in EHT (leaving this field, needed by wlan.internal.heCommonLDPCParameters)
        p.Rate = 0;
        p.NBPSCS = 0;
        p.NSD = 0;
        p.NCBPS = 0;
        p.NDBPS = 0;
        p.NSS = 0;
        p.DCM = false;
        p.ChannelCoding = 'LDPC';
        p.NCBPSSHORT = 0;
        p.NDBPSSHORT = 0;
        p.NCBPSLAST = 0;
        p.NCBPSLASTInit = 0;
        p.NDBPSLAST = 0;
        p.NDBPSLASTInit = 0;
        p.NPADPreFECMAC = 0;
        p.NPADPreFECPHY = 0;
        p.NPADPostFEC = 0;
        p.PreFECPaddingFactor = 0;
        p.PreFECPaddingFactorInit = 0;
        p.LDPCExtraSymbol = false;
        p.PSDULength = 0;
        
        % Replicate for all users and populate
        if coder.target('MATLAB')
            if numUsers==1
                userParams = struct;
            else
                userParams = repmat(p,numUsers,1);
            end
        else
            userParams = repmat(p,numUsers,1);
            coder.varsize('userParams(:).ChannelCoding');
        end
        
        for u = 1:numUsers
            userParams(u).NSYM = NSYM;
            userParams(u).NSYMInit = NSYMinitCommon; % Use the common MU one as per 36-51
            userParams(u).mSTBC = 1; % No STBC in EHT (leaving this field, needed by wlan.internal.heCommonLDPCParameters)

            userParams(u).Rate = R(u);
            userParams(u).NBPSCS = NBPSCS(u);
            userParams(u).NSD = NSD(u);
            userParams(u).NCBPS = NCBPS(u);
            userParams(u).NDBPS = NDBPS(u);
            userParams(u).NSS = NSS(u);
            userParams(u).DCM = dcm(u);
            userParams(u).ChannelCoding = channelCoding{u};

            userParams(u).NCBPSSHORT = NCBPSSHORT(u);
            userParams(u).NDBPSSHORT = NDBPSSHORT(u);
            userParams(u).NCBPSLAST = NCBPSLAST(u);
            userParams(u).NCBPSLASTInit = NCBPSLASTinit(u);
            userParams(u).NDBPSLAST = NDBPSLAST(u);
            userParams(u).NDBPSLASTInit = NDBPSLASTinit(u);
            userParams(u).NPADPreFECMAC = NPADPreFECMAC(u);
            userParams(u).NPADPreFECPHY = NPADPreFECPHY(u);
            userParams(u).NPADPostFEC = NPADPostFEC(u);
            userParams(u).PreFECPaddingFactor = a;
            userParams(u).PreFECPaddingFactorInit = ainitCommon;
            userParams(u).LDPCExtraSymbol = commonLDPCExtraSymbol;
            userParams(u).PSDULength = getPSDUlength(userParams(u));
        end
    end
    varargout{1} = userParams;
end

end

function psduLength = getPSDUlength(in)
    Nservice = 16;
    % IEEE P802.11be/D1.0, Section 36.4.3, Equation 36-111
    if strcmp(in.ChannelCoding,'BCC')
        Ntail = 6;
        psduLength = floor(((in.NSYM-1)*in.NDBPS+in.NDBPSLAST-Nservice-Ntail)/8);
    else % IEEE P802.11be/D1.0, Section 36.4.3, Equation 36-111
        psduLength = floor(((in.NSYMInit-1)*in.NDBPS+in.NDBPSLASTInit-Nservice)/8);
    end
end