--- a/src/contrib/pyviz.cc	Tue Sep 02 15:12:00 2008 +0100
+++ b/src/contrib/pyviz.cc	Wed Sep 03 14:26:49 2008 +0100
@@ -20,11 +20,37 @@
 
 #include "pyviz.h"
 #include "ns3/simulator.h"
+#include "ns3/config.h"
+#include "ns3/node-list.h"
+
+static
+std::vector<std::string>
+PathSplit (std::string str)
+{
+  std::vector<std::string> results;
+  size_t cutAt;
+  while ((cutAt = str.find_first_of('/')) != str.npos)
+    {
+      if(cutAt > 0)
+        {
+          results.push_back(str.substr(0,cutAt));
+        }
+      str = str.substr(cutAt+1);
+    }
+  if (str.length() > 0)
+    {
+      results.push_back(str);
+    }
+  return results;
+}
+
 
 namespace ns3 {
 
 PyViz::PyViz ()
 {
+  Config::Connect ("/NodeList/*/DeviceList/*/Tx", MakeCallback (&PyViz::TraceNetDevTx, this));
+  Config::Connect ("/NodeList/*/DeviceList/*/Rx", MakeCallback (&PyViz::TraceNetDevRx, this));
 }
 
 PyViz::~PyViz ()
@@ -34,11 +60,89 @@
 void
 PyViz::SimulatorRunUntil (Time time)
 {
+  m_transmissionSamples.clear ();
   while (!Simulator::IsFinished () && Simulator::Now () < time)
     {
       Simulator::RunOneEvent ();
     }
 }
 
+bool PyViz::TransmissionSampleKey::operator < (PyViz::TransmissionSampleKey const &other) const
+{
+  if (transmitter < other.transmitter)
+    return true;
+  if (transmitter > other.transmitter)
+    return false;
+  if (receiver < other.receiver)
+    return true;
+  if (receiver > other.receiver)
+    return false;
+  return (channel < other.channel);
+}
+
+void
+PyViz::TraceNetDevTx (std::string context, Ptr<const Packet> packet, Mac48Address address)
+{
+  std::vector<std::string> splitPath = PathSplit (context);
+  int nodeIndex = atoi (splitPath[1].c_str ());
+  int devIndex = atoi (splitPath[3].c_str ());
+  Ptr<Node> node = NodeList::GetNode (nodeIndex);
+  Ptr<NetDevice> device = node->GetDevice (devIndex);
+  TxRecord record = { node };
+  if (address == device->GetBroadcast ())
+    {
+      TransmissionSampleKey key = { node, NULL, device->GetChannel() };
+      TransmissionSampleValue &sample = m_transmissionSamples[key];
+      sample.bytes += packet->GetSize ();
+    }
+  else
+    {
+      m_txRecords[packet->GetUid ()] = record;
+    }
 }
 
+void
+PyViz::TraceNetDevRx (std::string context, Ptr<const Packet> packet, Mac48Address address)
+{
+  std::vector<std::string> splitPath = PathSplit (context);
+  int nodeIndex = atoi (splitPath[1].c_str ());
+  int devIndex = atoi (splitPath[3].c_str ());
+  Ptr<Node> node = NodeList::GetNode (nodeIndex);
+  Ptr<NetDevice> device = node->GetDevice (devIndex);
+  Ptr<Channel> channel = device->GetChannel ();
+
+  std::map<uint32_t, TxRecord>::iterator recordIter = m_txRecords.find (packet->GetUid ());
+  if (recordIter == m_txRecords.end ())
+    {
+      return;
+    }
+  
+  TxRecord &record = recordIter->second;
+  
+  TransmissionSampleKey key = { record.srcNode, node, channel };
+  TransmissionSampleValue &sample = m_transmissionSamples[key];
+  sample.bytes += packet->GetSize ();
+  m_txRecords.erase (recordIter);
+}
+
+PyViz::TransmissionSampleList
+PyViz::GetTransmissionSamples () const
+{
+  TransmissionSampleList list;
+  for (std::map<TransmissionSampleKey, TransmissionSampleValue>::const_iterator
+         iter = m_transmissionSamples.begin ();
+       iter !=  m_transmissionSamples.end ();
+       iter++)
+    {
+      TransmissionSample sample;
+      sample.transmitter = iter->first.transmitter;
+      sample.receiver = iter->first.receiver;
+      sample.channel = iter->first.channel;
+      sample.bytes = iter->second.bytes;
+      list.push_back (sample);
+    }
+  return list;
+}
+
+}
+

--- a/src/contrib/pyviz.h	Tue Sep 02 15:12:00 2008 +0100
+++ b/src/contrib/pyviz.h	Wed Sep 03 14:26:49 2008 +0100
@@ -24,6 +24,11 @@
 #define NS3_PYVIZ_H
 
 #include "ns3/nstime.h"
+#include "ns3/node.h"
+#include "ns3/channel.h"
+#include "ns3/packet.h"
+#include "ns3/mac48-address.h"
+#include <map>
 
 namespace ns3 {
 
@@ -35,7 +40,46 @@
   
   // Run simulation until a given (simulated, absolute) time is reached
   void SimulatorRunUntil (Time time);
-  
+
+  struct TransmissionSample
+  {
+    Ptr<Node> transmitter;
+    Ptr<Node> receiver; // NULL if broadcast
+    Ptr<Channel> channel;
+    uint32_t bytes;
+  };
+
+  typedef std::vector<TransmissionSample> TransmissionSampleList;
+
+  TransmissionSampleList GetTransmissionSamples () const;
+
+private:
+
+  struct TxRecord
+  {
+    Ptr<Node> srcNode;
+  };
+  std::map<uint32_t, TxRecord> m_txRecords;
+
+  struct TransmissionSampleKey
+  {
+    bool operator < (TransmissionSampleKey const &other) const;
+    Ptr<Node> transmitter;
+    Ptr<Node> receiver; // NULL if broadcast
+    Ptr<Channel> channel;
+  };
+
+  struct TransmissionSampleValue
+  {
+    TransmissionSampleValue () : bytes (0) {}
+    uint32_t bytes;
+  };
+
+  std::map<TransmissionSampleKey, TransmissionSampleValue> m_transmissionSamples;
+
+  void TraceNetDevTx (std::string context, Ptr<const Packet> packet, Mac48Address address);
+  void TraceNetDevRx (std::string context, Ptr<const Packet> packet, Mac48Address address);
+
 };