function ehtTxSpectralFlatnessMeasurement(chanEst,cfgEHT,pktNum) 
%ehtTxSpectralFlatnessMeasurement Featured example helper function
%
%   Measures and plots the spectral flatness using channel estimates.

% Copyright 2020-2021 The MathWorks, Inc.

% Function for EHT and direct spatial mapping only
validateattributes(cfgEHT,{'ehtMUConfig'},{'scalar'},mfilename, ...
    'format configuration object');
coder.internal.errorIf(~strcmpi(cfgEHT.RU{1}.SpatialMapping,'Direct'), ...
    'wlan:vhtTxSpectralFlatnessMeasurement:InvalidSpatialMapping');

% Calculate the deviation for each subcarrier of interest
ofdmInfo = ehtOFDMInfo('EHT-Data',cfgEHT);
[dev,testSC] = spectralFlatness(cfgEHT,chanEst,ofdmInfo);

% Indices of subcarriers tested within whole FFT window
ind = sort([testSC{1} testSC{2}]);

% Spectral flatness diagram
persistent hSF
if isempty(hSF) || ~isvalid(hSF)
    % Scale and position plots on the screen
    [repositionPlots,position] = helperPlotPositions();
    hSF = figure; 
    title('Spectral Flatness, Packet: 1');
    grid('on');
    hSF.Visible = 'Off';
    if repositionPlots
        hSF.Position = position(1,:);
    end
end

% Plot spectral flatness
Nfft = ofdmInfo.FFTLength;
figure(hSF);
hSFAx = gca;
hold(hSFAx,'off');
plot(hSFAx,ind-Nfft/2-1,dev,'o');
ylim(hSFAx,[min(-6.5,min(dev(:))), max(4.5,max(dev(:)))]);
title(hSFAx,['Spectral Flatness, Packet:' num2str(pktNum)]);
ylabel(hSFAx,'Deviation (dB)');
xlabel(hSFAx,'Subcarrier Index');
grid(hSFAx,'on');
hold(hSFAx,'on');

% Overlay lower limits
indRange = (ind(1):ind(end)).'; 
dBr = [-4 -6]; % Lower limit for two sets of test subcarriers
limlow = limitPlot(dBr,indRange,testSC,Nfft);
% Overlay upper limits
dBr = [+4 +4]; % Upper limit for two sets of test subcarriers
limup = limitPlot(dBr,indRange,testSC,Nfft);

% Create legend
numAnts = size(dev,2);
legendEntries = cell(numAnts+1,1);
legendEntries(1:numAnts) = arrayfun(@(x)sprintf('Antenna %d',x), ...
    1:numAnts,'UniformOutput',false);
legendEntries{numAnts+1} = 'Deviation limit';
legend(hSFAx,legendEntries,'location','northeastoutside');

drawnow; % Update plot

% Determine if limit exceeded
if any(any(bsxfun(@lt,dev,limlow(~isnan(limlow))))) || ...
        any(any(bsxfun(@gt,dev,limup(~isnan(limup)))))
    disp('    Spectral flatness failed');
else
    disp('    Spectral flatness passed');
end

% Plot spectral flatness limit
function lim = limitPlot(dBr,indRange,testSC,Nfft)
    lim = nan(numel(indRange),1);
    [~,ia] = intersect(indRange,testSC{1});
    lim(ia) = dBr(1)*ones(size(testSC{1}));
    [~,ia] = intersect(indRange,testSC{2});
    lim(ia) = dBr(2)*ones(size(testSC{2}));
    plot(hSFAx,indRange-Nfft/2-1,lim,'r-');
end

end

function [deviation,testSC] = spectralFlatness(cfgEHT,est,ofdmInfo)
    % Spectral flatness test assumes wired link between each transmit and
    % receive antenna, therefore only use the appropriate channel
    % estimates (an identity matrix for each subcarrier).
    [Nst,Nsts,~] = size(est);
    t = repmat(permute(eye(cfgEHT.User{1}.NumSpaceTimeStreams),[3 1 2]),Nst,1,1);
    estUse = reshape(est(t==1),Nst,[]);

    % Store channel estimates in subcarrier locations with nulls
    estFullFFT = zeros(ofdmInfo.FFTLength,Nsts);
    estFullFFT(ofdmInfo.ActiveFFTIndices,:) = estUse;

    % Subcarriers used for the measurement as per IEEE P802.11be/D1.0, Table 36-62.
    switch cfgEHT.ChannelBandwidth
        case 'CBW20'
            avSC = [-84:-2 2:84]; % Average subcarrier indices
        case 'CBW40'
            avSC = [-168:-3 3:168];
        case 'CBW80'
            avSC = [-344:-3 3:344];
        case 'CBW160'
            avSC = [-696:-515 -509:-12 12:509 515:696];
        otherwise
            avSC = [-1400:-1036 -1012:-515 -509:-12 12:509 515:1012 1036:1400];
    end

    testSC{1} = avSC+ofdmInfo.FFTLength/2+1;
    testSC{2} = setdiff(ofdmInfo.ActiveFrequencyIndices,avSC).'+ofdmInfo.FFTLength/2+1;

    % Calculate average magnitude of channel estimate
    avChanEst = 20*log10(mean(abs(estFullFFT(avSC+ofdmInfo.FFTLength/2+1,:))));

    % For each set of subcarrier indices calculate the deviation from the
    % average magnitude
    for i=1:numel(testSC)
        Eiavg = 20*log10(abs(estFullFFT(testSC{i},:)));
        chanEstDeviation{i} = bsxfun(@minus,Eiavg,avChanEst); %#ok<AGROW>
    end

    % Combine the subcarrier estimate deviations into one
    [~,idx] = sort([testSC{1} testSC{2}]);
    comb = [chanEstDeviation{1}; chanEstDeviation{2}];
    deviation = comb(idx,:,:);
end