/*
 * Copyright (c) 2020 Universita' degli Studi di Napoli Federico II
 *
 * SPDX-License-Identifier: GPL-2.0-only
 *
 * Author: Stefano Avallone <stavallo@unina.it>
 */

#include "ns3/ap-wifi-mac.h"
#include "ns3/attribute-container.h"
#include "ns3/boolean.h"
#include "ns3/config.h"
#include "ns3/ctrl-headers.h"
#include "ns3/eht-phy.h"
#include "ns3/enum.h"
#include "ns3/he-configuration.h"
#include "ns3/mobility-helper.h"
#include "ns3/multi-model-spectrum-channel.h"
#include "ns3/rng-seed-manager.h"
#include "ns3/spectrum-wifi-helper.h"
#include "ns3/sta-wifi-mac.h"
#include "ns3/test.h"
#include "ns3/tuple.h"
#include "ns3/wifi-net-device.h"
#include "ns3/wifi-psdu.h"

#include <algorithm>
#include <bitset>
#include <sstream>

using namespace ns3;

NS_LOG_COMPONENT_DEFINE("WifiPrimaryChannelsTest");

/**
 * @ingroup wifi-test
 * @ingroup tests
 *
 * @brief Test transmissions under different primary channel settings
 *
 * This test can be repeated for different widths of the operating channel. We
 * configure as many BSSes as the number of distinct 20 MHz subchannels in the
 * operating channel, so that each BSS is assigned a distinct primary20 channel.
 * For each BSS, we test the transmission of SU PPDUs, DL MU PPDUs and HE TB PPDUs
 * of all the widths (20 MHz, 40 MHz, etc.) allowed by the operating channel.
 * Transmissions of a given type take place simultaneously in BSSes that do not
 * operate on adjacent primary channels of the considered width (so that
 * transmissions do not interfere with each other). It is also possible to
 * select whether BSSes should be assigned (distinct) BSS colors or not.
 */
class WifiPrimaryChannelsTest : public TestCase
{
  public:
    /**
     * Constructor
     *
     * @param standard the standard
     * @param channelWidth the operating channel width
     * @param useDistinctBssColors whether to set distinct BSS colors to BSSes
     */
    WifiPrimaryChannelsTest(WifiStandard standard, MHz_u channelWidth, bool useDistinctBssColors);

    /**
     * Callback invoked when PHY receives a PSDU to transmit. Used to print
     * transmitted PSDUs for debug purposes.
     *
     * @param context the context
     * @param psduMap the PSDU map
     * @param txVector the TX vector
     * @param txPowerW the tx power in Watts
     */
    void Transmit(std::string context,
                  WifiConstPsduMap psduMap,
                  WifiTxVector txVector,
                  double txPowerW);
    /**
     * Have the AP of the given BSS transmit a SU PPDU using the given
     * transmission channel width
     *
     * @param bss the given BSS
     * @param txChannelWidth the given transmission channel width
     */
    void SendDlSuPpdu(uint8_t bss, MHz_u txChannelWidth);
    /**
     * Have the AP of the given BSS transmit a MU PPDU using the given
     * transmission channel width and RU type
     *
     * @param bss the given BSS
     * @param txChannelWidth the given transmission channel width
     * @param ruType the given RU type
     * @param nRus the number of RUs
     */
    void SendDlMuPpdu(uint8_t bss, MHz_u txChannelWidth, RuType ruType, std::size_t nRus);
    /**
     * Have the AP of the given BSS transmit a Basic Trigger Frame. This method calls
     * DoSendHeTbPpdu to actually have STAs transmit HE TB PPDUs using the given
     * transmission channel width and RU type
     *
     * @param bss the given BSS
     * @param txChannelWidth the given transmission channel width
     * @param ruType the given RU type
     * @param nRus the number of RUs
     */
    void SendHeTbPpdu(uint8_t bss, MHz_u txChannelWidth, RuType ruType, std::size_t nRus);
    /**
     * Have the STAs of the given BSS transmit an HE TB PPDU using the given
     * transmission channel width and RU type
     *
     * @param bss the given BSS
     * @param txChannelWidth the given transmission channel width
     * @param ruType the given RU type
     * @param nRus the number of RUs
     */
    void DoSendHeTbPpdu(uint8_t bss, MHz_u txChannelWidth, RuType ruType, std::size_t nRus);
    /**
     * Callback invoked when a station receives a DL PPDU.
     *
     * @param bss the BSS the receiving STA belongs to
     * @param station the receiving station
     * @param psdu the received PSDU
     * @param rxSignalInfo the info on the received signal (\see RxSignalInfo)
     * @param txVector TxVector of the received PSDU
     * @param perMpduStatus per MPDU reception status
     */
    void ReceiveDl(uint8_t bss,
                   uint8_t station,
                   Ptr<const WifiPsdu> psdu,
                   RxSignalInfo rxSignalInfo,
                   const WifiTxVector& txVector,
                   const std::vector<bool>& perMpduStatus);
    /**
     * Callback invoked when an AP receives an UL PPDU.
     *
     * @param bss the BSS the receiving AP belongs to
     * @param psdu the received PSDU
     * @param rxSignalInfo the info on the received signal (\see RxSignalInfo)
     * @param txVector TxVector of the received PSDU
     * @param perMpduStatus per MPDU reception status
     */
    void ReceiveUl(uint8_t bss,
                   Ptr<const WifiPsdu> psdu,
                   RxSignalInfo rxSignalInfo,
                   const WifiTxVector& txVector,
                   const std::vector<bool>& perMpduStatus);
    /**
     * Check that all stations associated with an AP.
     */
    void CheckAssociation();
    /**
     * Check that (i) all stations belonging to the given BSSes received the SU PPDUs
     * transmitted over the given channel width; and (ii) all stations belonging to
     * the other BSSes did not receive any frame if BSS Color is set (due to BSS Color
     * filtering) or if no transmission was performed on a channel adjacent to the one
     * they operate on, otherwise.
     *
     * @param txBss the set of BSSes that transmitted an SU PPDU
     * @param txChannelWidth the given transmission channel width
     */
    void CheckReceivedSuPpdus(std::set<uint8_t> txBss, MHz_u txChannelWidth);
    /**
     * Check that (i) all stations/APs belonging to the given BSSes received the DL/UL MU PPDUs
     * transmitted over the given channel width and RU width; and (ii) stations/APs belonging to
     * the other BSSes did not receive any frame if BSS Color is set (due to BSS Color
     * filtering) or if no transmission addressed to/from stations with the same AID was
     * performed on a channel adjacent to the one they operate on, otherwise.
     *
     * @param txBss the set of BSSes that transmitted an SU PPDU
     * @param txChannelWidth the given transmission channel width
     * @param ruType the given RU type
     * @param nRus the number of RUs
     * @param isDlMu true for DL MU PPDU, false for HE TB PPDU
     */
    void CheckReceivedMuPpdus(std::set<uint8_t> txBss,
                              MHz_u txChannelWidth,
                              RuType ruType,
                              std::size_t nRus,
                              bool isDlMu);
    /**
     * Check that (i) all stations belonging to the given BSSes received the transmitted
     * Trigger Frame; and (ii) all stations belonging to the other BSSes did not receive
     * any frame Trigger Frame (given that a Trigger Frame is transmitted on the primary20
     * channel and all the primary20 channels are distinct).
     *
     * @param txBss the set of BSSes that transmitted a Trigger Frame
     * @param txChannelWidth the given transmission channel width
     */
    void CheckReceivedTriggerFrames(std::set<uint8_t> txBss, MHz_u txChannelWidth);

  private:
    void DoSetup() override;
    void DoRun() override;

    WifiStandard m_standard;                      ///< standard
    MHz_u m_channelWidth;                         ///< operating channel width
    bool m_useDistinctBssColors;                  ///< true to set distinct BSS colors to BSSes
    uint8_t m_nBss;                               ///< number of BSSes
    uint8_t m_nStationsPerBss;                    ///< number of stations per AP
    std::vector<NetDeviceContainer> m_staDevices; ///< containers for stations' NetDevices
    NetDeviceContainer m_apDevices;               ///< container for AP's NetDevice
    std::vector<std::bitset<144>> m_received; /**< whether the last packet transmitted to/from each
                                                of  the (up to 144 per BSS) stations was received */
    std::vector<std::bitset<144>>
        m_processed;                /**< whether the last packet transmitted to/from each
                                      of the (up to 144 per BSS) stations was processed */
    Time m_time;                    ///< the time when the current action is executed
    Ptr<WifiPsdu> m_trigger;        ///< Basic Trigger Frame
    WifiTxVector m_triggerTxVector; ///< TX vector for Basic Trigger Frame
    Time m_triggerTxDuration;       ///< TX duration for Basic Trigger Frame
};

WifiPrimaryChannelsTest::WifiPrimaryChannelsTest(WifiStandard standard,
                                                 MHz_u channelWidth,
                                                 bool useDistinctBssColors)
    : TestCase("Check correct transmissions for various primary channel settings"),
      m_standard(standard),
      m_channelWidth(channelWidth),
      m_useDistinctBssColors(useDistinctBssColors)
{
}

void
WifiPrimaryChannelsTest::Transmit(std::string context,
                                  WifiConstPsduMap psduMap,
                                  WifiTxVector txVector,
                                  double txPowerW)
{
    for (const auto& psduPair : psduMap)
    {
        std::stringstream ss;

        if (psduPair.first != SU_STA_ID)
        {
            ss << " STA-ID " << psduPair.first;
        }
        ss << " " << psduPair.second->GetHeader(0).GetTypeString() << " seq "
           << psduPair.second->GetHeader(0).GetSequenceNumber() << " from "
           << psduPair.second->GetAddr2() << " to " << psduPair.second->GetAddr1();
        NS_LOG_INFO(ss.str());
    }
    NS_LOG_INFO(" TXVECTOR " << txVector);
}

void
WifiPrimaryChannelsTest::ReceiveDl(uint8_t bss,
                                   uint8_t station,
                                   Ptr<const WifiPsdu> psdu,
                                   RxSignalInfo rxSignalInfo,
                                   const WifiTxVector& txVector,
                                   const std::vector<bool>& perMpduStatus)
{
    if (psdu->GetNMpdus() == 1)
    {
        const WifiMacHeader& hdr = psdu->GetHeader(0);

        if (hdr.IsQosData() || hdr.IsTrigger())
        {
            NS_LOG_INFO("RECEIVED BY BSS=" << +bss << " STA=" << +station << "  " << *psdu);
            // the MAC received a PSDU from the PHY
            NS_TEST_EXPECT_MSG_EQ(m_received[bss].test(station),
                                  false,
                                  "Station [" << +bss << "][" << +station
                                              << "] received a frame twice");
            m_received[bss].set(station);
            // check if we are the intended destination of the PSDU
            auto dev = DynamicCast<WifiNetDevice>(m_staDevices[bss].Get(station));
            if ((hdr.IsQosData() && hdr.GetAddr1() == dev->GetMac()->GetAddress()) ||
                (hdr.IsTrigger() && hdr.GetAddr1() == Mac48Address::GetBroadcast()))
            {
                NS_TEST_EXPECT_MSG_EQ(m_processed[bss].test(station),
                                      false,
                                      "Station [" << +bss << "][" << +station
                                                  << "] processed a frame twice");
                m_processed[bss].set(station);
            }
        }
    }
}

void
WifiPrimaryChannelsTest::ReceiveUl(uint8_t bss,
                                   Ptr<const WifiPsdu> psdu,
                                   RxSignalInfo rxSignalInfo,
                                   const WifiTxVector& txVector,
                                   const std::vector<bool>& perMpduStatus)
{
    // if the BSS color is zero, this AP might receive the frame sent by another AP. Given that
    // stations only send TB PPDUs, we ignore this frame if the TX vector is not UL MU.
    if (psdu->GetNMpdus() == 1 && psdu->GetHeader(0).IsQosData() && txVector.IsUlMu())
    {
        auto dev = DynamicCast<WifiNetDevice>(m_apDevices.Get(bss));

        uint16_t staId = txVector.GetHeMuUserInfoMap().begin()->first;
        uint8_t station = staId - 1;
        NS_LOG_INFO("RECEIVED FROM BSSID=" << psdu->GetHeader(0).GetAddr3() << " STA=" << +station
                                           << "  " << *psdu);
        // the MAC received a PSDU containing a QoS data frame from the PHY
        NS_TEST_EXPECT_MSG_EQ(m_received[bss].test(station),
                              false,
                              "AP of BSS " << +bss << " received a frame from station " << +station
                                           << " twice");
        m_received[bss].set(station);
        // check if we are the intended destination of the PSDU
        if (psdu->GetHeader(0).GetAddr1() == dev->GetMac()->GetAddress())
        {
            NS_TEST_EXPECT_MSG_EQ(m_processed[bss].test(station),
                                  false,
                                  "AP of BSS " << +bss << " received a frame from station "
                                               << +station << " twice");
            m_processed[bss].set(station);
        }
    }
}

void
WifiPrimaryChannelsTest::DoSetup()
{
    RngSeedManager::SetSeed(1);
    RngSeedManager::SetRun(40);
    int64_t streamNumber = 100;
    uint8_t channelNum;

    // we create as many stations per BSS as the number of 26-tone RUs in a channel
    // of the configured width
    switch (static_cast<uint16_t>(m_channelWidth))
    {
    case 20:
        m_nStationsPerBss = 9;
        channelNum = 1;
        break;
    case 40:
        m_nStationsPerBss = 18;
        channelNum = 3;
        break;
    case 80:
        m_nStationsPerBss = (m_standard == WIFI_STANDARD_80211ax) ? 37 : 36;
        channelNum = 7;
        break;
    case 160:
        m_nStationsPerBss = (m_standard == WIFI_STANDARD_80211ax) ? 74 : 72;
        channelNum = 15;
        break;
    case 320:
        m_nStationsPerBss = 144;
        channelNum = 31;
        break;
    default:
        NS_ABORT_MSG("Channel width (" << m_channelWidth << ") not allowed");
    }

    // we create as many BSSes as the number of 20 MHz subchannels
    m_nBss = Count20MHzSubchannels(m_channelWidth);

    NodeContainer wifiApNodes;
    wifiApNodes.Create(m_nBss);

    std::vector<NodeContainer> wifiStaNodes(m_nBss);
    for (auto& container : wifiStaNodes)
    {
        container.Create(m_nStationsPerBss);
    }

    Ptr<MultiModelSpectrumChannel> spectrumChannel = CreateObject<MultiModelSpectrumChannel>();
    Ptr<FriisPropagationLossModel> lossModel = CreateObject<FriisPropagationLossModel>();
    spectrumChannel->AddPropagationLossModel(lossModel);
    Ptr<ConstantSpeedPropagationDelayModel> delayModel =
        CreateObject<ConstantSpeedPropagationDelayModel>();
    spectrumChannel->SetPropagationDelayModel(delayModel);

    SpectrumWifiPhyHelper phy;
    phy.SetChannel(spectrumChannel);

    WifiHelper wifi;
    wifi.SetStandard(m_standard);
    wifi.SetRemoteStationManager("ns3::ConstantRateWifiManager");

    WifiMacHelper mac;
    mac.SetType("ns3::StaWifiMac",
                "Ssid",
                SsidValue(Ssid("non-existent-ssid")),
                "MaxMissedBeacons",
                UintegerValue(20),
                "WaitBeaconTimeout",
                TimeValue(MicroSeconds(102400))); // same as BeaconInterval

    WifiPhy::ChannelSettingsValue channelValue;

    // Each BSS uses a distinct primary20 channel
    for (uint8_t bss = 0; bss < m_nBss; bss++)
    {
        channelValue.Set(
            WifiPhy::ChannelSegments{{channelNum, m_channelWidth, WIFI_PHY_BAND_6GHZ, bss}});
        phy.Set("ChannelSettings", channelValue);

        m_staDevices.push_back(wifi.Install(phy, mac, wifiStaNodes[bss]));
    }

    for (uint8_t bss = 0; bss < m_nBss; bss++)
    {
        channelValue.Set(
            WifiPhy::ChannelSegments{{channelNum, m_channelWidth, WIFI_PHY_BAND_6GHZ, bss}});
        phy.Set("ChannelSettings", channelValue);

        mac.SetType("ns3::ApWifiMac",
                    "Ssid",
                    SsidValue(Ssid("wifi-ssid-" + std::to_string(bss))),
                    "BeaconInterval",
                    TimeValue(MicroSeconds(102400)),
                    "EnableBeaconJitter",
                    BooleanValue(false));

        m_apDevices.Add(wifi.Install(phy, mac, wifiApNodes.Get(bss)));
    }

    // Assign fixed streams to random variables in use
    streamNumber = WifiHelper::AssignStreams(m_apDevices, streamNumber);
    for (uint8_t bss = 0; bss < m_nBss; bss++)
    {
        streamNumber = WifiHelper::AssignStreams(m_staDevices[bss], streamNumber);
    }

    // set BSS color
    if (m_useDistinctBssColors)
    {
        for (uint8_t bss = 0; bss < m_nBss; bss++)
        {
            auto dev = DynamicCast<WifiNetDevice>(m_apDevices.Get(bss));
            dev->GetHeConfiguration()->m_bssColor = bss + 1;
        }
    }

    MobilityHelper mobility;
    Ptr<ListPositionAllocator> positionAlloc = CreateObject<ListPositionAllocator>();

    positionAlloc->Add(Vector(0.0, 0.0, 0.0)); // all stations are co-located
    mobility.SetPositionAllocator(positionAlloc);

    mobility.SetMobilityModel("ns3::ConstantPositionMobilityModel");
    mobility.Install(wifiApNodes);
    for (uint8_t bss = 0; bss < m_nBss; bss++)
    {
        mobility.Install(wifiStaNodes[bss]);
    }

    m_received.resize(m_nBss);
    m_processed.resize(m_nBss);

    // pre-compute the Basic Trigger Frame to send
    auto apDev = DynamicCast<WifiNetDevice>(m_apDevices.Get(0));

    WifiMacHeader hdr;
    hdr.SetType(WIFI_MAC_CTL_TRIGGER);
    hdr.SetAddr1(Mac48Address::GetBroadcast());
    // Addr2 has to be set
    hdr.SetSequenceNumber(1);

    Ptr<Packet> pkt = Create<Packet>();
    CtrlTriggerHeader trigger;
    trigger.SetType(TriggerFrameType::BASIC_TRIGGER);
    pkt->AddHeader(trigger);

    m_triggerTxVector = WifiTxVector(OfdmPhy::GetOfdmRate6Mbps(),
                                     WIFI_MIN_TX_PWR_LEVEL,
                                     WIFI_PREAMBLE_LONG,
                                     NanoSeconds(800),
                                     1,
                                     1,
                                     0,
                                     MHz_u{20},
                                     false,
                                     false,
                                     false);
    m_trigger = Create<WifiPsdu>(pkt, hdr);

    m_triggerTxDuration = WifiPhy::CalculateTxDuration(m_trigger->GetSize(),
                                                       m_triggerTxVector,
                                                       apDev->GetMac()->GetWifiPhy()->GetPhyBand());
}

void
WifiPrimaryChannelsTest::DoRun()
{
    // schedule association requests at different times. One station's SSID is
    // set to the correct value before initialization, so that such a station
    // starts the scanning procedure by looking for the correct SSID
    Ptr<WifiNetDevice> dev;

    // association can be done in parallel over the multiple BSSes
    for (uint8_t bss = 0; bss < m_nBss; bss++)
    {
        dev = DynamicCast<WifiNetDevice>(m_staDevices[bss].Get(0));
        dev->GetMac()->SetSsid(Ssid("wifi-ssid-" + std::to_string(bss)));

        for (uint16_t i = 1; i < m_nStationsPerBss; i++)
        {
            dev = DynamicCast<WifiNetDevice>(m_staDevices[bss].Get(i));
            Simulator::Schedule(i * MicroSeconds(102400),
                                &WifiMac::SetSsid,
                                dev->GetMac(),
                                Ssid("wifi-ssid-" + std::to_string(bss)));
        }
    }

    // just before sending the beacon preceding the last association, increase the beacon
    // interval (to the max allowed value) so that beacons do not interfere with data frames
    for (uint8_t bss = 0; bss < m_nBss; bss++)
    {
        dev = DynamicCast<WifiNetDevice>(m_apDevices.Get(bss));
        auto mac = DynamicCast<ApWifiMac>(dev->GetMac());

        Simulator::Schedule((m_nStationsPerBss - 1) * MicroSeconds(102400),
                            &ApWifiMac::SetBeaconInterval,
                            mac,
                            MicroSeconds(1024 * 65535));
    }

    m_time = (m_nStationsPerBss + 1) * MicroSeconds(102400);

    Simulator::Schedule(m_time, &WifiPrimaryChannelsTest::CheckAssociation, this);

    // we are done with association. We now intercept frames received by the
    // PHY layer on stations and APs, which will no longer be passed to the FEM.
    for (uint8_t bss = 0; bss < m_nBss; bss++)
    {
        for (uint8_t i = 0; i < m_nStationsPerBss; i++)
        {
            auto dev = DynamicCast<WifiNetDevice>(m_staDevices[bss].Get(i));
            Simulator::Schedule(
                m_time,
                &WifiPhy::SetReceiveOkCallback,
                dev->GetPhy(),
                MakeCallback(&WifiPrimaryChannelsTest::ReceiveDl, this).Bind(bss, i));
        }
        auto dev = DynamicCast<WifiNetDevice>(m_apDevices.Get(bss));
        Simulator::Schedule(m_time,
                            &WifiPhy::SetReceiveOkCallback,
                            dev->GetPhy(),
                            MakeCallback(&WifiPrimaryChannelsTest::ReceiveUl, this).Bind(bss));
    }

    /*
     * We start generating (downlink) SU PPDUs.
     * First, APs operating on non-adjacent primary20 channels send a frame simultaneously
     * in their primary20. This is done in two rounds. As an example, we consider the
     * case of an 160 MHz operating channel:
     *
     *   AP0         AP2         AP4         AP6
     * |-----|     |-----|     |-----|     |-----|     |
     *
     *         AP1         AP3         AP5         AP7
     * |     |-----|     |-----|     |-----|     |-----|
     *
     * Then, we double the transmission channel width. We will have four rounds
     * of transmissions. We avoid using adjacent channels to avoid interference
     * among transmissions:
     *
     *      AP0                     AP4
     * |-----------|           |-----------|           |
     *      AP1                     AP5
     * |-----------|           |-----------|           |
     *                  AP2                     AP6
     * |           |-----------|           |-----------|
     *                  AP3                     AP7
     * |           |-----------|           |-----------|
     *
     * We double the transmission channel width again. We will have eight rounds
     * of transmissions:
     *
     *            AP0
     * |-----------------------|                       |
     *            AP1
     * |-----------------------|                       |
     *            AP2
     * |-----------------------|                       |
     *            AP3
     * |-----------------------|                       |
     *                                    AP4
     * |                       |-----------------------|
     *                                    AP5
     * |                       |-----------------------|
     *                                    AP6
     * |                       |-----------------------|
     *                                    AP7
     * |                       |-----------------------|
     *
     * We double the transmission channel width again. We will have eight rounds
     * of transmissions:
     *
     *                        AP0
     * |-----------------------------------------------|
     *                        AP1
     * |-----------------------------------------------|
     *                        AP2
     * |-----------------------------------------------|
     *                        AP3
     * |-----------------------------------------------|
     *                        AP4
     * |-----------------------------------------------|
     *                        AP5
     * |-----------------------------------------------|
     *                        AP6
     * |-----------------------------------------------|
     *                        AP7
     * |-----------------------------------------------|
     *
     * The transmission channel width reached the operating channel width, we are done.
     */

    Time roundDuration = MilliSeconds(5); // upper bound on the duration of a round

    // To have simultaneous transmissions on adjacent channels, just initialize
    // nRounds to 1 and nApsPerRound to m_channelWidth / 20. Of course, the test
    // will fail because some stations will not receive some frames due to interference
    uint16_t nRounds = 2;
    uint16_t nApsPerRound = Count20MHzSubchannels(m_channelWidth) / 2;
    for (MHz_u txChannelWidth{20}; txChannelWidth <= m_channelWidth;
         txChannelWidth *= 2, nRounds *= 2, nApsPerRound /= 2)
    {
        nRounds = std::min<uint16_t>(nRounds, m_nBss);
        nApsPerRound = std::max<uint16_t>(nApsPerRound, 1);

        for (uint16_t round = 0; round < nRounds; round++)
        {
            std::set<uint8_t> txBss;

            for (uint16_t i = 0; i < nApsPerRound; i++)
            {
                uint16_t ap = round + i * nRounds;
                txBss.insert(ap);
                Simulator::Schedule(m_time,
                                    &WifiPrimaryChannelsTest::SendDlSuPpdu,
                                    this,
                                    ap,
                                    txChannelWidth);
            }
            // check that the SU frames were correctly received
            Simulator::Schedule(m_time + roundDuration,
                                &WifiPrimaryChannelsTest::CheckReceivedSuPpdus,
                                this,
                                txBss,
                                txChannelWidth);
            m_time += roundDuration;
        }
    }

    /*
     * Repeat the same scheme as before with DL MU transmissions. For each transmission
     * channel width, every round is repeated as many times as the number of ways in
     * which we can partition the transmission channel width in equal sized RUs.
     */
    nRounds = 2;
    nApsPerRound = Count20MHzSubchannels(m_channelWidth) / 2;
    const auto mc = GetModulationClassForStandard(m_standard);
    for (MHz_u txChannelWidth{20}; txChannelWidth <= m_channelWidth;
         txChannelWidth *= 2, nRounds *= 2, nApsPerRound /= 2)
    {
        nRounds = std::min<uint16_t>(nRounds, m_nBss);
        nApsPerRound = std::max<uint16_t>(nApsPerRound, 1);

        for (uint16_t round = 0; round < nRounds; round++)
        {
            for (uint32_t type = 0; type < static_cast<uint32_t>(WifiRu::GetMaxRuType(mc)); ++type)
            {
                auto ruType = static_cast<RuType>(type);
                std::size_t nRus = WifiRu::GetNRus(txChannelWidth, ruType, mc);
                std::set<uint8_t> txBss;
                if (nRus > 0)
                {
                    for (uint16_t i = 0; i < nApsPerRound; i++)
                    {
                        uint16_t ap = round + i * nRounds;
                        txBss.insert(ap);
                        Simulator::Schedule(m_time,
                                            &WifiPrimaryChannelsTest::SendDlMuPpdu,
                                            this,
                                            ap,
                                            txChannelWidth,
                                            ruType,
                                            nRus);
                    }
                    // check that the MU frame was correctly received
                    Simulator::Schedule(m_time + roundDuration,
                                        &WifiPrimaryChannelsTest::CheckReceivedMuPpdus,
                                        this,
                                        txBss,
                                        txChannelWidth,
                                        ruType,
                                        nRus,
                                        /* isDlMu */ true);
                    m_time += roundDuration;
                }
            }
        }
    }

    /*
     * Repeat the same scheme as before with UL MU transmissions. For each transmission
     * channel width, every round is repeated as many times as the number of ways in
     * which we can partition the transmission channel width in equal sized RUs.
     */
    nRounds = 2;
    nApsPerRound = Count20MHzSubchannels(m_channelWidth) / 2;
    for (MHz_u txChannelWidth{20}; txChannelWidth <= m_channelWidth;
         txChannelWidth *= 2, nRounds *= 2, nApsPerRound /= 2)
    {
        nRounds = std::min<uint16_t>(nRounds, m_nBss);
        nApsPerRound = std::max<uint16_t>(nApsPerRound, 1);

        for (uint16_t round = 0; round < nRounds; round++)
        {
            for (uint32_t type = 0; type < static_cast<uint32_t>(WifiRu::GetMaxRuType(mc)); ++type)
            {
                auto ruType = static_cast<RuType>(type);
                std::size_t nRus = WifiRu::GetNRus(txChannelWidth, ruType, mc);
                std::set<uint8_t> txBss;
                if (nRus > 0)
                {
                    for (uint16_t i = 0; i < nApsPerRound; i++)
                    {
                        uint16_t ap = round + i * nRounds;
                        txBss.insert(ap);
                        Simulator::Schedule(m_time,
                                            &WifiPrimaryChannelsTest::SendHeTbPpdu,
                                            this,
                                            ap,
                                            txChannelWidth,
                                            ruType,
                                            nRus);
                    }
                    // check that Trigger Frames and TB PPDUs were correctly received
                    Simulator::Schedule(m_time + m_triggerTxDuration +
                                            MicroSeconds(10), /* during SIFS */
                                        &WifiPrimaryChannelsTest::CheckReceivedTriggerFrames,
                                        this,
                                        txBss,
                                        txChannelWidth);
                    Simulator::Schedule(m_time + roundDuration,
                                        &WifiPrimaryChannelsTest::CheckReceivedMuPpdus,
                                        this,
                                        txBss,
                                        txChannelWidth,
                                        ruType,
                                        nRus,
                                        /* isDlMu */ false);
                    m_time += roundDuration;
                }
            }
        }
    }

    // Trace PSDUs passed to the PHY on all devices
    Config::Connect("/NodeList/*/DeviceList/*/$ns3::WifiNetDevice/Phy/PhyTxPsduBegin",
                    MakeCallback(&WifiPrimaryChannelsTest::Transmit, this));

    Simulator::Stop(m_time);
    Simulator::Run();

    Simulator::Destroy();
}

void
WifiPrimaryChannelsTest::SendDlSuPpdu(uint8_t bss, MHz_u txChannelWidth)
{
    NS_LOG_INFO("*** BSS " << +bss << " transmits on primary " << txChannelWidth << " MHz channel");

    auto apDev = DynamicCast<WifiNetDevice>(m_apDevices.Get(bss));
    auto staDev = DynamicCast<WifiNetDevice>(m_staDevices[bss].Get(0));

    uint8_t bssColor = apDev->GetHeConfiguration()->m_bssColor;
    auto txVector = WifiTxVector(
        (m_standard == WIFI_STANDARD_80211ax) ? HePhy::GetHeMcs8() : EhtPhy::GetEhtMcs8(),
        WIFI_MIN_TX_PWR_LEVEL,
        (m_standard == WIFI_STANDARD_80211ax) ? WIFI_PREAMBLE_HE_SU : WIFI_PREAMBLE_EHT_MU,
        NanoSeconds(800),
        1,
        1,
        0,
        txChannelWidth,
        false,
        false,
        false,
        bssColor);
    WifiMacHeader hdr;
    hdr.SetType(WIFI_MAC_QOSDATA);
    hdr.SetQosTid(0);
    hdr.SetAddr1(staDev->GetMac()->GetAddress());
    hdr.SetAddr2(apDev->GetMac()->GetAddress());
    hdr.SetAddr3(apDev->GetMac()->GetBssid(0));
    hdr.SetSequenceNumber(1);
    Ptr<WifiPsdu> psdu = Create<WifiPsdu>(Create<Packet>(1000), hdr);
    apDev->GetPhy()->Send(WifiConstPsduMap({std::make_pair(SU_STA_ID, psdu)}), txVector);
}

void
WifiPrimaryChannelsTest::SendDlMuPpdu(uint8_t bss,
                                      MHz_u txChannelWidth,
                                      RuType ruType,
                                      std::size_t nRus)
{
    NS_LOG_INFO("*** BSS " << +bss << " transmits on primary " << txChannelWidth
                           << " MHz channel a DL MU PPDU "
                           << "addressed to " << nRus << " stations (RU type: " << ruType << ")");

    auto apDev = DynamicCast<WifiNetDevice>(m_apDevices.Get(bss));
    uint8_t bssColor = apDev->GetHeConfiguration()->m_bssColor;

    auto txVector = WifiTxVector(
        (m_standard == WIFI_STANDARD_80211ax) ? HePhy::GetHeMcs8() : EhtPhy::GetEhtMcs8(),
        WIFI_MIN_TX_PWR_LEVEL,
        (m_standard == WIFI_STANDARD_80211ax) ? WIFI_PREAMBLE_HE_MU : WIFI_PREAMBLE_EHT_MU,
        NanoSeconds(800),
        1,
        1,
        0,
        txChannelWidth,
        false,
        false,
        false,
        bssColor);
    if (m_standard == WIFI_STANDARD_80211be)
    {
        txVector.SetEhtPpduType(0);
    }
    WifiMacHeader hdr;
    hdr.SetType(WIFI_MAC_QOSDATA);
    hdr.SetQosTid(0);
    hdr.SetAddr2(apDev->GetMac()->GetAddress());
    hdr.SetAddr3(apDev->GetMac()->GetBssid(0));
    hdr.SetSequenceNumber(1);

    WifiConstPsduMap psduMap;

    auto numRus{nRus};
    if ((m_standard != WIFI_STANDARD_80211ax) && (ruType == RuType::RU_26_TONE))
    {
        // for the loop, take undefined RUs for EHT into account
        numRus += (txChannelWidth / MHz_u{80});
    }
    std::size_t staIdx{0};
    for (std::size_t i = 1; i <= numRus; ++i)
    {
        auto primary80OrLow80{true};
        auto primary160{true};
        auto index{i};
        if (txChannelWidth > MHz_u{80})
        {
            if (m_standard == WIFI_STANDARD_80211ax)
            {
                index = HeRu::GetIndexIn80MHzSegment(txChannelWidth, ruType, i);
                primary80OrLow80 = HeRu::GetPrimary80MHzFlag(txChannelWidth,
                                                             ruType,
                                                             i,
                                                             apDev->GetPhy()->GetPrimary20Index());
            }
            else
            {
                index = EhtRu::GetIndexIn80MHzSegment(txChannelWidth, ruType, i);
                const auto& [p160, p80OrLow80] =
                    EhtRu::GetPrimaryFlags(txChannelWidth,
                                           ruType,
                                           i,
                                           apDev->GetPhy()->GetPrimary20Index());
                primary160 = p160;
                primary80OrLow80 = p80OrLow80;
            }
        }
        if ((m_standard != WIFI_STANDARD_80211ax) && (txChannelWidth >= MHz_u{80}) &&
            (ruType == RuType::RU_26_TONE) && (index == 19))
        {
            // skip undefined RU for EHT
            continue;
        }
        const auto ru =
            (m_standard == WIFI_STANDARD_80211ax)
                ? WifiRu::RuSpec(HeRu::RuSpec{ruType, index, primary80OrLow80})
                : WifiRu::RuSpec(EhtRu::RuSpec{ruType, index, primary160, primary80OrLow80});

        NS_ABORT_IF((bss >= m_staDevices.size()) || (staIdx >= m_staDevices[bss].GetN()));
        auto staDev = DynamicCast<WifiNetDevice>(m_staDevices[bss].Get(staIdx));
        ++staIdx;

        NS_ABORT_IF(!staDev);
        const auto staId = DynamicCast<StaWifiMac>(staDev->GetMac())->GetAssociationId();

        txVector.SetHeMuUserInfo(staId, {ru, 8, 1});
        hdr.SetAddr1(staDev->GetMac()->GetAddress());
        psduMap[staId] = Create<const WifiPsdu>(Create<Packet>(1000), hdr);
    }
    txVector.SetSigBMode(VhtPhy::GetVhtMcs5());

    apDev->GetPhy()->Send(psduMap, txVector);
}

void
WifiPrimaryChannelsTest::SendHeTbPpdu(uint8_t bss,
                                      MHz_u txChannelWidth,
                                      RuType ruType,
                                      std::size_t nRus)
{
    NS_LOG_INFO("*** BSS " << +bss << " transmits a Basic Trigger Frame");

    auto apDev = DynamicCast<WifiNetDevice>(m_apDevices.Get(bss));

    m_trigger->GetHeader(0).SetAddr2(apDev->GetMac()->GetAddress());

    apDev->GetPhy()->Send(m_trigger, m_triggerTxVector);

    // schedule the transmission of HE TB PPDUs
    Simulator::Schedule(m_triggerTxDuration + apDev->GetPhy()->GetSifs(),
                        &WifiPrimaryChannelsTest::DoSendHeTbPpdu,
                        this,
                        bss,
                        txChannelWidth,
                        ruType,
                        nRus);
}

void
WifiPrimaryChannelsTest::DoSendHeTbPpdu(uint8_t bss,
                                        MHz_u txChannelWidth,
                                        RuType ruType,
                                        std::size_t nRus)
{
    auto apDev = DynamicCast<WifiNetDevice>(m_apDevices.Get(bss));
    uint8_t bssColor = apDev->GetHeConfiguration()->m_bssColor;

    WifiMacHeader hdr;
    hdr.SetType(WIFI_MAC_QOSDATA);
    hdr.SetQosTid(0);
    hdr.SetAddr1(apDev->GetMac()->GetAddress());
    hdr.SetAddr3(apDev->GetMac()->GetBssid(0));
    hdr.SetSequenceNumber(1);

    Time duration;
    uint16_t length = 0;
    WifiTxVector trigVector(
        (m_standard == WIFI_STANDARD_80211ax) ? HePhy::GetHeMcs8() : EhtPhy::GetEhtMcs8(),
        WIFI_MIN_TX_PWR_LEVEL,
        (m_standard == WIFI_STANDARD_80211ax) ? WIFI_PREAMBLE_HE_TB : WIFI_PREAMBLE_EHT_TB,
        NanoSeconds(3200),
        1,
        1,
        0,
        txChannelWidth,
        false,
        false,
        false,
        bssColor);

    auto numRus{nRus};
    if ((m_standard != WIFI_STANDARD_80211ax) && (ruType == RuType::RU_26_TONE))
    {
        // for the loop, take undefined RUs for EHT into account
        numRus += (txChannelWidth / MHz_u{80});
    }
    std::size_t staIdx{0};
    for (std::size_t i = 1; i <= numRus; ++i)
    {
        NS_LOG_INFO("*** BSS " << +bss << " STA " << i - 1 << " transmits on primary "
                               << txChannelWidth
                               << " MHz channel an HE TB PPDU (RU type: " << ruType << ")");

        auto primary80OrLow80{true};
        auto primary160{true};
        auto index{i};
        if (txChannelWidth > MHz_u{80})
        {
            if (m_standard == WIFI_STANDARD_80211ax)
            {
                index = HeRu::GetIndexIn80MHzSegment(txChannelWidth, ruType, i);
                primary80OrLow80 = HeRu::GetPrimary80MHzFlag(txChannelWidth,
                                                             ruType,
                                                             i,
                                                             apDev->GetPhy()->GetPrimary20Index());
            }
            else
            {
                index = EhtRu::GetIndexIn80MHzSegment(txChannelWidth, ruType, i);
                const auto& [p160, p80OrLow80] =
                    EhtRu::GetPrimaryFlags(txChannelWidth,
                                           ruType,
                                           i,
                                           apDev->GetPhy()->GetPrimary20Index());
                primary160 = p160;
                primary80OrLow80 = p80OrLow80;
            }
        }
        if ((m_standard != WIFI_STANDARD_80211ax) && (txChannelWidth >= MHz_u{80}) &&
            (ruType == RuType::RU_26_TONE) && (index == 19))
        {
            // skip undefined RU for EHT
            continue;
        }
        const auto ru =
            (m_standard == WIFI_STANDARD_80211ax)
                ? WifiRu::RuSpec(HeRu::RuSpec{ruType, index, primary80OrLow80})
                : WifiRu::RuSpec(EhtRu::RuSpec{ruType, index, primary160, primary80OrLow80});

        NS_ABORT_IF((bss >= m_staDevices.size()) || (staIdx >= m_staDevices[bss].GetN()));
        auto staDev = DynamicCast<WifiNetDevice>(m_staDevices[bss].Get(staIdx));
        ++staIdx;

        const auto staId = DynamicCast<StaWifiMac>(staDev->GetMac())->GetAssociationId();
        WifiTxVector txVector(
            (m_standard == WIFI_STANDARD_80211ax) ? HePhy::GetHeMcs8() : EhtPhy::GetEhtMcs8(),
            WIFI_MIN_TX_PWR_LEVEL,
            (m_standard == WIFI_STANDARD_80211ax) ? WIFI_PREAMBLE_HE_TB : WIFI_PREAMBLE_EHT_TB,
            NanoSeconds(3200),
            1,
            1,
            0,
            txChannelWidth,
            false,
            false,
            false,
            bssColor);
        txVector.SetHeMuUserInfo(staId, {ru, 8, 1});
        trigVector.SetHeMuUserInfo(staId, {ru, 8, 1});

        hdr.SetAddr2(staDev->GetMac()->GetAddress());
        Ptr<const WifiPsdu> psdu = Create<const WifiPsdu>(Create<Packet>(1000), hdr);

        if (duration.IsZero())
        {
            // calculate just once
            duration = WifiPhy::CalculateTxDuration(psdu->GetSize(),
                                                    txVector,
                                                    staDev->GetMac()->GetWifiPhy()->GetPhyBand(),
                                                    staId);
            std::tie(length, duration) = HePhy::ConvertHeTbPpduDurationToLSigLength(
                duration,
                txVector,
                staDev->GetMac()->GetWifiPhy()->GetPhyBand());
        }
        txVector.SetLength(length);

        staDev->GetPhy()->Send(WifiConstPsduMap{{staId, psdu}}, txVector);
    }

    // AP's PHY expects to receive a TRIGVECTOR (just once)
    trigVector.SetLength(length);
    auto apHePhy = std::static_pointer_cast<HePhy>(apDev->GetPhy()->GetLatestPhyEntity());
    apHePhy->SetTrigVector(trigVector, duration);
}

void
WifiPrimaryChannelsTest::CheckAssociation()
{
    for (uint8_t bss = 0; bss < m_nBss; bss++)
    {
        auto dev = DynamicCast<WifiNetDevice>(m_apDevices.Get(bss));
        auto mac = DynamicCast<ApWifiMac>(dev->GetMac());
        NS_TEST_EXPECT_MSG_EQ(mac->GetStaList(SINGLE_LINK_OP_ID).size(),
                              m_nStationsPerBss,
                              "Not all the stations completed association");
    }
}

void
WifiPrimaryChannelsTest::CheckReceivedSuPpdus(std::set<uint8_t> txBss, MHz_u txChannelWidth)
{
    for (uint8_t bss = 0; bss < m_nBss; bss++)
    {
        if (txBss.find(bss) != txBss.end())
        {
            // Every station in the BSS of an AP that transmitted the frame hears (i.e.,
            // passes to the MAC) the frame
            for (uint8_t sta = 0; sta < m_nStationsPerBss; sta++)
            {
                NS_TEST_EXPECT_MSG_EQ(m_received[bss].test(sta),
                                      true,
                                      "Station [" << +bss << "][" << +sta
                                                  << "] did not receive the SU frame on primary"
                                                  << txChannelWidth << " channel");
            }
            // only the first station actually processed the frames
            NS_TEST_EXPECT_MSG_EQ(m_processed[bss].test(0),
                                  true,
                                  "Station [" << +bss << "][0]"
                                              << " did not process the SU frame on primary"
                                              << txChannelWidth << " channel");
            for (uint8_t sta = 1; sta < m_nStationsPerBss; sta++)
            {
                NS_TEST_EXPECT_MSG_EQ(m_processed[bss].test(sta),
                                      false,
                                      "Station [" << +bss << "][" << +sta
                                                  << "] processed the SU frame on primary"
                                                  << txChannelWidth << " channel");
            }
        }
        else
        {
            // There was no transmission in this BSS. If BSS Color filtering is enabled or no frame
            // transmission overlaps with the primary20 channel of this BSS, stations in this BSS
            // did not hear any frame.
            if (m_useDistinctBssColors ||
                std::none_of(txBss.begin(), txBss.end(), [&](const uint8_t& txAp) {
                    auto txApPhy = DynamicCast<WifiNetDevice>(m_apDevices.Get(txAp))->GetPhy();
                    auto thisApPhy = DynamicCast<WifiNetDevice>(m_apDevices.Get(bss))->GetPhy();
                    return txApPhy->GetOperatingChannel().GetPrimaryChannelIndex(txChannelWidth) ==
                           thisApPhy->GetOperatingChannel().GetPrimaryChannelIndex(txChannelWidth);
                }))
            {
                for (uint8_t sta = 0; sta < m_nStationsPerBss; sta++)
                {
                    NS_TEST_EXPECT_MSG_EQ(m_received[bss].test(sta),
                                          false,
                                          "Station [" << +bss << "][" << +sta
                                                      << "] received the SU frame on primary"
                                                      << txChannelWidth << " channel");
                }
            }
            else
            {
                // all stations heard the frame but no station processed it
                for (uint8_t sta = 0; sta < m_nStationsPerBss; sta++)
                {
                    NS_TEST_EXPECT_MSG_EQ(m_received[bss].test(sta),
                                          true,
                                          "Station [" << +bss << "][" << +sta
                                                      << "] did not receive the SU frame on primary"
                                                      << txChannelWidth << " channel");
                    NS_TEST_EXPECT_MSG_EQ(m_processed[bss].test(sta),
                                          false,
                                          "Station [" << +bss << "][" << +sta
                                                      << "] processed the SU frame on primary"
                                                      << txChannelWidth << " channel");
                }
            }
        }
        // reset bitmaps
        m_received[bss].reset();
        m_processed[bss].reset();
    }
}

void
WifiPrimaryChannelsTest::CheckReceivedMuPpdus(std::set<uint8_t> txBss,
                                              MHz_u txChannelWidth,
                                              RuType ruType,
                                              std::size_t nRus,
                                              bool isDlMu)
{
    for (uint8_t bss = 0; bss < m_nBss; bss++)
    {
        if (txBss.find(bss) != txBss.end())
        {
            // There was a transmission in this BSS.
            // [DL] Due to AID filtering, only stations that are addressed by the MU PPDU do hear
            // the frame [UL] The AP hears a TB PPDU sent by all and only the solicited stations
            for (std::size_t sta = 0; sta < nRus; sta++)
            {
                NS_TEST_EXPECT_MSG_EQ(
                    m_received[bss].test(sta),
                    true,
                    (isDlMu ? "A DL MU PPDU transmitted to" : "An HE TB PPDU transmitted by")
                        << " station [" << +bss << "][" << +sta << "] on primary" << txChannelWidth
                        << " channel, RU type " << ruType << " was not received");
            }
            for (uint8_t sta = nRus; sta < m_nStationsPerBss; sta++)
            {
                NS_TEST_EXPECT_MSG_EQ(m_received[bss].test(sta),
                                      false,
                                      (isDlMu ? "A DL MU PPDU" : "An HE TB PPDU")
                                          << " transmitted on primary" << txChannelWidth
                                          << " channel, RU type " << ruType << " was received "
                                          << (isDlMu ? "by" : "from") << " station [" << +bss
                                          << "][" << +sta << "]");
            }
            // [DL] Only the addressed stations actually processed the frames
            // [UL] The AP processed the frames sent by all and only the addressed stations
            for (std::size_t sta = 0; sta < nRus; sta++)
            {
                NS_TEST_EXPECT_MSG_EQ(
                    m_processed[bss].test(sta),
                    true,
                    (isDlMu ? "A DL MU PPDU transmitted to" : "An HE TB PPDU transmitted by")
                        << " station [" << +bss << "][" << +sta << "] on primary" << txChannelWidth
                        << " channel, RU type " << ruType << " was not processed");
            }
            for (uint8_t sta = nRus; sta < m_nStationsPerBss; sta++)
            {
                NS_TEST_EXPECT_MSG_EQ(m_processed[bss].test(sta),
                                      false,
                                      (isDlMu ? "A DL MU PPDU" : "An HE TB PPDU")
                                          << " transmitted on primary" << txChannelWidth
                                          << " channel, RU type " << ruType << " was received "
                                          << (isDlMu ? "by" : "from") << " station [" << +bss
                                          << "][" << +sta << "] and processed");
            }
        }
        else
        {
            // There was no transmission in this BSS.
            // [DL] If BSS Color filtering is enabled or no frame transmission overlaps with
            // the primary20 channel of this BSS, stations in this BSS did not hear any frame.
            // [UL] The AP did not hear any TB PPDU because no TRIGVECTOR was passed to the PHY
            if (!isDlMu || m_useDistinctBssColors ||
                std::none_of(txBss.begin(), txBss.end(), [&](const uint8_t& txAp) {
                    auto txApPhy = DynamicCast<WifiNetDevice>(m_apDevices.Get(txAp))->GetPhy();
                    auto thisApPhy = DynamicCast<WifiNetDevice>(m_apDevices.Get(bss))->GetPhy();
                    return txApPhy->GetOperatingChannel().GetPrimaryChannelIndex(txChannelWidth) ==
                           thisApPhy->GetOperatingChannel().GetPrimaryChannelIndex(txChannelWidth);
                }))
            {
                for (uint8_t sta = 0; sta < m_nStationsPerBss; sta++)
                {
                    NS_TEST_EXPECT_MSG_EQ(m_received[bss].test(sta),
                                          false,
                                          (isDlMu ? "A DL MU PPDU" : "An HE TB PPDU")
                                              << " transmitted on primary" << txChannelWidth
                                              << " channel, RU type " << ruType << " was received "
                                              << (isDlMu ? "by" : "from") << " station [" << +bss
                                              << "][" << +sta << "]");
                }
            }
            else
            {
                // [DL] Stations having the same AID of the stations addressed by the MU PPDU
                // received the frame
                for (std::size_t sta = 0; sta < nRus; sta++)
                {
                    NS_TEST_EXPECT_MSG_EQ(
                        m_received[bss].test(sta),
                        true,
                        (isDlMu ? "A DL MU PPDU transmitted to" : "An HE TB PPDU transmitted by")
                            << " station [" << +bss << "][" << +sta << "] on primary"
                            << txChannelWidth << " channel, RU type " << ruType
                            << " was not received");
                }
                for (uint8_t sta = nRus; sta < m_nStationsPerBss; sta++)
                {
                    NS_TEST_EXPECT_MSG_EQ(m_received[bss].test(sta),
                                          false,
                                          (isDlMu ? "A DL MU PPDU" : "An HE TB PPDU")
                                              << " transmitted on primary" << txChannelWidth
                                              << " channel, RU type " << ruType << " was received "
                                              << (isDlMu ? "by" : "from") << " station [" << +bss
                                              << "][" << +sta << "]");
                }
                // no station processed the frame
                for (uint8_t sta = 0; sta < m_nStationsPerBss; sta++)
                {
                    NS_TEST_EXPECT_MSG_EQ(m_processed[bss].test(sta),
                                          false,
                                          (isDlMu ? "A DL MU PPDU" : "An HE TB PPDU")
                                              << " transmitted on primary" << txChannelWidth
                                              << " channel, RU type " << ruType << " was received "
                                              << (isDlMu ? "by" : "from") << " station [" << +bss
                                              << "][" << +sta << "] and processed");
                }
            }
        }
        // reset bitmaps
        m_received[bss].reset();
        m_processed[bss].reset();
    }
}

void
WifiPrimaryChannelsTest::CheckReceivedTriggerFrames(std::set<uint8_t> txBss, MHz_u txChannelWidth)
{
    for (uint8_t bss = 0; bss < m_nBss; bss++)
    {
        if (txBss.find(bss) != txBss.end())
        {
            // Every station in the BSS of an AP that transmitted the Trigger Frame hears (i.e.,
            // passes to the MAC) and processes the frame
            for (uint8_t sta = 0; sta < m_nStationsPerBss; sta++)
            {
                NS_TEST_EXPECT_MSG_EQ(m_received[bss].test(sta),
                                      true,
                                      "Station [" << +bss << "][" << +sta
                                                  << "] did not receive the Trigger Frame "
                                                     "soliciting a transmission on primary"
                                                  << txChannelWidth << " channel");
                NS_TEST_EXPECT_MSG_EQ(m_processed[bss].test(sta),
                                      true,
                                      "Station [" << +bss << "][" << +sta
                                                  << "] did not process the Trigger Frame "
                                                     "soliciting a transmission on primary"
                                                  << txChannelWidth << " channel");
            }
        }
        else
        {
            // Given that a Trigger Frame is transmitted on the primary20 channel and all the
            // primary20 channels are distinct, stations in other BSSes did not hear the frame
            for (uint8_t sta = 0; sta < m_nStationsPerBss; sta++)
            {
                NS_TEST_EXPECT_MSG_EQ(
                    m_received[bss].test(sta),
                    false,
                    "Station ["
                        << +bss << "][" << +sta
                        << "] received the Trigger Frame soliciting a transmission on primary"
                        << txChannelWidth << " channel");
            }
        }
        // reset bitmaps
        m_received[bss].reset();
        m_processed[bss].reset();
    }
}

/**
 * @ingroup wifi-test
 * @ingroup tests
 *
 * @brief Test functions returning the indices of primary and secondary channels
 *        of different width.
 */
class Wifi20MHzChannelIndicesTest : public TestCase
{
  public:
    /**
     * Constructor
     */
    Wifi20MHzChannelIndicesTest();
    ~Wifi20MHzChannelIndicesTest() override = default;

    /**
     * Check that the indices of the 20 MHz channels included in all the primary
     * and secondary channels are correct when setting the given primary20 channel.
     *
     * @param primary20 the index of the primary20 channel to configure
     * @param secondary20 the expected index of the secondary20 channel
     * @param primary40 the expected indices of the 20 MHz channels in the primary40 channel
     * @param secondary40 the expected indices of the 20 MHz channels in the secondary40 channel
     * @param primary80 the expected indices of the 20 MHz channels in the primary80 channel
     * @param secondary80 the expected indices of the 20 MHz channels in the secondary80 channel
     * @param primary160 the expected indices of the 20 MHz channels in the primary160 channel
     * @param secondary160 the expected indices of the 20 MHz channels in the secondary160 channel
     */
    void RunOne(uint8_t primary20,
                const std::set<uint8_t>& secondary20,
                const std::set<uint8_t>& primary40,
                const std::set<uint8_t>& secondary40,
                const std::set<uint8_t>& primary80,
                const std::set<uint8_t>& secondary80,
                const std::set<uint8_t>& primary160,
                const std::set<uint8_t>& secondary160);

  private:
    void DoRun() override;

    WifiPhyOperatingChannel m_channel; //!< operating channel
};

Wifi20MHzChannelIndicesTest::Wifi20MHzChannelIndicesTest()
    : TestCase("Check computation of primary and secondary channel indices")
{
}

void
Wifi20MHzChannelIndicesTest::RunOne(uint8_t primary20,
                                    const std::set<uint8_t>& secondary20,
                                    const std::set<uint8_t>& primary40,
                                    const std::set<uint8_t>& secondary40,
                                    const std::set<uint8_t>& primary80,
                                    const std::set<uint8_t>& secondary80,
                                    const std::set<uint8_t>& primary160,
                                    const std::set<uint8_t>& secondary160)
{
    auto printToStr = [](const std::set<uint8_t>& s) {
        std::stringstream ss;
        ss << "{";
        for (const auto& index : s)
        {
            ss << +index << " ";
        }
        ss << "}";
        return ss.str();
    };

    m_channel.SetPrimary20Index(primary20);

    auto actualPrimary20 = m_channel.GetAll20MHzChannelIndicesInPrimary(MHz_u{20});
    NS_TEST_ASSERT_MSG_EQ((actualPrimary20 == std::set<uint8_t>{primary20}),
                          true,
                          "Expected Primary20 {" << +primary20 << "}"
                                                 << " differs from actual "
                                                 << printToStr(actualPrimary20));

    auto actualSecondary20 = m_channel.GetAll20MHzChannelIndicesInSecondary(actualPrimary20);
    NS_TEST_ASSERT_MSG_EQ((actualSecondary20 == secondary20),
                          true,
                          "Expected Secondary20 " << printToStr(secondary20)
                                                  << " differs from actual "
                                                  << printToStr(actualSecondary20));

    auto actualPrimary40 = m_channel.GetAll20MHzChannelIndicesInPrimary(MHz_u{40});
    NS_TEST_ASSERT_MSG_EQ((actualPrimary40 == primary40),
                          true,
                          "Expected Primary40 " << printToStr(primary40) << " differs from actual "
                                                << printToStr(actualPrimary40));

    auto actualSecondary40 = m_channel.GetAll20MHzChannelIndicesInSecondary(primary40);
    NS_TEST_ASSERT_MSG_EQ((actualSecondary40 == secondary40),
                          true,
                          "Expected Secondary40 " << printToStr(secondary40)
                                                  << " differs from actual "
                                                  << printToStr(actualSecondary40));

    auto actualPrimary80 = m_channel.GetAll20MHzChannelIndicesInPrimary(MHz_u{80});
    NS_TEST_ASSERT_MSG_EQ((actualPrimary80 == primary80),
                          true,
                          "Expected Primary80 " << printToStr(primary80) << " differs from actual "
                                                << printToStr(actualPrimary80));

    auto actualSecondary80 = m_channel.GetAll20MHzChannelIndicesInSecondary(primary80);
    NS_TEST_ASSERT_MSG_EQ((actualSecondary80 == secondary80),
                          true,
                          "Expected Secondary80 " << printToStr(secondary80)
                                                  << " differs from actual "
                                                  << printToStr(actualSecondary80));

    auto actualPrimary160 = m_channel.GetAll20MHzChannelIndicesInPrimary(MHz_u{160});
    NS_TEST_ASSERT_MSG_EQ((actualPrimary160 == primary160),
                          true,
                          "Expected Primary160 " << printToStr(primary160)
                                                 << " differs from actual "
                                                 << printToStr(actualPrimary160));

    auto actualSecondary160 = m_channel.GetAll20MHzChannelIndicesInSecondary(primary160);
    NS_TEST_ASSERT_MSG_EQ((actualSecondary160 == secondary160),
                          true,
                          "Expected Secondary160 " << printToStr(secondary160)
                                                   << " differs from actual "
                                                   << printToStr(actualSecondary160));
}

void
Wifi20MHzChannelIndicesTest::DoRun()
{
    /* 20 MHz channel */
    m_channel.SetDefault(MHz_u{20}, WIFI_STANDARD_80211be, WIFI_PHY_BAND_6GHZ);
    RunOne(0, {}, {}, {}, {}, {}, {}, {});

    /* 40 MHz channel */
    m_channel.SetDefault(MHz_u{40}, WIFI_STANDARD_80211be, WIFI_PHY_BAND_6GHZ);
    RunOne(0, {1}, {0, 1}, {}, {}, {}, {}, {});
    RunOne(1, {0}, {0, 1}, {}, {}, {}, {}, {});

    /* 80 MHz channel */
    m_channel.SetDefault(MHz_u{80}, WIFI_STANDARD_80211be, WIFI_PHY_BAND_6GHZ);
    RunOne(0, {1}, {0, 1}, {2, 3}, {0, 1, 2, 3}, {}, {}, {});
    RunOne(1, {0}, {0, 1}, {2, 3}, {0, 1, 2, 3}, {}, {}, {});
    RunOne(2, {3}, {2, 3}, {0, 1}, {0, 1, 2, 3}, {}, {}, {});
    RunOne(3, {2}, {2, 3}, {0, 1}, {0, 1, 2, 3}, {}, {}, {});

    /* 160 MHz channel */
    m_channel.SetDefault(MHz_u{160}, WIFI_STANDARD_80211be, WIFI_PHY_BAND_6GHZ);
    RunOne(0, {1}, {0, 1}, {2, 3}, {0, 1, 2, 3}, {4, 5, 6, 7}, {0, 1, 2, 3, 4, 5, 6, 7}, {});
    RunOne(1, {0}, {0, 1}, {2, 3}, {0, 1, 2, 3}, {4, 5, 6, 7}, {0, 1, 2, 3, 4, 5, 6, 7}, {});
    RunOne(2, {3}, {2, 3}, {0, 1}, {0, 1, 2, 3}, {4, 5, 6, 7}, {0, 1, 2, 3, 4, 5, 6, 7}, {});
    RunOne(3, {2}, {2, 3}, {0, 1}, {0, 1, 2, 3}, {4, 5, 6, 7}, {0, 1, 2, 3, 4, 5, 6, 7}, {});
    RunOne(4, {5}, {4, 5}, {6, 7}, {4, 5, 6, 7}, {0, 1, 2, 3}, {0, 1, 2, 3, 4, 5, 6, 7}, {});
    RunOne(5, {4}, {4, 5}, {6, 7}, {4, 5, 6, 7}, {0, 1, 2, 3}, {0, 1, 2, 3, 4, 5, 6, 7}, {});
    RunOne(6, {7}, {6, 7}, {4, 5}, {4, 5, 6, 7}, {0, 1, 2, 3}, {0, 1, 2, 3, 4, 5, 6, 7}, {});
    RunOne(7, {6}, {6, 7}, {4, 5}, {4, 5, 6, 7}, {0, 1, 2, 3}, {0, 1, 2, 3, 4, 5, 6, 7}, {});

    /* 320 MHz channel */
    m_channel.SetDefault(MHz_u{320}, WIFI_STANDARD_80211be, WIFI_PHY_BAND_6GHZ);
    RunOne(0,
           {1},
           {0, 1},
           {2, 3},
           {0, 1, 2, 3},
           {4, 5, 6, 7},
           {0, 1, 2, 3, 4, 5, 6, 7},
           {8, 9, 10, 11, 12, 13, 14, 15});
    RunOne(1,
           {0},
           {0, 1},
           {2, 3},
           {0, 1, 2, 3},
           {4, 5, 6, 7},
           {0, 1, 2, 3, 4, 5, 6, 7},
           {8, 9, 10, 11, 12, 13, 14, 15});
    RunOne(2,
           {3},
           {2, 3},
           {0, 1},
           {0, 1, 2, 3},
           {4, 5, 6, 7},
           {0, 1, 2, 3, 4, 5, 6, 7},
           {8, 9, 10, 11, 12, 13, 14, 15});
    RunOne(3,
           {2},
           {2, 3},
           {0, 1},
           {0, 1, 2, 3},
           {4, 5, 6, 7},
           {0, 1, 2, 3, 4, 5, 6, 7},
           {8, 9, 10, 11, 12, 13, 14, 15});
    RunOne(4,
           {5},
           {4, 5},
           {6, 7},
           {4, 5, 6, 7},
           {0, 1, 2, 3},
           {0, 1, 2, 3, 4, 5, 6, 7},
           {8, 9, 10, 11, 12, 13, 14, 15});
    RunOne(5,
           {4},
           {4, 5},
           {6, 7},
           {4, 5, 6, 7},
           {0, 1, 2, 3},
           {0, 1, 2, 3, 4, 5, 6, 7},
           {8, 9, 10, 11, 12, 13, 14, 15});
    RunOne(6,
           {7},
           {6, 7},
           {4, 5},
           {4, 5, 6, 7},
           {0, 1, 2, 3},
           {0, 1, 2, 3, 4, 5, 6, 7},
           {8, 9, 10, 11, 12, 13, 14, 15});
    RunOne(7,
           {6},
           {6, 7},
           {4, 5},
           {4, 5, 6, 7},
           {0, 1, 2, 3},
           {0, 1, 2, 3, 4, 5, 6, 7},
           {8, 9, 10, 11, 12, 13, 14, 15});
    RunOne(8,
           {9},
           {8, 9},
           {10, 11},
           {8, 9, 10, 11},
           {12, 13, 14, 15},
           {8, 9, 10, 11, 12, 13, 14, 15},
           {0, 1, 2, 3, 4, 5, 6, 7});
    RunOne(9,
           {8},
           {8, 9},
           {10, 11},
           {8, 9, 10, 11},
           {12, 13, 14, 15},
           {8, 9, 10, 11, 12, 13, 14, 15},
           {0, 1, 2, 3, 4, 5, 6, 7});
    RunOne(10,
           {11},
           {10, 11},
           {8, 9},
           {8, 9, 10, 11},
           {12, 13, 14, 15},
           {8, 9, 10, 11, 12, 13, 14, 15},
           {0, 1, 2, 3, 4, 5, 6, 7});
    RunOne(11,
           {10},
           {10, 11},
           {8, 9},
           {8, 9, 10, 11},
           {12, 13, 14, 15},
           {8, 9, 10, 11, 12, 13, 14, 15},
           {0, 1, 2, 3, 4, 5, 6, 7});
    RunOne(12,
           {13},
           {12, 13},
           {14, 15},
           {12, 13, 14, 15},
           {8, 9, 10, 11},
           {8, 9, 10, 11, 12, 13, 14, 15},
           {0, 1, 2, 3, 4, 5, 6, 7});
    RunOne(13,
           {12},
           {12, 13},
           {14, 15},
           {12, 13, 14, 15},
           {8, 9, 10, 11},
           {8, 9, 10, 11, 12, 13, 14, 15},
           {0, 1, 2, 3, 4, 5, 6, 7});
    RunOne(14,
           {15},
           {14, 15},
           {12, 13},
           {12, 13, 14, 15},
           {8, 9, 10, 11},
           {8, 9, 10, 11, 12, 13, 14, 15},
           {0, 1, 2, 3, 4, 5, 6, 7});
    RunOne(15,
           {14},
           {14, 15},
           {12, 13},
           {12, 13, 14, 15},
           {8, 9, 10, 11},
           {8, 9, 10, 11, 12, 13, 14, 15},
           {0, 1, 2, 3, 4, 5, 6, 7});
}

/**
 * @ingroup wifi-test
 * @ingroup tests
 *
 * @brief wifi primary channels test suite
 */
class WifiPrimaryChannelsTestSuite : public TestSuite
{
  public:
    WifiPrimaryChannelsTestSuite();
};

WifiPrimaryChannelsTestSuite::WifiPrimaryChannelsTestSuite()
    : TestSuite("wifi-primary-channels", Type::UNIT)
{
    for (const auto standard : {WIFI_STANDARD_80211ax, WIFI_STANDARD_80211be})
    {
        // Test cases for 20 MHz can be added, but are not that useful (there would be a single BSS)
        AddTestCase(new WifiPrimaryChannelsTest(standard, MHz_u{40}, true),
                    TestCase::Duration::QUICK);
        AddTestCase(new WifiPrimaryChannelsTest(standard, MHz_u{40}, false),
                    TestCase::Duration::QUICK);
        AddTestCase(new WifiPrimaryChannelsTest(standard, MHz_u{80}, true),
                    TestCase::Duration::EXTENSIVE);
        AddTestCase(new WifiPrimaryChannelsTest(standard, MHz_u{80}, false),
                    TestCase::Duration::EXTENSIVE);
        AddTestCase(new WifiPrimaryChannelsTest(standard, MHz_u{160}, true),
                    TestCase::Duration::TAKES_FOREVER);
        AddTestCase(new WifiPrimaryChannelsTest(standard, MHz_u{160}, false),
                    TestCase::Duration::TAKES_FOREVER);
    }
    // TODO: activate WifiPrimaryChannelsTest for 320 MHz channels once PHY operations are optimized
    // to run faster, otherwise run-time is too long
    /*AddTestCase(new WifiPrimaryChannelsTest(WIFI_STANDARD_80211be, MHz_u{320}, true),
                TestCase::Duration::TAKES_FOREVER);
    AddTestCase(new WifiPrimaryChannelsTest(WIFI_STANDARD_80211be, MHz_u{320}, false),
                TestCase::Duration::TAKES_FOREVER);*/
    AddTestCase(new Wifi20MHzChannelIndicesTest(), TestCase::Duration::QUICK);
}

static WifiPrimaryChannelsTestSuite g_wifiPrimaryChannelsTestSuite; ///< the test suite
