function info = ehtOFDMInfo(fieldname,cfg,varargin)
%ehtOFDMInfo OFDM information for EHT
%
%   INFO = ehtOFDMInfo(FIELDNAME,CFG) returns a structure containing OFDM
%   information for the specified field, configuration and resource unit
%   (RU) of interest.
%
%   INFO is a structure with the following fields:
%     FFTLength              - The FFT length
%     CPLength               - The cyclic prefix length
%     NumTones               - The number of active subcarriers
%     NumSubchannels         - The number of 20 MHz subchannels
%     ActiveFrequencyIndices - Indices of active subcarriers relative to DC
%                              in the range [-NFFT/2, NFFT/2-1]
%     ActiveFFTIndices       - Indices of active subcarriers within the FFT
%                              in the range [1, NFFT]
%     DataIndices            - Indices of data within the active 
%                              subcarriers in the range [1, NumTones]
%     PilotIndices           - Indices of pilots within the active
%                              subcarriers in the range [1, NumTones]
%
%   FIELDNAME is the field to demodulate and must be 'EHT-LTF', 'EHT-Data',
%   'L-LTF', 'L-SIG', 'RL-SIG', 'U-SIG', or 'EHT-SIG'
%
%   CFG is a format configuration object of type <a
%   href="matlab:help('ehtMUConfig')">ehtMUConfig</a>.
%
%   INFO = ehtOFDMInfo(...,RUNUMBER) returns OFDM info for the
%   resource unit (RU) of interest for a multi-user configuration. When not
%   provided the RUNUMBER is assumed to be one.
%
%   INFO = ehtOFDMInfo(...,'OversamplingFactor',OSF) returns the OFDM
%   information to modulate or demodulate a waveform oversampled by the
%   factor OSF. OSF must be greater than or equal to 1. If not specified, 1
%   is assumed.
%
%   Example:
%   % Get OFDM info for a 320 MHz channel bandwidth.
%
%      % Create an EHT object for 320 MHz
%      cfg = ehtMUConfig('CBW320');
%
%      info = ehtOFDMInfo('EHT-Data',cfg,1);
%      disp(info.NumTones)
%
%   See also ehtDemodulate.

%   Copyright 2020-2021 The MathWorks, Inc.

narginchk(2,5);
fieldname = validatestring(fieldname,{'EHT-Data','EHT-LTF','EHT-SIG','U-SIG','RL-SIG','L-SIG','L-LTF'},mfilename,'field');
validateattributes(cfg,{'ehtMUConfig'},{'scalar'},mfilename,'format configuration object');
[cbw,guardInterval,rusize,ruindex,osf] = parseInputs(fieldname,cfg,varargin{:});

numSubchannels = cbw/20;

switch fieldname
    case {'EHT-Data','EHT-LTF'}
        [ruMappingInd,freqInd] = ehtOccupiedSubcarrierIndices(cbw,rusize,ruindex);
        dataIndices = ruMappingInd.Data;
        pilotIndices = ruMappingInd.Pilot;

        fftLength = 256*numSubchannels;
        switch guardInterval
            case 0.8
                cpLength = 16*numSubchannels;
            case 1.6
                cpLength = 32*numSubchannels;
            otherwise % 3.2
                assert(guardInterval==3.2)
                cpLength = 64*numSubchannels;
        end
    case {'U-SIG','EHT-SIG','RL-SIG','L-SIG'}
        [freqInd,pilotIdx] = preEHTToneIndices(numSubchannels);
        [dataIndices,pilotIndices] = occupiedIndices(freqInd,pilotIdx);
        fftLength = 64*numSubchannels;
        cpLength = 16*numSubchannels;
    otherwise % 'L-LTF'
        assert(strcmp(fieldname,'L-LTF'))
        if numSubchannels==16 % 320 MHz
            pilotIdx20 = [-21; -7; 7; 21];
            freqInd20 = [-26:-1 1:26].';
            freqIndTmp = bsxfun(@plus,freqInd20,(0:64:960)-(256+128+64+32));
            pilotIdxTmp = bsxfun(@plus,pilotIdx20,(0:64:448)-(256+128+64+32));
            % Indices returned as column vectors
            freqInd = freqIndTmp(:);
            pilotIdx = pilotIdxTmp(:);
        else
            [freqInd,pilotIdx] = wlan.internal.nonHTToneIndices(numSubchannels);
        end
        [dataIndices,pilotIndices] = occupiedIndices(freqInd,pilotIdx);
        fftLength = 64*numSubchannels;
        cpLength = [32*numSubchannels 0];
end

% Form structure
info = struct;
info.FFTLength = fftLength*osf;
info.CPLength = cpLength*osf;
info.NumSubchannels = numSubchannels;
info.NumTones = numel(freqInd);
info.ActiveFrequencyIndices = freqInd;
info.ActiveFFTIndices = freqInd+fftLength*osf/2+1;
info.DataIndices = dataIndices;
info.PilotIndices = pilotIndices;

if any(floor(info.CPLength)~=info.CPLength)
    error('The specified oversampling factor %f results in a non-integer cyclic prefix length',osf)
end

end

function [cbw,guardInterval,rusize,ruindex,osf] = parseInputs(fieldname,cfg,varargin)
% Parse inputs to return parameters

    [chanBW,guardInterval,rusize,ruindex,osf] = parseMUSignatures(fieldname,cfg,varargin{:});
    cbw = wlan.internal.cbwStr2Num(chanBW);
end

function [chanBW,guardInterval,rusize,ruindex,osf] = parseMUSignatures(fieldname,cfg,varargin)
% Parse signature with ehtMUConfig object

    chanBW = cfg.ChannelBandwidth;
    guardInterval = cfg.GuardInterval;
    allocInfo = ruInfo(cfg);
    ruNumber = 1; % Default
    osf = 1; % Default
    if nargin>2
        if any(strcmp(fieldname,{'EHT-LTF','EHT-Data'}))
            if nargin==3 % ehtOFDMInfo(FIELDNAME,CFG)
                ruNumber = varargin{1};
                validateRUNumber(ruNumber,allocInfo.NumRUs);
            elseif nargin==4 % ehtOFDMInfo(FIELDNAME,CFG,'OversamplingFactor',...)
                osf = wlan.internal.parseOSF(varargin{:});
            else % ehtOFDMInfo(FIELDNAME,CFG,ruNumber,'OversamplingFactor',...)
                ruNumber = varargin{1};
                validateRUNumber(ruNumber,allocInfo.NumRUs);
                osf = wlan.internal.parseOSF(varargin{2:end});
            end
        else
            % Allow for optional RU number before NV pairs (even though it is not used)
            osf = parseOSFWithOptionalNumeric(varargin{:});
        end
    end
    rusize = allocInfo.RUSizes{ruNumber};
    ruindex = allocInfo.RUIndices{ruNumber};
end

function osf = parseOSFWithOptionalNumeric(varargin)
    if nargin>0 && isnumeric(varargin{1})
        % Ignore optional numeric
        osf = wlan.internal.parseOSF(varargin{2:end});
    else
        osf = wlan.internal.parseOSF(varargin{1:end});
    end
end

function validateRUNumber(ruNumber,numRUs)
    validateattributes(ruNumber,{'numeric'},{'scalar','integer','>',0,'<=',numRUs},mfilename,'RU Number');
end

function [freqInd,pilotIdx] = preEHTToneIndices(numSubchannels)
%preEHTToneIndices Pre-EHT tone indices
%   [FREQIND,PILOTIND] = preEHTToneIndices(NUMSUBCHANNELS) returns the
%   subcarrier frequency indices for active subcarrier and pilots for
%   pre-EHT fields, given the number of 20 MHz subchannels.

    pilotIdx20 = [-21; -7; 7; 21];
    freqInd20 = [-28:-1 1:28].';
    % Replicate indices for each 20 MHz subchannel
    switch numSubchannels
        case 1
            pilotIdxTmp = pilotIdx20;
            freqIndTmp = freqInd20;
        case 2
            freqIndTmp = bsxfun(@plus,freqInd20,[0 64]-32);
            pilotIdxTmp = bsxfun(@plus,pilotIdx20,[0 64]-32);
        case 4
            freqIndTmp = bsxfun(@plus,freqInd20,(0:64:192)-(64+32));
            pilotIdxTmp = bsxfun(@plus,pilotIdx20,(0:64:192)-(64+32));
        case 8
            freqIndTmp = bsxfun(@plus,freqInd20,(0:64:448)-(128+64+32));
            pilotIdxTmp = bsxfun(@plus,pilotIdx20,(0:64:448)-(128+64+32));
        otherwise % 16
            assert(numSubchannels==16)
            freqIndTmp = bsxfun(@plus,freqInd20,(0:64:960)-(256+128+64+32));
            pilotIdxTmp = bsxfun(@plus,pilotIdx20,(0:64:960)-(256+128+64+32));  
    end
    % Return as column vectors
    freqInd = freqIndTmp(:);
    pilotIdx = pilotIdxTmp(:);
end

function [dataInd,pilotInd] = occupiedIndices(freqInd,pilotIdx)
%occupiedIndices Pre-EHT tone indices
%   [DATAIND,PILOTIND] = occupiedIndices(FREQIND,PILOTIND) returns the
%   pilot and data indices within active subcarrier indices given the 
%   active frequency indices and pilot indices.

    idx = ismember(freqInd,pilotIdx);
    seqInd = (1:numel(freqInd))';
    pilotInd = seqInd(idx);
    dataInd = seqInd(~idx);
end