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

// This example is used to validate 802.11n MIMO.
//
// It outputs plots of the throughput versus the distance
// for every HT MCS value and from 1 to 4 MIMO streams.
//
// The simulation assumes a single station in an infrastructure network:
//
//  STA     AP
//    *     *
//    |     |
//   n1     n2
//
// The user can choose whether UDP or TCP should be used and can configure
// some 802.11n parameters (frequency, channel width and guard interval).
//
// An important configuration parameter is preamble detection.  It is enabled
// by default (to match the default ns-3 configuration) but will dominate
// performance at low SNRs, causing the different MCS to appear to have
// the same range (because regardless of the MCS, the preamble detection
// thresholds do not change).

#include "ns3/boolean.h"
#include "ns3/command-line.h"
#include "ns3/config.h"
#include "ns3/double.h"
#include "ns3/gnuplot.h"
#include "ns3/ht-phy.h"
#include "ns3/internet-stack-helper.h"
#include "ns3/ipv4-address-helper.h"
#include "ns3/ipv4-global-routing-helper.h"
#include "ns3/mobility-helper.h"
#include "ns3/on-off-helper.h"
#include "ns3/packet-sink-helper.h"
#include "ns3/packet-sink.h"
#include "ns3/ssid.h"
#include "ns3/string.h"
#include "ns3/udp-client-server-helper.h"
#include "ns3/udp-server.h"
#include "ns3/uinteger.h"
#include "ns3/yans-wifi-channel.h"
#include "ns3/yans-wifi-helper.h"

using namespace ns3;

int
main(int argc, char* argv[])
{
    std::ofstream file("80211n-mimo-throughput.plt");

    std::vector<std::string> modes = {
        "HtMcs0",  "HtMcs1",  "HtMcs2",  "HtMcs3",  "HtMcs4",  "HtMcs5",  "HtMcs6",  "HtMcs7",
        "HtMcs8",  "HtMcs9",  "HtMcs10", "HtMcs11", "HtMcs12", "HtMcs13", "HtMcs14", "HtMcs15",
        "HtMcs16", "HtMcs17", "HtMcs18", "HtMcs19", "HtMcs20", "HtMcs21", "HtMcs22", "HtMcs23",
        "HtMcs24", "HtMcs25", "HtMcs26", "HtMcs27", "HtMcs28", "HtMcs29", "HtMcs30", "HtMcs31",
    };

    bool udp{true};
    Time simulationTime{"5s"};
    double frequency{5.0}; // whether 2.4 or 5.0 GHz
    double step{5};        // meters
    bool shortGuardInterval{false};
    bool channelBonding{false};
    bool preambleDetection{true};

    CommandLine cmd(__FILE__);
    cmd.AddValue("step", "Granularity of the results to be plotted in meters", step);
    cmd.AddValue("simulationTime", "Simulation time per step", simulationTime);
    cmd.AddValue("channelBonding",
                 "Enable/disable channel bonding (channel width = 20 MHz if false, channel width = "
                 "40 MHz if true)",
                 channelBonding);
    cmd.AddValue("preambleDetection", "Enable/disable preamble detection model", preambleDetection);
    cmd.AddValue("shortGuardInterval", "Enable/disable short guard interval", shortGuardInterval);
    cmd.AddValue("frequency",
                 "Whether working in the 2.4 or 5.0 GHz band (other values gets rejected)",
                 frequency);
    cmd.AddValue("udp", "UDP if set to 1, TCP otherwise", udp);
    cmd.Parse(argc, argv);

    Gnuplot plot = Gnuplot("80211n-mimo-throughput.eps");

    for (uint32_t i = 0; i < modes.size(); i++) // MCS
    {
        std::cout << modes[i] << std::endl;
        Gnuplot2dDataset dataset(modes[i]);
        for (double d = 0; d <= 100;) // distance
        {
            std::cout << "Distance = " << d << "m: " << std::endl;
            uint32_t payloadSize; // 1500 byte IP packet
            if (udp)
            {
                payloadSize = 1472; // bytes
            }
            else
            {
                payloadSize = 1448; // bytes
                Config::SetDefault("ns3::TcpSocket::SegmentSize", UintegerValue(payloadSize));
            }

            uint8_t nStreams = 1 + (i / 8); // number of MIMO streams

            NodeContainer wifiStaNode;
            wifiStaNode.Create(1);
            NodeContainer wifiApNode;
            wifiApNode.Create(1);

            YansWifiChannelHelper channel = YansWifiChannelHelper::Default();
            YansWifiPhyHelper phy;
            phy.SetChannel(channel.Create());
            if (!preambleDetection)
            {
                phy.DisablePreambleDetectionModel();
            }

            // Set MIMO capabilities
            phy.Set("Antennas", UintegerValue(nStreams));
            phy.Set("MaxSupportedTxSpatialStreams", UintegerValue(nStreams));
            phy.Set("MaxSupportedRxSpatialStreams", UintegerValue(nStreams));
            phy.Set("ChannelSettings",
                    StringValue(std::string("{0, ") + (channelBonding ? "40, " : "20, ") +
                                (frequency == 2.4 ? "BAND_2_4GHZ" : "BAND_5GHZ") + ", 0}"));

            WifiMacHelper mac;
            WifiHelper wifi;
            if (frequency == 5.0)
            {
                wifi.SetStandard(WIFI_STANDARD_80211n);
            }
            else if (frequency == 2.4)
            {
                wifi.SetStandard(WIFI_STANDARD_80211n);
                Config::SetDefault("ns3::LogDistancePropagationLossModel::ReferenceLoss",
                                   DoubleValue(40.046));
            }
            else
            {
                std::cout << "Wrong frequency value!" << std::endl;
                return 0;
            }

            StringValue ctrlRate;
            if (frequency == 2.4)
            {
                ctrlRate = StringValue("ErpOfdmRate24Mbps");
            }
            else
            {
                ctrlRate = StringValue("OfdmRate24Mbps");
            }
            wifi.SetRemoteStationManager("ns3::ConstantRateWifiManager",
                                         "DataMode",
                                         StringValue(modes[i]),
                                         "ControlMode",
                                         ctrlRate);

            Ssid ssid = Ssid("ns3-80211n");

            mac.SetType("ns3::StaWifiMac", "Ssid", SsidValue(ssid));

            NetDeviceContainer staDevice;
            staDevice = wifi.Install(phy, mac, wifiStaNode);

            mac.SetType("ns3::ApWifiMac", "Ssid", SsidValue(ssid));

            NetDeviceContainer apDevice;
            apDevice = wifi.Install(phy, mac, wifiApNode);

            // Set guard interval
            Config::Set("/NodeList/*/DeviceList/*/$ns3::WifiNetDevice/HtConfiguration/"
                        "ShortGuardIntervalSupported",
                        BooleanValue(shortGuardInterval));

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

            positionAlloc->Add(Vector(0.0, 0.0, 0.0));
            positionAlloc->Add(Vector(d, 0.0, 0.0));
            mobility.SetPositionAllocator(positionAlloc);

            mobility.SetMobilityModel("ns3::ConstantPositionMobilityModel");

            mobility.Install(wifiApNode);
            mobility.Install(wifiStaNode);

            /* Internet stack*/
            InternetStackHelper stack;
            stack.Install(wifiApNode);
            stack.Install(wifiStaNode);

            Ipv4AddressHelper address;
            address.SetBase("192.168.1.0", "255.255.255.0");
            Ipv4InterfaceContainer staNodeInterface;
            Ipv4InterfaceContainer apNodeInterface;

            staNodeInterface = address.Assign(staDevice);
            apNodeInterface = address.Assign(apDevice);

            /* Setting applications */
            const auto maxLoad = HtPhy::GetDataRate(i,
                                                    channelBonding ? MHz_u{40} : MHz_u{20},
                                                    NanoSeconds(shortGuardInterval ? 400 : 800),
                                                    nStreams);
            ApplicationContainer serverApp;
            if (udp)
            {
                // UDP flow
                uint16_t port = 9;
                UdpServerHelper server(port);
                serverApp = server.Install(wifiStaNode.Get(0));
                serverApp.Start(Seconds(0));
                serverApp.Stop(simulationTime + Seconds(1));
                const auto packetInterval = payloadSize * 8.0 / maxLoad;

                UdpClientHelper client(staNodeInterface.GetAddress(0), port);
                client.SetAttribute("MaxPackets", UintegerValue(4294967295U));
                client.SetAttribute("Interval", TimeValue(Seconds(packetInterval)));
                client.SetAttribute("PacketSize", UintegerValue(payloadSize));
                ApplicationContainer clientApp = client.Install(wifiApNode.Get(0));
                clientApp.Start(Seconds(1));
                clientApp.Stop(simulationTime + Seconds(1));
            }
            else
            {
                // TCP flow
                uint16_t port = 50000;
                Address localAddress(InetSocketAddress(Ipv4Address::GetAny(), port));
                PacketSinkHelper packetSinkHelper("ns3::TcpSocketFactory", localAddress);
                serverApp = packetSinkHelper.Install(wifiStaNode.Get(0));
                serverApp.Start(Seconds(0));
                serverApp.Stop(simulationTime + Seconds(1));

                OnOffHelper onoff("ns3::TcpSocketFactory", Ipv4Address::GetAny());
                onoff.SetAttribute("OnTime",
                                   StringValue("ns3::ConstantRandomVariable[Constant=1]"));
                onoff.SetAttribute("OffTime",
                                   StringValue("ns3::ConstantRandomVariable[Constant=0]"));
                onoff.SetAttribute("PacketSize", UintegerValue(payloadSize));
                onoff.SetAttribute("DataRate", DataRateValue(maxLoad));
                AddressValue remoteAddress(InetSocketAddress(staNodeInterface.GetAddress(0), port));
                onoff.SetAttribute("Remote", remoteAddress);
                ApplicationContainer clientApp = onoff.Install(wifiApNode.Get(0));
                clientApp.Start(Seconds(1));
                clientApp.Stop(simulationTime + Seconds(1));
            }

            Ipv4GlobalRoutingHelper::PopulateRoutingTables();

            Simulator::Stop(simulationTime + Seconds(1));
            Simulator::Run();

            double throughput = 0;
            if (udp)
            {
                // UDP
                double totalPacketsThrough =
                    DynamicCast<UdpServer>(serverApp.Get(0))->GetReceived();
                throughput = totalPacketsThrough * payloadSize * 8 /
                             simulationTime.GetMicroSeconds(); // Mbit/s
            }
            else
            {
                // TCP
                double totalPacketsThrough =
                    DynamicCast<PacketSink>(serverApp.Get(0))->GetTotalRx();
                throughput = totalPacketsThrough * 8 / simulationTime.GetMicroSeconds(); // Mbit/s
            }
            dataset.Add(d, throughput);
            std::cout << throughput << " Mbit/s" << std::endl;
            d += step;
            Simulator::Destroy();
        }
        plot.AddDataset(dataset);
    }

    plot.SetTerminal("postscript eps color enh \"Times-BoldItalic\"");
    plot.SetLegend("Distance (Meters)", "Throughput (Mbit/s)");
    plot.SetExtra("set xrange [0:100]\n\
set yrange [0:600]\n\
set ytics 0,50,600\n\
set style line 1 dashtype 1 linewidth 5\n\
set style line 2 dashtype 1 linewidth 5\n\
set style line 3 dashtype 1 linewidth 5\n\
set style line 4 dashtype 1 linewidth 5\n\
set style line 5 dashtype 1 linewidth 5\n\
set style line 6 dashtype 1 linewidth 5\n\
set style line 7 dashtype 1 linewidth 5\n\
set style line 8 dashtype 1 linewidth 5\n\
set style line 9 dashtype 2 linewidth 5\n\
set style line 10 dashtype 2 linewidth 5\n\
set style line 11 dashtype 2 linewidth 5\n\
set style line 12 dashtype 2 linewidth 5\n\
set style line 13 dashtype 2 linewidth 5\n\
set style line 14 dashtype 2 linewidth 5\n\
set style line 15 dashtype 2 linewidth 5\n\
set style line 16 dashtype 2 linewidth 5\n\
set style line 17 dashtype 3 linewidth 5\n\
set style line 18 dashtype 3 linewidth 5\n\
set style line 19 dashtype 3 linewidth 5\n\
set style line 20 dashtype 3 linewidth 5\n\
set style line 21 dashtype 3 linewidth 5\n\
set style line 22 dashtype 3 linewidth 5\n\
set style line 23 dashtype 3 linewidth 5\n\
set style line 24 dashtype 3 linewidth 5\n\
set style line 25 dashtype 4 linewidth 5\n\
set style line 26 dashtype 4 linewidth 5\n\
set style line 27 dashtype 4 linewidth 5\n\
set style line 28 dashtype 4 linewidth 5\n\
set style line 29 dashtype 4 linewidth 5\n\
set style line 30 dashtype 4 linewidth 5\n\
set style line 31 dashtype 4 linewidth 5\n\
set style line 32 dashtype 4 linewidth 5\n\
set style increment user");
    plot.GenerateOutput(file);
    file.close();

    return 0;
}
