diff --git a/ns-3_wifi_tests/wifi-test b/ns-3_wifi_tests/wifi-test
index 3c3eb62..c0432bb 100755
Binary files a/ns-3_wifi_tests/wifi-test and b/ns-3_wifi_tests/wifi-test differ
diff --git a/ns-3_wifi_tests/wifi-test.cc b/ns-3_wifi_tests/wifi-test.cc
index b6b0485..52ccf88 100644
--- a/ns-3_wifi_tests/wifi-test.cc
+++ b/ns-3_wifi_tests/wifi-test.cc
@@ -18,6 +18,8 @@
 #include "ns3/ipv4-address-helper.h"
 #include "ns3/ipv4-global-routing-helper.h"
 #include "ns3/constant-position-mobility-model.h"
+#include "ns3/energy-module.h"
+#include "ns3/wifi-radio-energy-model-helper.h"
 
 // C++ library
 #include <iostream>
@@ -125,6 +127,66 @@ void applyScenarios(Cell cell,int sensorsPktSize, int sensorsSendInterval){
   sinkApp.Start (Seconds (0));
 }
 
+void
+TotalEnergy (std::string context,double oldValue, double newValue)
+{
+  NS_LOG_UNCOND ("Energy Value of node " << context << " at time " <<Simulator::Now ().GetSeconds ()
+                 << "\t"<< newValue);
+}
+
+
+void setupEnergy(Cell cell){
+  NodeContainer nodes(cell.first.first,cell.first.second);
+  NetDeviceContainer nodesNetDev(cell.second.first,cell.second.second);
+
+  // Install energy source
+  BasicEnergySourceHelper edgeBasicSourceHelper;
+  edgeBasicSourceHelper.Set ("BasicEnergySourceInitialEnergyJ", DoubleValue (2.9009));
+  edgeBasicSourceHelper.Set ("BasicEnergySupplyVoltageV", DoubleValue (3.3));
+  EnergySourceContainer apEdgeNodesSources = edgeBasicSourceHelper.Install (cell.first.first);
+  EnergySourceContainer wifiEdgeNodesSources = edgeBasicSourceHelper.Install (cell.first.second);
+
+  // Install device energy model
+  WifiRadioEnergyModelHelper radioEnergyHelper;
+  radioEnergyHelper.Set ("TxCurrentA", DoubleValue (0.38));
+  radioEnergyHelper.Set ("RxCurrentA", DoubleValue (0.313));
+  radioEnergyHelper.Set ("IdleCurrentA", DoubleValue (0.273));
+  DeviceEnergyModelContainer edgeApDeviceModels = radioEnergyHelper.Install (cell.second.first, apEdgeNodesSources);
+  DeviceEnergyModelContainer edgeDeviceModels = radioEnergyHelper.Install (cell.second.second, wifiEdgeNodesSources);
+
+  
+  // Trace
+  DeviceEnergyModelContainer energyModels(edgeApDeviceModels, edgeDeviceModels);
+  DeviceEnergyModelContainer::Iterator it=energyModels.Begin();
+
+  int i=0;
+  while(it!=energyModels.End()){
+    (*it)->TraceConnect ("TotalEnergyConsumption", std::to_string(i),MakeCallback (&TotalEnergy));
+    it++;
+    i++;
+  }
+
+
+
+
+
+
+
+  //   Ptr<BasicEnergySource> basicSourcePtr0 = DynamicCast<BasicEnergySource> (wifiEdgeNodesSources.Get (0));
+
+  // //basicSourcePtr0->TraceConnectWithoutContext ("RemainingEnergy", MakeCallback (&RemainingEnergy));
+  // //device energy model
+
+  // Ptr<DeviceEnergyModel> basicRadioModelPtr0 =
+  // basicSourcePtr0->FindDeviceEnergyModels ("ns3::WifiRadioEnergyModel").Get (0);
+
+  // NS_ASSERT (basicRadioModelPtr0 != NULL);
+  // basicRadioModelPtr0->TraceConnectWithoutContext ("TotalEnergyConsumption", MakeCallback (&TotalEnergy));
+
+  
+}
+
+
 int main(int argc, char* argv[]){
 
   LogComponentEnable("UdpEchoClientApplication", LOG_LEVEL_INFO);
@@ -132,7 +194,7 @@ int main(int argc, char* argv[]){
 
   uint32_t sensorsFrequency=1;
   uint32_t sensorsPktSize=150;
-  uint32_t sensorsNumber=10;
+  uint32_t sensorsNumber=2;
 
   CommandLine cmd;
   cmd.AddValue ("sensorsSendInterval", "Number of temperature measurement per second (default 1)", sensorsFrequency);
@@ -143,9 +205,10 @@ int main(int argc, char* argv[]){
 
   Cell c=createCell(sensorsNumber);
   applyScenarios(c,sensorsPktSize,sensorsFrequency);
-
+  setupEnergy(c);
+  
   // Run simulators
-  Simulator::Stop (Seconds (10));
+  Simulator::Stop (Seconds (20));
   Simulator::Run ();
 
   // Destroy