#include "ns3/applications-module.h"
#include "ns3/core-module.h"
#include "ns3/dce-module.h"
#include "ns3/internet-module.h"
#include "ns3/mobility-module.h"
#include "ns3/point-to-point-module.h"
#include "ns3/wifi-module.h"

/*
 * This examples uses Linux Stack + ns3 Native Application.
 *
 * The traffic goes through a simple network composed of
 * Wireless Client (802.11n) + Access Point Router + Core Network Router + End Server.
 *
 */

NS_LOG_COMPONENT_DEFINE("dce-cradle-80211n-network");

using namespace ns3;
using namespace std;

Ptr<PacketSink> sink;                           /* TCP Receiver (Traffic Receiver Application) */
NodeContainer apWifiNode, staWifiNodes;         /* Nodes Container for the Wireless Network */
NetDeviceContainer apWifiNetDevice;             /* Network Devices for the Access Point in the simulation */
NetDeviceContainer staWifiNetDevices;           /* Network Devices for Wifi Nodes in the simulation */

NodeContainer coreNetworkNodes, endServerNodes, nodes;

/*** Simulation Attributes ***/
uint32_t aggregationType = 0;                   /* Level of packet aggregation. */
uint32_t msduAggregationLevel = 100;            /* Level of MSDU aggregation. */
uint32_t mpduAggregationLevel = 100;            /* Level of MPDY aggregation. */
string phyMode = "HtMcs7";       		/* Type of the Physical Layer. */
double dbiGain = 10;                            /* Antenna Gain. */
bool verbose = false;                           /* Print Wifi Device Logging Information. */
bool pcap_tracing = false;                      /* PCAP Tracing is enabled or not. */
double simulationTime = 10.1;                   /* Simulation time in seconds. */

uint64_t lastTotalRx = 0;

void Progress()
{
    Time now = Simulator::Now() - Seconds(1);
    double cur = (sink->GetTotalRx() - lastTotalRx) * (double)8/1e6;
    cout << now.GetSeconds() << '\t' << cur << endl;
    lastTotalRx = sink->GetTotalRx();
    Simulator::Schedule(Seconds(1), Progress);
}

/**
 * Log packet drops.
 */
static void MacTxpDrop(Ptr<const Packet> p)
{
    NS_LOG_UNCOND("Mac Drop at " << Simulator::Now().GetSeconds());
}

static void PhyTxpDrop(Ptr<const Packet> p)
{
    NS_LOG_UNCOND("Phy Drop at " << Simulator::Now().GetSeconds());
}

void PrintTcpFlags(std::string key, std::string value)
{
    cout << key << "=" << value;
}

int main(int argc, char *argv[])
{
    string applicationType = "bulk";                    /* Type of the Tx application */
    uint32_t maxPackets = 0;                            /* Maximum Number of Packets */
    string dataRate = "1Gbps";                          /* Application Layer Data Rate. */
    uint32_t payloadSize = 1400;                        /* Transport Layer Payload size in bytes. */
    string socketType = "ns3::LinuxTcpSocketFactory";   /* Socket Type (TCP/UDP/SCTP/DCCP) */
    string tcpVariant = "cubic";                        /* TCP Variant Type. */
    string bufferSize = "";                             /* TCP Buffer Size. */
    uint32_t queueSize = 1000;                          /* Wifi Mac Queue Size. */
    uint32_t nWifiSta = 1;                              /* Number of Wifi STAs. */
    string distanceToRx = "1.0";                        /* The distance between transmitter and receiver in meters. */

    /* Command line argument parser setup. */
    CommandLine cmd;

    cmd.AddValue("applicationType", "Type of the Tx Application: onoff, bulk", applicationType);
    cmd.AddValue("maxPackets", "Maximum number of packets to send", maxPackets);
    cmd.AddValue("payloadSize", "Payload size in bytes", payloadSize);
    cmd.AddValue("dataRate", "Payload size in bytes", dataRate);
    cmd.AddValue("socketType", "Type of the Socket (ns3::LinuxTcpSocketFactory, ns3::LinuxUdpSocketFactory, ns3::LinuxDccpSocketFactory, ns3::LinuxSctpSocketFactory)", socketType);
    cmd.AddValue("tcpVariant", "Transport protocol to use: reno, bic, cubic, westwood, highspeed, hybla, htcp, vegas, scalable, olia, lia, wvegas", tcpVariant);
    cmd.AddValue("bufferSize", "Send/Receiver buffer size", bufferSize);
    cmd.AddValue("aggregationType", "Level of Aggregation: 0=No, 1=A-MSDU, 2=A-MPDU, 3=TwoLevel", aggregationType);
    cmd.AddValue("msduAggregationLevel", "Level of MSDU aggregation", msduAggregationLevel);
    cmd.AddValue("mpduAggregationLevel", "Level of MPDY aggregation", mpduAggregationLevel);
    cmd.AddValue("queueSize", "The size of the Wifi Mac Queue", queueSize);
    cmd.AddValue("phyMode", "802.11ad PHY Mode", phyMode);
    cmd.AddValue("dbiGain", "Gain of the antennas", dbiGain);
    cmd.AddValue("verbose", "turn on all WifiNetDevice log components", verbose);
    cmd.AddValue("dist", "distance between nodes", distanceToRx);
    cmd.AddValue("simulationTime", "Simulation time in Seconds", simulationTime);
    cmd.AddValue("pcap", "Enable PCAP Tracing", pcap_tracing);
    cmd.Parse(argc, argv);

    /* Global params: no fragmentation, no RTS/CTS, fixed rate for all packets */
    Config::SetDefault("ns3::WifiRemoteStationManager::FragmentationThreshold", StringValue("999999"));
    Config::SetDefault("ns3::WifiRemoteStationManager::RtsCtsThreshold", StringValue("999999"));
    Config::SetDefault("ns3::WifiMacQueue::MaxPacketNumber", UintegerValue(queueSize));

    /* Instantiate the Output. */
    cout << "Simulator Time" << '\t' << "Application Goodput (Mbps)" << endl;

    /* Instantiate DCE Manager to use Liux Stack */
    DceManagerHelper dceManager;
    dceManager.SetTaskManagerAttribute("FiberManagerType", StringValue("UcontextFiberManager"));
    dceManager.SetNetworkStack("ns3::LinuxSocketFdFactory", "Library", StringValue("liblinux.so"));

    /* Instantiate DCE Application Helper to use DCE Applications */
    DceApplicationHelper dce;
    dce.SetStackSize(1 << 16);

    /*** Create Network Topology ***/
    PointToPointHelper p2pHelper;

    /* Create Core Network */
    NodeContainer coreNetworkNodes;
    coreNetworkNodes.Create(2);

    NetDeviceContainer coreNetworkDevices;
    p2pHelper.SetDeviceAttribute("DataRate", StringValue("1Gbps"));
    p2pHelper.SetChannelAttribute("Delay", TimeValue(MilliSeconds(20)));
    p2pHelper.SetQueue("ns3::DropTailQueue", "MaxPackets", UintegerValue(5000));
    coreNetworkDevices = p2pHelper.Install(coreNetworkNodes);

    /* Create End Server Node */
    NodeContainer endServerNodes;
    endServerNodes.Create(1);

    NetDeviceContainer endServerDevices;
    p2pHelper.SetDeviceAttribute("DataRate", StringValue("1Gbps"));
    p2pHelper.SetChannelAttribute("Delay", TimeValue(MilliSeconds(20)));
    p2pHelper.SetQueue("ns3::DropTailQueue", "MaxPackets", UintegerValue(5000));
    endServerDevices = p2pHelper.Install(coreNetworkNodes.Get(1), endServerNodes.Get(0));

    /* Create Wifi AP */
    WifiHelper wifi;

    /* Basic setup */
    wifi.SetStandard(WIFI_PHY_STANDARD_80211n_5GHZ);

    /* Turn on logging */
    if (verbose)
    {
        wifi.EnableLogComponents();
        LogComponentEnable("dce-cradle-80211n-network", LOG_LEVEL_ALL);
    }

    /**** Set up Channel ****/
    YansWifiChannelHelper wifiChannel ;
    /* Simple propagation delay model */
    wifiChannel.SetPropagationDelay("ns3::ConstantSpeedPropagationDelayModel");
    /* Friis model with standard-specific wavelength */
    wifiChannel.AddPropagationLoss("ns3::FriisPropagationLossModel", "Frequency", DoubleValue(5e9));

    /**** Allocate a default Adhoc Wifi MAC ****/
    /* Add a QoS upper mac */
    HtWifiMacHelper wifiMac = HtWifiMacHelper::Default ();

    /* Set Wifi MAC Type. */
    wifiMac.SetType("ns3::AdhocWifiMac");

    /* Enable Maximum Aggregation */
    if ((aggregationType == 1) || (aggregationType == 3))
        wifiMac.SetMsduAggregatorForAc(AC_BE, "ns3::MsduStandardAggregator", "MaxAmsduSize", UintegerValue(7935 * msduAggregationLevel));
    if ((aggregationType == 2) || (aggregationType == 3))
    {
        wifiMac.SetBlockAckThresholdForAc(AC_BE, 64);
        wifiMac.SetMpduAggregatorForAc(AC_BE, "ns3::MpduStandardAggregator", "MaxAmpduSize", UintegerValue(65535 * mpduAggregationLevel));
    }

    /**** SETUP ALL NODES ****/
    YansWifiPhyHelper wifiPhy = YansWifiPhyHelper::Default();
    /* Nodes will be added to the channel we set up earlier */
    wifiPhy.SetChannel(wifiChannel.Create());
    /* All nodes transmit at 10 dBm == 10 mW, no adaptation */
//    wifiPhy.Set("TxPowerStart", DoubleValue(10.0));
//    wifiPhy.Set("TxPowerEnd", DoubleValue(10.0));
//    wifiPhy.Set("TxPowerLevels", UintegerValue(1));
//    wifiPhy.Set("TxGain", DoubleValue(0));
//    wifiPhy.Set("RxGain", DoubleValue(0));
//    /* Sensitivity model includes implementation loss and noise figure */
//    wifiPhy.Set("RxNoiseFigure", DoubleValue(3));
//    /* Set CCA threshold for all nodes to 0dBm=1mW; i.e. don't carrier sense */
//    wifiPhy.Set("CcaMode1Threshold", DoubleValue(-79));
//    wifiPhy.Set("EnergyDetectionThreshold", DoubleValue(-79 + 3));
    /* Set the phy layer error model */
    wifiPhy.SetErrorRateModel("ns3::NistErrorRateModel");
    wifiPhy.Set("ShortGuardEnabled", BooleanValue(true));
    wifiPhy.Set("ChannelWidth", UintegerValue(40));
    wifi.SetRemoteStationManager("ns3::ConstantRateWifiManager", "DataMode", StringValue(phyMode),
                                                                 "ControlMode", StringValue(phyMode));

    /* Create 802.11n Wifi Nodes */
    apWifiNode.Add(coreNetworkNodes.Get(0));
    staWifiNodes.Create(nWifiSta);

    /* Make the wifi nodes and set them up with the phy and the mac */
    NodeContainer wifiNodes;
    wifiNodes.Add(apWifiNode);
    wifiNodes.Add(staWifiNodes);

    /* Accummulate end nodes in one node container */
    nodes.Add(staWifiNodes);
    nodes.Add(coreNetworkNodes);
    nodes.Add(endServerNodes);

    /* Install 802.11n Network Access Point */
    apWifiNetDevice = wifi.Install(wifiPhy, wifiMac, apWifiNode);

    /* Install 802.11n Network Wireless Stations */
    staWifiNetDevices = wifi.Install(wifiPhy, wifiMac, staWifiNodes);

    /* Aggregate all 802.11n devices in one container */
    NetDeviceContainer wifiDevices;
    wifiDevices.Add(apWifiNetDevice);
    wifiDevices.Add(staWifiNetDevices);

    /**** Set up node positions ****/
    MobilityHelper mobility;

    /* Allocate constant position for Access Point */
    Ptr<ListPositionAllocator> positionAlloc = CreateObject<ListPositionAllocator>();
    positionAlloc->Add(Vector(0, 0, 0));
    mobility.SetPositionAllocator(positionAlloc);
    mobility.SetMobilityModel("ns3::ConstantPositionMobilityModel");
    mobility.Install(apWifiNode);

    /* Allocate Mobility model for Wireless Stations */
    mobility.SetPositionAllocator("ns3::RandomBoxPositionAllocator",
                                  "X", StringValue("ns3::ConstantRandomVariable[Constant=" + distanceToRx + "]"),
                                  "Y", StringValue("ns3::ConstantRandomVariable[Constant=0.0]"),
                                  "Z", StringValue("ns3::ConstantRandomVariable[Constant=0.0]"));

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

    /* Give all nodes a Linux Stack */
    LinuxStackHelper stack;
    stack.Install(nodes);
    dceManager.Install(nodes);

    /* And interfaces IP addresses */
    Ipv4AddressHelper ipv4;
    Ipv4InterfaceContainer wifiDevicesInterfaces;

    ipv4.SetBase("10.1.1.0", "255.255.255.0");
    wifiDevicesInterfaces = ipv4.Assign(wifiDevices);

    ipv4.SetBase("10.1.2.0", "255.255.255.0");
    ipv4.Assign(coreNetworkDevices);

    ipv4.SetBase("10.1.3.0", "255.255.255.0");
    ipv4.Assign(endServerDevices);

    /* Setup IP routes */
    Ipv4GlobalRoutingHelper::PopulateRoutingTables();
    LinuxStackHelper::PopulateRoutingTables();

    /* Instantiate the Output. */
    stack.SysctlGet(staWifiNodes.Get(0), Seconds(0), ".net.ipv4.tcp_available_congestion_control", &PrintTcpFlags);
    stack.SysctlGet(staWifiNodes.Get(0), Seconds(0), ".net.ipv4.tcp_congestion_control", &PrintTcpFlags);
    stack.SysctlGet(staWifiNodes.Get(0), Seconds(0), ".net.ipv4.tcp_sack", &PrintTcpFlags);
    stack.SysctlGet(staWifiNodes.Get(0), Seconds(0), ".net.ipv4.tcp_timestamps", &PrintTcpFlags);
    stack.SysctlGet(staWifiNodes.Get(0), Seconds(0), ".net.ipv4.tcp_window_scaling", &PrintTcpFlags);
    stack.SysctlGet(staWifiNodes.Get(0), Seconds(0), ".net.ipv4.tcp_mtu_probing", &PrintTcpFlags);
    stack.SysctlGet(staWifiNodes.Get(0), Seconds(0), ".net.ipv4.tcp_rmem", &PrintTcpFlags);
    stack.SysctlGet(staWifiNodes.Get(0), Seconds(0), ".net.ipv4.tcp_wmem", &PrintTcpFlags);
    stack.SysctlGet(staWifiNodes.Get(0), Seconds(0), ".net.core.rmem_max", &PrintTcpFlags);
    stack.SysctlGet(staWifiNodes.Get(0), Seconds(0), ".net.core.wmem_max", &PrintTcpFlags);
    stack.SysctlGet(staWifiNodes.Get(0), Seconds(0), ".net.core.rmem_default", &PrintTcpFlags);
    stack.SysctlGet(staWifiNodes.Get(0), Seconds(0), ".net.core.wmem_default", &PrintTcpFlags);
    stack.SysctlGet(staWifiNodes.Get(0), Seconds(0), ".net.core.netdev_max_backlog", &PrintTcpFlags);

    if (bufferSize.length() != 0)
    {
        stack.SysctlSet(nodes, ".net.ipv4.tcp_congestion_control", tcpVariant);
        /* Allow testing with buffers up */
        stack.SysctlSet(nodes, ".net.core.rmem_max", bufferSize);
        stack.SysctlSet(nodes, ".net.core.wmem_max", bufferSize);
        stack.SysctlSet(nodes, ".net.core.wmem_default", bufferSize);
        stack.SysctlSet(nodes, ".net.core.rmem_default", bufferSize);

        /* Increase Linux autotuning TCP buffer limit to */
        stack.SysctlSet(nodes, ".net.ipv4.tcp_rmem", bufferSize + " " + bufferSize + " " + bufferSize);
        stack.SysctlSet(nodes, ".net.ipv4.tcp_wmem", bufferSize + " " + bufferSize + " " + bufferSize);
        /* Increase the length of the processor input queue */
        stack.SysctlSet(nodes, ".net.core.netdev_max_backlog", "30000");

        stack.SysctlGet(staWifiNodes.Get(0), Seconds(1), ".net.ipv4.tcp_congestion_control", &PrintTcpFlags);
        stack.SysctlGet(staWifiNodes.Get(0), Seconds(1), ".net.ipv4.tcp_rmem", &PrintTcpFlags);
        stack.SysctlGet(staWifiNodes.Get(0), Seconds(1), ".net.ipv4.tcp_wmem", &PrintTcpFlags);
        stack.SysctlGet(staWifiNodes.Get(0), Seconds(1), ".net.core.rmem_max", &PrintTcpFlags);
        stack.SysctlGet(staWifiNodes.Get(0), Seconds(1), ".net.core.wmem_max", &PrintTcpFlags);
        stack.SysctlGet(staWifiNodes.Get(0), Seconds(1), ".net.core.wmem_default", &PrintTcpFlags);
        stack.SysctlGet(staWifiNodes.Get(0), Seconds(1), ".net.core.rmem_default", &PrintTcpFlags);
        stack.SysctlGet(staWifiNodes.Get(0), Seconds(1), ".net.core.netdev_max_backlog", &PrintTcpFlags);
    }

    /* Receiver (Installed on the Client) */
    PacketSinkHelper sinkHelper(socketType, InetSocketAddress(Ipv4Address::GetAny(), 9999));
    ApplicationContainer sinkApp = sinkHelper.Install(staWifiNodes.Get(0));
    sink = DynamicCast<PacketSink>(sinkApp.Get(0));
    sinkApp.Start(Seconds(4.0));

    if (applicationType == "onoff")
    {
        /* The total number of bytes to send. Once these bytes are sent, no packet is sent again, even in on state. The value zero means that there is no limit. */
        Config::SetDefault("ns3::OnOffApplication::MaxBytes", UintegerValue(payloadSize * maxPackets));
        Config::SetDefault("ns3::OnOffApplication::PacketSize", UintegerValue(payloadSize));
        Config::SetDefault("ns3::OnOffApplication::OnTime", StringValue("ns3::ConstantRandomVariable[Constant=1e6]"));
        Config::SetDefault("ns3::OnOffApplication::OffTime", StringValue("ns3::ConstantRandomVariable[Constant=0]"));
        Config::SetDefault("ns3::OnOffApplication::DataRate", DataRateValue(DataRate(dataRate)));

        /* Sender (Installed on the End Serrver) */
        OnOffHelper src(socketType, InetSocketAddress("10.1.1.2", 9999));
        ApplicationContainer srcApp = src.Install(endServerNodes.Get(0));
        srcApp.Start(Seconds(5.0));
    }
    else if (applicationType == "bulk")
    {
        BulkSendHelper bulk = BulkSendHelper("ns3::LinuxTcpSocketFactory", InetSocketAddress("10.1.1.2", 9999));
        bulk.SetAttribute("MaxBytes", UintegerValue(0));
        bulk.SetAttribute("SendSize", UintegerValue(payloadSize));
        ApplicationContainer srcApp = bulk.Install(endServerNodes.Get(0));
        srcApp.Start(Seconds(5.0));
    }

    /* Trace Congestion Window Changes for all Sockets. */
    Config::ConnectWithoutContext("/NodeList/*/DeviceList/*/$ns3::PointToPointNetDevice/MacTxDrop", MakeCallback(&MacTxpDrop));
    Config::ConnectWithoutContext("/NodeList/*/DeviceList/*/$ns3::PointToPointNetDevice/PhyTxDrop", MakeCallback(&PhyTxpDrop));

    if (pcap_tracing)
    {
        wifiPhy.SetPcapDataLinkType(YansWifiPhyHelper::DLT_IEEE802_11_RADIO);
        wifiPhy.EnablePcap("Traces/WirelessStation", staWifiNodes);
        wifiPhy.EnablePcap("Traces/AccessPoint", wifiDevices.Get(1));
        p2pHelper.EnablePcap("Traces/AccessPoint", apWifiNode, false);
        p2pHelper.EnablePcap("Traces/Router", coreNetworkNodes, false);
        p2pHelper.EnablePcap("Traces/EndServer", endServerNodes, false);
    }

    /* Run Simulation */
    Simulator::Schedule(Seconds(6.0), Progress);
    Simulator::Stop(Seconds(simulationTime));
    Simulator::Run();

    /* Print Results */
    Time now = Simulator::Now() - Seconds(2);
    cout << "Statistics at the Receiver:"  << endl;
    cout << "Average goodput seen by the application = " << (sink->GetTotalRx() * (double)8/1e6)/now.GetSeconds() << endl;
    cout << "End Simulation at " << now.GetSeconds() << endl;

    Simulator::Destroy();

    return 0;
}