function y = ehtConstellationMap(x,NBPSCS,DCM,varargin)
%ehtConstellationMap EHT constellation mapping
%
%   Y = ehtConstellationMap(X,NBPSCS,DCM) performs constellation mapping of
%   X, and optionally dual carrier modulation using the number of coded
%   bits per subcarrier per spatial stream (NBPSCS) to one of the following
%   modulations:
%     BPSK, QPSK, 16QAM, or 64QAM as per IEEE Std 802.11-2016, Section
%     17.3.5.8; 256QAM as per IEEE Std 802.11-2016, Section 21.3.10.9.1;
%     1024QAM as per IEEE P802.11ax/D4.1, Section 27.3.11.9.
%     4096QAM as per IEEE P802.11be/D1.0, Section 36.3.13.7.
%   Constellation mapping is performed column-wise.
%
%   Y is an Nsd-by-Nsym-by-Nss-by-Nseg array containing the mapped symbols.
%   Nsd is the number of data carrying subcarriers, Nsym is the number of
%   OFDM symbols, Nss is the number of spatial streams, and Nseg is the
%   number of segments.
%
%   X is an Nsd*NBPSCS/(Nss*Nseg)-by-Nsym-by-Nss-by-Nseg array containing
%   the bits to map.
%
%   NBPSCS is a scalar specifying the number of coded bits per subcarrier
%   per spatial stream. It must be equal to 1 (BPSK), 2 (QPSK), 4 (16QAM),
%   6 (64QAM), 8 (256QAM), 10 (1024QAM), 12 (4096QAM). NBPSCS is equal to
%   log2(M), where M is the modulation order.
%
%   DCM is a logical representing if dual carrier modulation is used.
%
%   Y = ehtConstellationMap(...,PHASE) rotates the constellation points
%   counter-clockwise by the specified amount, PHASE, in radians. PHASE can
%   be a scalar, matrix or multidimensional array. PHASE and mapped symbols
%   must have compatible sizes. PHASE and mapped symbols have compatible
%   sizes if, for every dimension, the dimension sizes are either the same
%   or one of them is 1. For more information, see <a href="matlab:doc('wlanconstellationmap')">wlanConstellationMap</a>
%   documentation.
%
%   Example:
%   % Generate 4096QAM samples
%
%       % Create input bits
%       bits = randi([0 1],4800,1,'int8');
%
%       % Constellation mapping
%       NBPSCS = 12;
%       DCM = 0;
%       mappedData = ehtConstellationMap(bits,NBPSCS,DCM);
%
%   See also ehtConstellationDemap

%   Copyright 2020-2021 The MathWorks, Inc.

% Validate NBPSCS
if (~isnumeric(NBPSCS) || ~isscalar(NBPSCS) || ~any(NBPSCS==[1 2 4 6 8 10 12]))
    error('NBPSCS must be equal to 1, 2, 4, 6, 8, 10, or 12.');
end
% Validate Input
validateattributes(x,{'int8','double'},{'binary'},mfilename,'Input bits');
coder.internal.errorIf(mod(size(x,1),NBPSCS)~=0,'wlan:wlanConstellationMap:xSizeBit');

% Return an empty matrix if x is empty
if isempty(x)
    y = zeros(size(x),'double');
    return;
end

if DCM
    % DCM - IEEE P802.11ax/D4.1, Section 27.3.11.9
    switch NBPSCS
        case 1 % BPSK
            % Frequency dependent phase shift on upper half
            [NSD,NSYM,~,~] = size(x);
            lowerMappedData = constellationMap(x,NBPSCS);
            k = repmat((0:NSD-1).',1,NBPSCS*NSYM);
            upperMappedData = lowerMappedData.*exp(1i*pi*(k+NSD));
        case 2 % QPSK
            % Conjugate upper half
            lowerMappedData = constellationMap(x,NBPSCS);
            upperMappedData = conj(lowerMappedData);
        otherwise % NBPSCS = 4, 16-QAM
            assert(NBPSCS==4)
            NBPSS = size(x,1);
            % Permute bits before mapping on upper half
            lowerMappedData = constellationMap(x,NBPSCS);
            permInd = reshape(bsxfun(@plus,repmat([2 1 4 3].',1,NBPSS/4),(0:4:NBPSS-1)),NBPSS,1);
            upperMappedData = constellationMap(x(permInd,:,:,:),NBPSCS);
    end
    yt = [lowerMappedData; upperMappedData];
else
    % Constellation mapping
    yt = constellationMap(x,NBPSCS);
end

if (nargin==4) % Apply phase rotation
    phase = varargin{1};
    wlan.internal.validatePhase(phase,size(yt),mfilename);
    y = complex(bsxfun(@times,yt,exp(1i*phase))); % Force complex so it is maintained for BPSK
else
    y = yt;
end

end

function y = constellationMap(x,NBPSCS,varargin)
%constellationMap Constellation mapping
%
%   Y = constellationMap(X,NBPSCS) maps the input bits (X) using the
%   number of coded bits per subcarrier per spatial stream (NBPSCS).
%   
%   Y is a complex vector, matrix or multidimensional array containing the
%   mapped symbols. Y has the same size of X except for the number of rows,
%   which is equal to the number of rows of X divided by NBPSCS.
%
%   X is a binary 'int8' or 'double' vector, matrix or multidimensional 
%   array containing the input bits to map into symbols. The number of rows
%   of X must be an integer multiple of NBPSCS.
%
%   NBPSCS is the number of coded bits per subcarrier per spatial stream.
%
%   Y = constellationMap(...,PHASE) rotates the constellation points
%   counter-clockwise by the specified amount, PHASE, in radians. PHASE can
%   be a scalar, matrix or multidimensional array. PHASE and mapped symbols
%   must have compatible sizes. PHASE and mapped symbols have compatible
%   sizes if, for every dimension, the dimension sizes are either the same
%   or one of them is 1. For more information, see <a href="matlab:doc('wlanconstellationmap')">wlanConstellationMap</a>
%   documentation.

% Separate out BPSK from other QAM modulations
if NBPSCS==1
    % As per IEEE Std 802.11-2016 Section 17.3.5.8, Fig. 17-10 
    constellation = [-1; 1];
    yt = complex(constellation(x+1)); % Force to complex
else
    % Get symbol mapping
    symMap = comm.internal.qam.wlanSymbolMap(NBPSCS);
    if ndims(x)>3
        % qam.modulate does not support inputs with more than 3 dimensions. 
        % Reshape into a column
        yTmp = comm.internal.qam.modulate(x(:),2^NBPSCS,'custom',symMap,1,1,[]);
        inpSize = size(x);
        yt = reshape(yTmp,[inpSize(1)/NBPSCS inpSize(2:end)]);
    else
        yt = comm.internal.qam.modulate(x,2^NBPSCS,'custom',symMap,1,1,[]);
    end
end

if (nargin==3) % Apply phase rotation
    phase = varargin{1};
    wlan.internal.validatePhase(phase,size(yt),mfilename);
    y = complex(bsxfun(@times,yt,exp(1i*phase))); % Force complex so it is maintained for BPSK
else
    y = yt;
end
end