#include <iostream.h>
#include "TTree.h"
#include "TFile.h"
#include "TCanvas.h"
#include "TBranch.h"
#include "TBenchmark.h"
#include "TSystem.h"
#include "TH1.h"
#include "TH2.h"

void MakeDelay(int time_in_sec) {
   gSystem->Sleep(time_in_sec*1000);
}
    
void PurgeMemory(bool show = false) {
  const int RAMSIZE = 512;
  
  if (show) cout << "Purge " << RAMSIZE << " MB of memory" << endl;
    
  Double_t** buffers = (Double_t**) malloc(RAMSIZE*4);
  
  for(int n=0;n<RAMSIZE;n++) {
     buffers[n] = new Double_t[125000];
     if (show && (n % 10 == 0)) cout << ".";
     memset(buffers[n],123,1000000);
  }
  
  if (show) cout << "Done" << endl;
  
  for(int n=0;n<RAMSIZE;n++) 
     delete[] buffers[n];
  
  free(buffers);
}

void ProduceTree(const char* filename, const char* treename, 
                 Int_t numentries, Int_t compression,
                 int numbranches, int branchsize, int buffersize) {
   cout << "Producing tree " << treename << " ..." << endl;                  
                     
   TFile f(filename,"recreate");
   f.SetCompressionLevel(compression);
   
   TTree t(treename, treename);
   
   Float_t* data = new Float_t[numbranches * branchsize];
   for (int i=0;i<numbranches * branchsize;i++)
      data[i] = 1.2837645786 * i;
      
   for (int nbr=0;nbr<numbranches;nbr++) {
      TString brname = "Branch";
      brname+=nbr;
      
      TString format = brname;
      format+="[";
      format+=branchsize;
      format+="]/F";
      
      t.Branch(brname.Data(),&(data[nbr*branchsize]),format.Data(), buffersize*branchsize*sizeof(Float_t));
   }   
   
   for (int n=0;n<numentries;n++)
     t.Fill();
   
   t.Write();
   
   delete[] data;
}

void TestTree(const char* filename, const char* treename, 
              int numbranches, int branchsize, int activebranches,
              float& RealTime, float& CpuTime) {
   TFile f(filename);
   TTree* t = (TTree*) f.Get(treename);
   
   Float_t* data = new Float_t[numbranches*branchsize];
   
   for(int nbr=0;nbr<numbranches;nbr++) {
      TString brname = "Branch";
      brname+=nbr;
      t->SetBranchAddress(brname,&(data[nbr*branchsize]));
   }
   
   if (activebranches<=0) t->SetBranchStatus("*",1); else {
      t->SetBranchStatus("*",0);
      for (int nbr=0;nbr<activebranches;nbr++) {
        TString brname = "Branch";
        brname+=nbr;
        t->SetBranchStatus(brname,1);
      }
   }
   
   Int_t counter = 0;
   
   gBenchmark->Reset();
   gBenchmark->Start("TestTree");
   
   while (t->GetEntry(counter++));
   
   gBenchmark->Stop("TestTree");
   
   gBenchmark->Show("TestTree");
   
   RealTime = gBenchmark->GetRealTime("TestTree");
   CpuTime = gBenchmark->GetCpuTime("TestTree");
   
   delete[] data;
}



void SingleTest(const char* filename, const char* treename,
                int NumBranches, int NumBuffersStep,
                Int_t NumberOfEvents, Int_t SizeOfBranch) {
    
   TFile f(filename,"recreate");
   TTree t(treename,"results of tree perfomance measurements");
   
   Int_t BufferSize, ActiveBranches;
   Float_t RealTime, CpuTime;
   
   t.Branch("BufferSize",&BufferSize,"BufferSize/I");    
   t.Branch("ActiveBranches",&ActiveBranches,"ActiveBranches/I");
   t.Branch("RealTime",&RealTime,"RealTime/F");
   t.Branch("CpuTime",&CpuTime,"CpuTime/F");
   
   BufferSize = 1;

   for(int bufcounter=0;bufcounter<NumBuffersStep;bufcounter++) {
      ProduceTree("TreeFile.root","TestTree", NumberOfEvents, 0, NumBranches, SizeOfBranch, BufferSize);
    
      for(ActiveBranches=1;ActiveBranches<=NumBranches;ActiveBranches++) {
         PurgeMemory();  
         cout << "Buffer size = " << BufferSize << " ActiveBranches = " << ActiveBranches << endl;
         MakeDelay(5);
         TestTree("TreeFile.root","TestTree",NumBranches, SizeOfBranch, ActiveBranches, RealTime, CpuTime);
         f.cd();
         t.Fill();
      }
      
      BufferSize*=2;
   }
    
   f.cd(); 
   t.Write();
}

void ShowTestResults(const char* filename, const char* treename, 
                     const char* CanvasName, const char* HistogramTitle, 
                     int NumBranches, int NumBuffersStep) {
   TFile f(filename);
   TTree* t = (TTree*) f.Get(treename);
   
   Int_t BufferSize, ActiveBranches;
   Float_t RealTime, CpuTime;
   
   t->SetBranchAddress("BufferSize",&BufferSize);    
   t->SetBranchAddress("ActiveBranches",&ActiveBranches);
   t->SetBranchAddress("RealTime",&RealTime);
   t->SetBranchAddress("CpuTime",&CpuTime);


   TH2D* histoReal = new TH2D("RealTime",TString(HistogramTitle) + " Real time",
                          NumBranches,0.5,0.5+NumBranches,NumBuffersStep,-0.5,-0.5+NumBuffersStep);
   histoReal->GetXaxis()->SetTitle("number of active branches");
   histoReal->GetYaxis()->SetTitle("relative buffer size (logarithm)");
   histoReal->GetZaxis()->SetTitle("time (s)");
   histoReal->SetDirectory(0);
   histoReal->SetStats(kFALSE);
   			  

   TH2D* histoCPU = new TH2D("CPUtime",TString(HistogramTitle) + " CPU time",
                          NumBranches,0.5,0.5+NumBranches,NumBuffersStep,-0.5,-0.5+NumBuffersStep);
   histoCPU->GetXaxis()->SetTitle("number of active branches");
   histoCPU->GetYaxis()->SetTitle("relative buffer size (logarithm)");
   histoCPU->GetYaxis()->SetTitle("time (s)");
   histoCPU->SetDirectory(0);
   histoCPU->SetStats(kFALSE);
   
   int counter=0;
   while (t->GetEntry(counter++)) {
      float value = TMath::Log(BufferSize)/TMath::Log(2.);
      histoReal->Fill(ActiveBranches,value,RealTime);
      histoCPU->Fill(ActiveBranches,value,CpuTime);
         
   }
   
   TString cname(CanvasName);
   TCanvas* c1 = new TCanvas(cname+"_Real","Results of tree tests");
   c1->SetPhi(145);
   c1->SetTheta(15);
   histoReal->Draw("LEGO");
   c1->SaveAs(cname+"_Real.gif");
   
//   TCanvas* c2 = new TCanvas(cname+"_CPU","Results of tree tests");
//   c2->SetPhi(145);
//   c2->SetTheta(15);
//   histoCPU->Draw("LEGO");
//   c2->SaveAs(cname+"_CPU.gif");
}
                    
    
        
void PerfomanceTest() {
   SingleTest("Tree_small.root","Perfomance", 10,10, 1000000, 10);  // 1000000 events x 10 branches x 10 doubles/branch
   SingleTest("Tree_large.root","Perfomance", 10,10, 10000, 1000);  // 10000 events x 10 branches x 1000 doubles/branch
    
   ShowTestResults("Tree_small.root","Perfomance","Tree_small","1000000 events x 10 branches x 10 doubles/branch", 10,10);
   ShowTestResults("Tree_large.root","Perfomance","Tree_large","10000 events x 10 branches x 1000 doubles/branch", 10,10);
}