/*
 * Copyright (c) 2024
 *
 * SPDX-License-Identifier: GPL-2.0-only
 *
 * Author: Sébastien Deronne <sebastien.deronne@gmail.com>
 */

#ifndef WIFI_RU_H
#define WIFI_RU_H

#include "ns3/eht-ru.h"
#include "ns3/he-ru.h"

#include <variant>

namespace ns3
{

/**
 * This class handles RU variants.
 */
class WifiRu
{
  public:
    /// variant of the RU specification
    using RuSpec = std::variant<HeRu::RuSpec, EhtRu::RuSpec>;

    /**
     * Struct providing a function call operator to compare two RUs.
     */
    struct RuSpecCompare
    {
        /**
         * Constructor.
         *
         * @param channelWidth the channel width
         * @param p20Index the index of the primary20 channel
         */
        RuSpecCompare(MHz_u channelWidth, uint8_t p20Index);
        /**
         * Function call operator.
         *
         * @param lhs left hand side RU
         * @param rhs right hand side RU
         * @return true if the left hand side RU has its leftmost tone at a lower
         *         frequency than the leftmost tone of the right hand side RU,
         *         false otherwise
         */
        bool operator()(const RuSpec& lhs, const RuSpec& rhs) const;

      private:
        MHz_u m_channelWidth; ///< The channel width
        uint8_t m_p20Index;   ///< Primary20 channel index
    };

    /**
     * Get the type of a given RU
     *
     * @param ru the RU variant
     * @return the RU type
     */
    static RuType GetRuType(RuSpec ru);

    /**
     * Get the index of a given RU
     *
     * @param ru the RU variant
     * @return the RU index
     */
    static std::size_t GetIndex(RuSpec ru);

    /**
     * Get the RU PHY index
     *
     * @param ru the RU variant
     * @param bw the width of the channel of which the RU is part
     * @param p20Index the index of the primary20 channel
     * @return the RU PHY index
     */
    static std::size_t GetPhyIndex(RuSpec ru, MHz_u bw, uint8_t p20Index);

    /**
     * Get the approximate bandwidth occupied by a RU.
     *
     * @param ruType the RU type
     * @return the approximate bandwidth occupied by the RU
     */
    static MHz_u GetBandwidth(RuType ruType);

    /**
     * Get the RU corresponding to the approximate bandwidth.
     *
     * @param bandwidth the approximate bandwidth occupied by the RU
     * @return the RU type
     */
    static RuType GetRuType(MHz_u bandwidth);

    /**
     * Get the number of distinct RUs of the given type (number of tones)
     * available in a PPDU of the given bandwidth.
     *
     * @param bw the bandwidth of the PPDU
     * @param ruType the RU type (number of tones)
     * @param mc the modulation class of the PPDU
     * @return the number of distinct RUs available
     */
    static std::size_t GetNRus(MHz_u bw, RuType ruType, WifiModulationClass mc);

    /**
     * Get the subcarrier group of the RU having the given PHY index among all the
     * RUs of the given type (number of tones) available in a PPDU of the
     * given bandwidth. A subcarrier group is defined as one or more pairs
     * indicating the lowest frequency index and the highest frequency index.
     *
     * @param bw the bandwidth of the PPDU
     * @param ruType the RU type (number of tones)
     * @param phyIndex the PHY index (starting at 1) of the RU
     * @param mc the modulation class of the PPDU
     * @return the subcarrier range of the specified RU
     */
    static SubcarrierGroup GetSubcarrierGroup(MHz_u bw,
                                              RuType ruType,
                                              std::size_t phyIndex,
                                              WifiModulationClass mc);

    /**
     * Get the RU_ALLOCATION value for equal size RUs
     *
     * @param ruType equal size RU type (generated by GetEqualSizedRusForStations)
     * @param isOdd if number of stations is an odd number
     * @param hasUsers whether it contributes to User field(s) in the content channel this RU
     * Allocation belongs to
     * @param mc the modulation class of the PPDU
     * @return RU_ALLOCATION value
     */
    static uint16_t GetEqualizedRuAllocation(RuType ruType,
                                             bool isOdd,
                                             bool hasUsers,
                                             WifiModulationClass mc);

    /**
     * Get the RU specs based on RU_ALLOCATION
     *
     * @param ruAllocation 9 bit RU_ALLOCATION value
     * @param mc the modulation class of the PPDU
     * @return RU spec associated with the RU_ALLOCATION
     */
    static std::vector<RuSpec> GetRuSpecs(uint16_t ruAllocation, WifiModulationClass mc);

    /**
     * Get the set of distinct RUs of the given type (number of tones)
     * available in an MU PPDU of the given bandwidth.
     *
     * @param bw the bandwidth of the PPDU
     * @param ruType the RU type (number of tones)
     * @param mc the modulation class of the PPDU
     * @return the set of distinct RUs available
     */
    static std::vector<RuSpec> GetRusOfType(MHz_u bw, RuType ruType, WifiModulationClass mc);

    /**
     * Get the set of 26-tone RUs that can be additionally allocated if the given
     * bandwidth is split in RUs of the given type.
     *
     * @param bw the bandwidth of the PPDU
     * @param ruType the RU type (number of tones)
     * @param mc the modulation class of the PPDU
     * @return the set of 26-tone RUs that can be additionally allocated
     */
    static std::vector<RuSpec> GetCentral26TonesRus(MHz_u bw,
                                                    RuType ruType,
                                                    WifiModulationClass mc);

    /**
     * Check whether the given RU overlaps with the given set of RUs.
     *
     * @param bw the bandwidth of the PPDU
     * @param ru the given RU allocation
     * @param v the given set of RUs
     * @return true if the given RU overlaps with the given set of RUs.
     */
    static bool DoesOverlap(MHz_u bw, RuSpec ru, const std::vector<RuSpec>& v);

    /**
     * Given the channel bandwidth and the number of stations candidate for being
     * assigned an RU, maximize the number of candidate stations that can be assigned
     * an RU subject to the constraint that all the stations must be assigned an RU
     * of the same size (in terms of number of tones).
     *
     * @param bandwidth the channel bandwidth
     * @param nStations the number of candidate stations. On return, it is set to
     *                  the number of stations that are assigned an RU
     * @param nCentral26TonesRus the number of additional 26-tone RUs that can be allocated if the
     * returned RU size is greater than 26 tones
     * @param mc the modulation class of the PPDU
     * @return the RU type
     */
    static RuType GetEqualSizedRusForStations(MHz_u bandwidth,
                                              std::size_t& nStations,
                                              std::size_t& nCentral26TonesRus,
                                              WifiModulationClass mc);

    /**
     * Find the RU allocation of the given RU type overlapping the given
     * reference RU allocation.
     *
     * @param bw the bandwidth of the PPDU
     * @param referenceRu the reference RU allocation
     * @param searchedRuType the searched RU type
     * @return the searched RU allocation.
     */
    static RuSpec FindOverlappingRu(MHz_u bw, RuSpec referenceRu, RuType searchedRuType);

    /**
     * Get whether a given RU variant is a HE RU
     *
     * @param ru the RU variant
     * @return true if the RU is a HE RU, false otherwise
     */
    static bool IsHe(RuSpec ru);

    /**
     * Get whether a given RU variant is a EHT RU
     *
     * @param ru the RU variant
     * @return true if the RU is a EHT RU, false otherwise
     */
    static bool IsEht(RuSpec ru);

    /**
     * Get the maximum RU type supported by a given modulation class.
     *
     * @param mc the modulation class
     * @return the RU type
     */
    static RuType GetMaxRuType(WifiModulationClass mc);
};

/**
 * @brief Stream insertion operator.
 *
 * @param os the stream
 * @param ru the RU variant
 * @returns a reference to the stream
 */
std::ostream& operator<<(std::ostream& os, const WifiRu::RuSpec& ru);

} // namespace ns3

#endif /* WIFI_RU_H */
