#include "ns3/core-module.h"
#include "ns3/common-module.h"
#include "ns3/node-module.h"
#include "ns3/wifi-module.h"
#include "ns3/mobility-module.h"
#include "ns3/simulator-module.h"
 
#include "ns3/helper-module.h"
#include "ns3/contrib-module.h"
 
#include "ns3/ipv4-global-routing-helper.h"
#include "ns3/onoff-application.h"
 
#include "ns3/vector.h"
 
#include "ns3/visualizer.h"
 
using namespace ns3;
 
NS_LOG_COMPONENT_DEFINE("WifiPairsThroughputSimulation");
 
struct SimulationConfig
{
  unsigned int PACKET_SIZE;
  int NUM_PAIRS; // number of pairs in an experiment
  int NUM_STEPS; // number of throughputs to test
  std::string START; // first throughput tested
  std::string STEP; // throughput shift between 2 tests
  SimulationConfig(const unsigned int ps,
		   const int np,
		   const int ns,
		   const int start,
		   const int step,
		   const std::string unit) : 
    PACKET_SIZE(ps), NUM_PAIRS(np), NUM_STEPS(ns)
  {
    std::ostringstream oss1, oss2;
    oss1<<start<<unit;oss2<<step<<unit;
    START=oss1.str();STEP=oss2.str();
  }
};
// packet size, num pairs, num steps, start, step, unit
const SimulationConfig CONFIGS[] = {
  SimulationConfig(1000,7,26,62,62,"Kbps"),
  SimulationConfig(1000,7,26,62,62,"Kbps"),
//  SimulationConfig(1000,7,30,62,62,"Kbps"),
//  SimulationConfig(1000,7,30,62,62,"Kbps"),
//  SimulationConfig(1000,8,24,250,250,"Kbps"),
//  SimulationConfig(1000,8,24,250,250,"Kbps"),
}; const unsigned int NUM_CONFIGS = sizeof(CONFIGS)/sizeof(SimulationConfig);
const std::string OUTPUT_IMAGE_NAME = "wifi_udp_throughput.png";
const double DX_m = 3.0; // space between a AP and a MS
const double DY_m = 3.0; // space between 2 pairs
const unsigned int NUM_WIFI_CHANNELS = 1; // all the pairs are put in # chans
const int SAMPLING_PERIOD_s = 20;
const unsigned int NUM_RUNS_BY_CONFIG = 1;
const WifiPhyStandard USED_WIFI_STANDARD = WIFI_PHY_STANDARD_80211g;
 
//output parameters
const bool OUTPUT_PCAP = true;
const bool IFTRUEOUTPUTAP_PCAP = true;
 
template <typename T> class Point2d
{
public:
  T x,y;
  Point2d(T _x, T _y) : x(_x), y(_y) {};
};
 
class Pair : public Object 
{ // An AP/MS pair that will store the "#bits transmitted" data
public:
  Ptr<Node> client;
  Ptr<Node> server;
  uint64_t bits;
  std::vector<Point2d<double> > output;
  Pair() {};
  void ReceivePacket(Ptr<Socket> socket);
};
 
void Pair::ReceivePacket(Ptr<Socket> socket)
{ // simply refresh the #bits transmitted data
  Ptr<Packet> packet;
 
  while(packet = socket->Recv())
    {
      bits += packet->GetSize()*8;
    }
}
 
class ChannelChooser
{ 
  /* naive divide a segment and put numbers algo
     we actually consider the channels as coords on a 2d line
     then we create 2d segments on that line "unused channels"
     and always break the biggest one in 2, putting a (increasing)
     number on the middle point
     then we have all the points of our segment ordered
     we just loop on it to linearily attribute channels =) */
private: 
  class Segment1d : public Point2d<unsigned int> {
  public:
    Segment1d(unsigned int a, unsigned int b) : Point2d<unsigned int>(a,b) {}
    int size() { return y - x + 1; }
  };
 
  std::vector<Segment1d> segs_to_divide;
  std::vector<unsigned int> channel_attrib_vector; // used a vector because cba managing mem by myself here
  unsigned int num_chans;
 
  void divide(std::vector<Segment1d>::iterator it);
  std::vector<Segment1d>::iterator look_for_biggest_segment();
public:
  ChannelChooser(unsigned int s); // choose between s channels
  unsigned int get(unsigned int i); // my i-th channel will be ...
};
 
ChannelChooser::ChannelChooser(unsigned int s) : num_chans(s) {
  segs_to_divide.push_back(Segment1d(0,s-1));
  divide(segs_to_divide.begin());
}
 
void ChannelChooser::divide(std::vector<Segment1d>::iterator it)
{
  switch(it->size())
    {
    case 0:
      segs_to_divide.erase(it);
      break;
    case 1:
      segs_to_divide.erase(it);
      channel_attrib_vector.push_back(it->x);
      break;
    default:
      {
	unsigned int a = it->x;
	unsigned int b = it->y;
	unsigned int pivot = (a+b) / 2;
	// first we tag the pivot
	segs_to_divide.erase(it);
	channel_attrib_vector.push_back(pivot);
	// btw be it pair or impair the right size is systematically bigger
	segs_to_divide.push_back(Segment1d(pivot+1,b));
	if (pivot != a) // if size == 2 then pivot == a...
	  segs_to_divide.push_back(Segment1d(a,pivot-1)); 
      }
    }
  if(segs_to_divide.size() > 0)
    divide(look_for_biggest_segment()); // dat tail call o//
}
 
std::vector<ChannelChooser::Segment1d>::iterator ChannelChooser::look_for_biggest_segment()
{
  // ugly bubble sort, we could do a lot better but not worth the effort
  int max_size = 0;
  std::vector<Segment1d>::iterator res;
  for(std::vector<Segment1d>::iterator it = segs_to_divide.begin(); it < segs_to_divide.end(); it++)
    {
      if(it->size() < max_size)
	continue;
      max_size = it->size();
      res = it;
    }
  return res;
}
 
unsigned int ChannelChooser::get(unsigned int i)
{
  return channel_attrib_vector[i % num_chans];
}
 
class ExperimentWifiPairs
{
  /* The class that will actually run an experiment given appropriate params
     It is quite stateless, but not completely :p
     It is quite easy to make it completely stateless
     Just diffuse *step* where it needs to be
   */
private:
  void AdvanceRate(const Ptr<Pair> p);
  Ptr<Socket> SetupPacketReceive(Ptr<Pair> p, const Address &);
  uint32_t step; // bps(throughput) shift between 2 sampling in an experiment
public:
  ExperimentWifiPairs() {};
  std::vector <Ptr<Pair> > Run(const std::string sstart,
			       const std::string sstep,
			       const int numstep,
			       const int numpairs,
			       const unsigned int packet_size,
			       const YansWifiChannelHelper &channel);
};
 
Ptr<Socket> ExperimentWifiPairs::SetupPacketReceive(Ptr<Pair> p, const Address & address)
{
  TypeId tid = TypeId::LookupByName("ns3::UdpSocketFactory");
  Ptr<Socket> sink = Socket::CreateSocket(p->server, tid);
  sink->Bind(address);
  sink->SetRecvCallback(MakeCallback(&Pair::ReceivePacket, p));
  return sink;
}
 
void ExperimentWifiPairs::AdvanceRate(const Ptr<Pair> p)
{
  DataRateValue drv = DataRateValue(DataRate(0));
  double measured_mbps = p->bits/1000000.0/SAMPLING_PERIOD_s;
  p->bits = 0;
 
  Ptr<Application> onoff = p->client->GetApplication(0);
  onoff->GetAttribute("DataRate", drv);
 
  double  theoric_mbps;
  theoric_mbps = drv.Get().GetBitRate()/ 1000000.0;
  onoff->SetAttribute("DataRate", DataRateValue(DataRate(drv.Get().GetBitRate()+step)));
 
  p->output.push_back(Point2d<double>(theoric_mbps, measured_mbps));
 
  Simulator::Schedule(Seconds(SAMPLING_PERIOD_s), &ExperimentWifiPairs::AdvanceRate, this, p);
}
 
std::vector<Ptr<Pair> > ExperimentWifiPairs::Run(const std::string sstart,
						 const std::string sstep,
						 const int numstep,
						 const int numpairs,
						 const unsigned int packet_size,
						 const YansWifiChannelHelper &channel)
{
  step = DataRate(sstep).GetBitRate();
 
  std::vector<Ptr<Pair> > res;
 
  NodeContainer APnodes,MSnodes;
  NetDeviceContainer APdevices, MSdevices;
  Ipv4InterfaceContainer MSinterfaces, APinterfaces;
  MobilityHelper mobility;
  InternetStackHelper stack;
  Ipv4AddressHelper address;
 
  ChannelChooser chan = ChannelChooser(NUM_WIFI_CHANNELS);
 
  APnodes.Create(numpairs);
  MSnodes.Create(numpairs);
 
  NqosWifiMacHelper mac = NqosWifiMacHelper::Default();  
  YansWifiPhyHelper phy = YansWifiPhyHelper::Default();  
  phy.SetChannel(channel.Create());
 
  // now going to the MAC layer
  WifiHelper wifi = WifiHelper::Default();
  wifi.SetStandard(USED_WIFI_STANDARD); // USED... is a global in this file
  //wifi.SetRemoteStationManager("ns3::ArfWifiManager");
 
  // then we add mobility to the mobile stations 
  mobility.SetPositionAllocator("ns3::GridPositionAllocator",
				"MinX", DoubleValue(0.0),
				"MinY", DoubleValue(0.0),
				"DeltaX", DoubleValue(DX_m),
				"DeltaY", DoubleValue(DY_m),
				"GridWidth", UintegerValue(2),
				"LayoutType", StringValue("RowFirst")
				);
  mobility.SetMobilityModel("ns3::ConstantPositionMobilityModel");
 
  for(int i = 0; i < numpairs; i++)
    {
      std::ostringstream oss;
      oss<<"exp_ssid_"<<i;
      std::string ssid = oss.str();
 
      phy.Set("ChannelNumber", UintegerValue(chan.get(i)));
 
      // creating a MAC layer and devices for the MS
      mac.SetType("ns3::StaWifiMac",
		  "Ssid", SsidValue(ssid),
		  "ActiveProbing", BooleanValue(false)
		  ,"QosSupported", BooleanValue(false)
		  );
      MSdevices.Add(wifi.Install(phy, mac, MSnodes.Get(i)));
      //and same for the AP
 
      mac.SetType("ns3::ApWifiMac",
		  "Ssid", SsidValue(ssid),
		  "BeaconGeneration", BooleanValue(true),
		  "BeaconInterval", TimeValue(Seconds(2.5))
		  ,"QosSupported", BooleanValue(false)
                  );
      APdevices.Add(wifi.Install(phy, mac, APnodes.Get(i)));
 
      mobility.Install(APnodes.Get(i));
      mobility.Install(MSnodes.Get(i));
    }
 
  // now we load internet stacks into our nodes
  stack.Install(APnodes);
  stack.Install(MSnodes);
 
  address.SetBase("192.128.0.0", "255.255.255.0");
 
  MSinterfaces = address.Assign(MSdevices);
  APinterfaces = address.Assign(APdevices);
 
  // now the routes
  ns3::Ipv4GlobalRoutingHelper::PopulateRoutingTables();
  // and our network is configured!!
 
  for(int i = 0; i < numpairs; i++)
    {
      Ptr<Pair> p = CreateObject<Pair>();
      p->server = APnodes.Get(i);
      p->client = MSnodes.Get(i);
      p->bits = 0;
      res.push_back(p);
 
      Ipv4Address serverAddress = APinterfaces.GetAddress(i, 0);
 
      OnOffHelper onoff("ns3::UdpSocketFactory", Address(InetSocketAddress(serverAddress,22)));
      onoff.SetAttribute("OnTime", RandomVariableValue(ConstantVariable(SAMPLING_PERIOD_s*(numstep+1)+1.5)));
      onoff.SetAttribute("OffTime", RandomVariableValue(ConstantVariable(0)));
      onoff.SetAttribute("DataRate", DataRateValue(DataRate(sstart)));
      onoff.SetAttribute("PacketSize", UintegerValue(packet_size)); // bytes/s
 
      ApplicationContainer apps = onoff.Install(p->client);
      apps.Start(Seconds(1.5));
      apps.Stop(Seconds(SAMPLING_PERIOD_s*(numstep+1)+1.5));
 
      Simulator::Schedule(Seconds(SAMPLING_PERIOD_s+1.5), &ExperimentWifiPairs::AdvanceRate, this, p);
      Ptr<Socket> recvSink = SetupPacketReceive(p, Address(InetSocketAddress(serverAddress,22)));
    }
 
  if(OUTPUT_PCAP)
    phy.EnablePcap("iperf_sim", (IFTRUEOUTPUTAP_PCAP ? APdevices : MSdevices));
 
  Simulator::Stop(Seconds(SAMPLING_PERIOD_s*(numstep+1)+2));
  //Simulator::Run();
  Visualizer::Run();
  Simulator::Destroy();
 
  return res;
}
 
// should divide this one when i get time
void processRunsData(const std::vector<std::vector<Ptr<Pair> > > p,
		     Gnuplot2dDataset &d)
{
  const int num_runs = p.size();
  if(num_runs == 0)
    return;
  const int num_pairs = p.back().size();
  if(num_pairs == 0)
    return;
  const int num_samples = p.back().back()->output.size();
  if(num_samples == 0)
    return;
  double acc,iacc;
  // look at how to optimize this loop later
  // it can VERY probably be HEAVY optimized
  // though I don't know if this small summing
  // operations needs such optim
  for(int i = 0; i < num_samples; i++)
    { // but here we will process the data for 1 run
      acc = 0.0;
      for(int j = 0; j < num_pairs; j++)
	{
	  iacc = 0.0;
	  for(int k = 0; k < num_runs; k++)
	    iacc += p[k][j]->output[i].y; // acc of one pair in the runs
	  acc += iacc / num_runs; // we mean the iacc on the runs
	}
      d.Add(p.back().back()->output[i].x*num_pairs, acc);
    }
}
 
int main(int argc, char *argv[])
{
  Gnuplot gnuplot = Gnuplot(OUTPUT_IMAGE_NAME);
 
  ExperimentWifiPairs experiment;
  Gnuplot2dDataset ds;
  YansWifiChannelHelper channel;
 
  //YansWifiChannelHelper channel;
  //channel.SetPropagationDelay("ns3::ConstantSpeedPropagationDelayModel", "Speed", DoubleValue(299792458));
  //channel.AddPropagationLoss("ns3::RandomPropagationLossModel", "Variable", RandomVariableValue(ConstantVariable(50)));
 
  typedef std::vector<Ptr<Pair> > ConfigRun;
  for(unsigned int i = 0; i < NUM_CONFIGS; i++)
    {
      std::vector<ConfigRun> config_runs;
      channel = YansWifiChannelHelper::Default();
      std::ostringstream oss;
      oss<<CONFIGS[i].NUM_PAIRS<<" pairs "
	 <<CONFIGS[i].PACKET_SIZE<<"bits/packet";
      for(unsigned int j = 0; j < NUM_RUNS_BY_CONFIG; j++)
	{
	  SeedManager::SetRun(i*NUM_RUNS_BY_CONFIG+j);
	  config_runs.push_back(experiment.Run(CONFIGS[i].START,
					       CONFIGS[i].STEP,
					       CONFIGS[i].NUM_STEPS,
					       CONFIGS[i].NUM_PAIRS,
					       CONFIGS[i].PACKET_SIZE,
					       channel));
	}
      ds = Gnuplot2dDataset(oss.str());
      ds.SetStyle(Gnuplot2dDataset::LINES_POINTS);
      processRunsData(config_runs, ds);
      gnuplot.AddDataset(ds);
    }
  std::ostringstream oss;
  const std::string ws = 
    (USED_WIFI_STANDARD == WIFI_PHY_STANDARD_80211a ? "802.11a" :
     (USED_WIFI_STANDARD == WIFI_PHY_STANDARD_80211b ? "802.11b" :
      (USED_WIFI_STANDARD == WIFI_PHY_STANDARD_80211g ? "802.11g" :
       "???")));
  oss<<"WiFi UDP pairs Mbps/Mbps "<<NUM_RUNS_BY_CONFIG<<" runs";
  oss<<" - "<<SAMPLING_PERIOD_s<<"s/point ";
  oss<<ws<<" "<<NUM_WIFI_CHANNELS<<" channels used";
  gnuplot.SetTitle(oss.str());
 
  gnuplot.GenerateOutput(std::cout);
 
  return 0;
}