#include <iostream>
#include <fstream>
#include <sstream>

#include <sys/types.h>
#include <sys/stat.h>
#include <dirent.h>

#include "ns3/log.h"

#include "ns3/data-rate.h"
#include "ns3/core-module.h"
#include "ns3/mobility-module.h"

#include "ns3/internet-module.h"
#include "ns3/internet-stack-helper.h"
#include "ns3/ipv4-address-helper.h"
#include "ns3/ipv4-interface-container.h"
#include "ns3/ipv4-flow-classifier.h"

#include "ns3/wave-net-device.h"
#include "ns3/wave-mac-helper.h"
#include "ns3/wave-helper.h"

#include "ns3/applications-module.h"
#include "ns3/flow-monitor-helper.h"

#define NODE_COUNT 2
#define DEFAULT_SIM_TIME 10.0

using namespace ns3;

#define MODULE_NAME "WaveTestSuite"

//!< Specifies max data rate with which UDP stream will be send
//!< Should be higher or equal WiFi data rate (m_dataRate)
#define MAX_UDP_DATA_RATE "30000kb/s"

NS_LOG_COMPONENT_DEFINE (MODULE_NAME);

class WaveTestSuite
{
public:
  WaveTestSuite(std::string dataRate);

  void Run ();

private:
  /// >>
  void CreateNodes ();
  void SetUpMobilityModel ();
  void CreateDevices ();
  void InstallInternetStack ();
  void InstallApplications (double simTime);
  void Report ();
  void PrepareAnimationInterface ();
  void PrepareReportGeneration ();
  /// <<

  bool PrepareFolderForResults ();
  void FinalizeReport ();
  bool CheckParams ();

  int32_t m_distance;
  std::string m_dataRate;

  std::ostringstream m_prefix;
  std::ostringstream m_folderName;
  std::ofstream m_reportFile;
  NodeContainer m_nodes;
  NetDeviceContainer m_devices;
  Ipv4InterfaceContainer m_ipv4Interfaces;

  FlowMonitorHelper m_flowMonHelper;
  Ptr<FlowMonitor> m_monitor;

  double m_rxBt;
  double m_lastT;

  double m_reportTime;
};

WaveTestSuite::WaveTestSuite (std::string dataRate):
  m_distance (200),
  m_dataRate (dataRate),
  m_prefix(MODULE_NAME),
  m_folderName("./"),
  m_rxBt (0),
  m_lastT (0),
  m_reportTime (1.)
{
}

void
WaveTestSuite::CreateNodes ()
{
  m_nodes = NodeContainer();
  m_nodes.Create (NODE_COUNT);
}

void
WaveTestSuite::SetUpMobilityModel ()
{
  MobilityHelper mobility;

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

  positionAlloc->Add (Vector (0, 0, 0.0));
  positionAlloc->Add (Vector (0, m_distance, 0.0));
  mobility.SetPositionAllocator (positionAlloc);
  mobility.SetMobilityModel ("ns3::ConstantPositionMobilityModel");
  mobility.Install (m_nodes);
}

void
WaveTestSuite::CreateDevices ()
{
  YansWifiChannelHelper waveChannel = YansWifiChannelHelper::Default ();

  YansWavePhyHelper wavePhy =  YansWavePhyHelper::Default ();
  wavePhy.SetChannel (waveChannel.Create ());
  wavePhy.SetPcapDataLinkType (YansWifiPhyHelper::DLT_IEEE802_11);

  /// 44.8 dBm as defined in 802.11p for US
  wavePhy.Set ("TxPowerStart", DoubleValue (44.8) );
  wavePhy.Set ("TxPowerEnd", DoubleValue (44.8) );

  QosWaveMacHelper waveMac = QosWaveMacHelper::Default ();
  WaveHelper waveHelper = WaveHelper::Default ();
  waveHelper.SetRemoteStationManager ("ns3::ConstantRateWifiManager",
                                      "DataMode", StringValue (m_dataRate),
                                      "ControlMode",StringValue (m_dataRate),
                                      "NonUnicastMode", StringValue (m_dataRate));

  m_devices = waveHelper.Install (wavePhy, waveMac, m_nodes);

  const SchInfo schInfo = SchInfo (SCH3, true, EXTENDED_CONTINUOUS);
  const TxProfile txProfile = TxProfile (SCH3);
// >>>
// Replace the previous line with the next one to check that user is able to specify Data Rate
//  const TxProfile txProfile = TxProfile (SCH3, true, 4, m_dataRate);
// <<<
  for (NetDeviceContainer::Iterator i = m_devices.Begin (); i != m_devices.End (); ++i)
    {
      Ptr<NetDevice> device = *i;

      WaveNetDevice * waveNetDevice = static_cast<WaveNetDevice*>(GetPointer (device));

      Simulator::Schedule(Seconds (0.0), &WaveNetDevice::StartSch, waveNetDevice, schInfo);
      Simulator::Schedule(Seconds (0.0), &WaveNetDevice::RegisterTxProfile, waveNetDevice, txProfile);
    }

//  wavePhy.EnablePcapAll ("waveTxProfileTest.pcap");
}

void
WaveTestSuite::InstallInternetStack ()
{
  InternetStackHelper internet;
  internet.Install (m_nodes);

  NS_LOG_INFO ("Assign IP Addresses.");

  Ipv4AddressHelper ipv4;
  ipv4.SetBase ("10.1.1.0", "255.255.255.0");

  m_ipv4Interfaces = ipv4.Assign (m_devices);
}

void
WaveTestSuite::InstallApplications (double simTime)
{
//  BulkSendHelper sourceHelper ("ns3::UdpSocketFactory", InetSocketAddress (m_ipv4Interfaces.GetAddress (m_nodes.GetN () - 1), 50001));
//  sourceHelper.SetAttribute ("SendSize", UintegerValue (1024));
//  sourceHelper.SetAttribute ("MaxBytes", UintegerValue (0));
//  sourceHelper.SetAttribute ("DataRate", DataRateValue (DataRate (MAX_UDP_DATA_RATE)));
//  ApplicationContainer sourceApp = sourceHelper.Install (m_nodes.Get (0));

  OnOffHelper onOffHelper ("ns3::UdpSocketFactory", InetSocketAddress (m_ipv4Interfaces.GetAddress (m_nodes.GetN () - 1), 50001));
  onOffHelper.SetAttribute ("DataRate", DataRateValue (DataRate (MAX_UDP_DATA_RATE)));
  onOffHelper.SetAttribute ("PacketSize", UintegerValue (1024));
  onOffHelper.SetAttribute ("MaxBytes", UintegerValue (0));
  onOffHelper.SetAttribute ("OffTime", StringValue("ns3::ConstantRandomVariable[Constant=0.0]"));
  onOffHelper.SetAttribute ("OnTime", StringValue("ns3::ConstantRandomVariable[Constant=10.0]"));

  ApplicationContainer sourceApp = onOffHelper.Install (m_nodes.Get (0));
  PacketSinkHelper dlPacketSinkHelper ("ns3::UdpSocketFactory", InetSocketAddress (Ipv4Address::GetAny (), 50001));
  sourceApp.Add(dlPacketSinkHelper.Install (m_nodes.Get (m_nodes.GetN () - 1)));
  sourceApp.Start (Seconds (0.));
  sourceApp.Stop (Seconds (simTime));
}

void
WaveTestSuite::Report ()
{
  m_monitor->CheckForLostPackets ();

  Ptr<Ipv4FlowClassifier> classifier = DynamicCast<Ipv4FlowClassifier> (m_flowMonHelper.GetClassifier ());
  m_monitor->SerializeToXmlFile(m_folderName.str () + "flowmon_" + m_prefix.str () + ".xml", true, true);

  std::map<FlowId, FlowMonitor::FlowStats> stats = m_monitor->GetFlowStats ();

  std::ostringstream reportStream;

  for (std::map<FlowId, FlowMonitor::FlowStats>::const_iterator iter = stats.begin ();iter != stats.end ();iter++)
    {
      Ipv4FlowClassifier::FiveTuple t = classifier->FindFlow (iter->first);

      if (iter->first == 1)
        {
          if(iter->second.rxPackets==0)
            {
              reportStream << "Time: " << Simulator::Now () << " Src Addr " << t.sourceAddress << " Throughput: 0 Kbps\n";
            }
          else
            {
              float th = iter->second.rxBytes * 8.0 / (iter->second.timeLastRxPacket.GetSeconds() - iter->second.timeFirstRxPacket.GetSeconds()) / 1024;
              float th2;
              if (iter->second.timeLastRxPacket.GetSeconds() - m_lastT > 0)
                {
                  th2 = (iter->second.rxBytes - m_rxBt) * 8.0 / (iter->second.timeLastRxPacket.GetSeconds() - m_lastT) / 1024;
                }
              else
                {
                  // No packets were received since last report
                  th2 = 0;
                }

              reportStream << "Time: " << Simulator::Now () << " Src Addr " << t.sourceAddress << " AveThroughput: " << th << " Kbps";
              reportStream << " Throughput: " << th2 << " Kbps";
            }

          m_rxBt = iter->second.rxBytes;
          m_lastT = iter->second.timeLastRxPacket.GetSeconds();
        }
    }

  m_reportFile << reportStream.str () << "\n";
  NS_LOG_INFO (reportStream.str ());

  Simulator::Schedule(Seconds(m_reportTime), &WaveTestSuite::Report, this);
}

void
WaveTestSuite::PrepareReportGeneration ()
{
  m_monitor = m_flowMonHelper.Install(m_nodes);
  m_reportFile.open (m_folderName.str () + "report_" + m_prefix.str () + ".txt");

  Simulator::Schedule (Seconds (m_reportTime), &WaveTestSuite::Report, this);
}

void
WaveTestSuite::FinalizeReport ()
{
  m_reportFile.close ();
}

void
WaveTestSuite::Run ()
{
  CreateNodes ();
  SetUpMobilityModel ();
  CreateDevices ();
  InstallInternetStack ();
  InstallApplications (DEFAULT_SIM_TIME);

  PrepareReportGeneration ();

  Simulator::Stop (Seconds (DEFAULT_SIM_TIME));
  Simulator::Run ();
  Simulator::Destroy ();

  FinalizeReport ();
}

int main (int argc, char *argv[])
{
  WaveTestSuite testSuite ("OfdmRate3MbpsBW10MHz");

  LogComponentEnable(MODULE_NAME, LOG_LEVEL_ALL);
  LogComponentEnable(MODULE_NAME, LOG_PREFIX_ALL);

  testSuite.Run ();
}