#include "TRandom.h"
#include "TFile.h"
#include "TTree.h"
#include "TBranch.h"
#include "TRandom.h"
#include "TTimeStamp.h"

#include <iostream>
#include "math.h"

using namespace std;

TRandom *RandomGenerator;
ULong64_t Entries, Times;
TString DemoFileName;

// BEGIN TREE CREATION
void CreateTree(ULong64_t Entries, Int_t CompressionLevel) {
  UChar_t Value;
  TFile *DemoFile=new TFile(DemoFileName, "RECREATE");
  DemoFile->SetCompressionLevel(CompressionLevel);
  TTree *DemoTree=new TTree("DemoTree", "DemoTree");
  DemoTree->Branch("Value", &Value, "Value/b");
  for (ULong64_t Counter=0; Counter<Entries; Counter++) {
    Value=(UChar_t)(ceil(RandomGenerator->Rndm()*256)-1);
    DemoTree->Fill();
  }
  DemoTree->Write();
  delete DemoTree;
  DemoFile->Close();
  delete DemoFile;
}
// END TREE CREATION

// BEGIN TEST SEQUENTIAL READ
void TestSequentialRead(ULong64_t FromEvent, ULong64_t ToEvent, ULong64_t Delta) {
  TFile *DemoFile=new TFile(DemoFileName, "READONLY");
  TTree *DemoTree=(TTree*)DemoFile->Get("DemoTree");
  ULong64_t Counter=FromEvent;
  while (Counter!=ToEvent) {
    DemoTree->GetEntry(Counter);
    Counter+=Delta;
  }
  DemoFile->Close();
  delete DemoFile;
}
// END TEST SEQUENTIAL READ

// BEGIN TEST RANDOM READ
void TestRandomRead(ULong64_t Min, ULong64_t Max, ULong64_t Times) {
  UChar_t Value;
  TFile *DemoFile=new TFile(DemoFileName, "READONLY");
  TTree *DemoTree=(TTree*)DemoFile->Get("DemoTree");
  DemoTree->SetBranchAddress("Value", &Value);
  ULong64_t Choice;
  for (ULong64_t Counter=0; Counter<Times; Counter++) {
    Choice=(ULong64_t)(ceil(RandomGenerator->Rndm()*(Max-Min+1))-1+Min);
    DemoTree->GetEntry(Choice);
  }
  DemoFile->Close();
  delete DemoFile;
}
// END TESTRANDOM READ

// BEGIN TEST SEQUENCE
void FullTest(Int_t CompressionLevel) {
  TTimeStamp StartTime, EndTime;
  Double_t Elapsed, TotalTime;
  cout << "--- Compression level: " << CompressionLevel << endl << endl;
  TString TestName[]={"Tree creation", "Full sequential read", "Reverse sequential read", "Sequential read first half", "Sequential read last half", "Random read full range", "Random read first half", "Random read last half"};
  ULong64_t Runs[]={1, 5, 5, 5, 5, 3, 3, 3};
  ULong64_t TestEntries[]={Entries, Entries, Entries, Entries/2, Entries/2, Times, Times, Times};
  
  for (Int_t TestType=0; TestType<=7; TestType++) {
    TotalTime=0.;
    cout << TestName[TestType] << " (Entries to parse: " << TestEntries[TestType] << ") (Runs: " << Runs[TestType] << ")" << endl;
    for (ULong64_t Counter=0; Counter<Runs[TestType]; Counter++) {
      StartTime.Set();
      if (TestType==0) { CreateTree(Entries, CompressionLevel); }
      if (TestType==1) { TestSequentialRead(0, Entries-1, 1); }
      if (TestType==2) { TestSequentialRead(Entries-1, 0, -1); }
      if (TestType==3) { TestSequentialRead(0, Entries/2-1, 1); }
      if (TestType==4) { TestSequentialRead(Entries/2, Entries-1, 1); }
      if (TestType==5) { TestRandomRead(0, Entries-1, Times); }
      if (TestType==6) { TestRandomRead(0, Entries/2-1, Times); }
      if (TestType==7) { TestRandomRead(Entries/2, Entries-1, Times); }
      EndTime.Set();
      Elapsed=EndTime.AsDouble()-StartTime.AsDouble();
      TotalTime=TotalTime+Elapsed;
      cout << "- Run " << Counter+1 << ": " << Elapsed << " s" << endl;
    }
    cout << "Mean: " << TotalTime/Runs[TestType] << " s" << endl << endl;
  }
}
// END TEST SEQUENCE

// BEGIN FUNCTION TESTTREE
void TestTree() {
  // Settings
  DemoFileName="demo.root";
  Entries=100000000;
  Times=10000;

  // Initialization
  RandomGenerator=new TRandom(0);
  
  // Run
  FullTest(0);
  FullTest(1);
  FullTest(9);
}
// END FUNCTION TESTTREE

int main(int, char**) {
  TestTree();
}