#include "TFile.h"
#include "TTree.h"
#include "TRandom.h"
#include "TCanvas.h"
#include "TMath.h"
#include "TSystem.h"
#include "Riostream.h"

void tree3w(Int_t nhugebranch, Bool_t fillzero) {
// Example of a Tree where branches are variable length arrays
// A second Tree is created and filled in parallel.
// Run this script with
//   .x tree3.C
// In the function treer, the first Tree is open.
// The second Tree is declared friend of the first tree.
// TTree::Draw is called with variables from both Trees.
//
//  Author: Rene Brun

   const Int_t kMaxTrack = 500;
   Int_t ntrack;
   Int_t stat[kMaxTrack];
   Int_t sign[kMaxTrack];
   Float_t px[kMaxTrack];
   Float_t py[kMaxTrack];
   Float_t pz[kMaxTrack];
   Float_t pt[kMaxTrack];
   Float_t zv[kMaxTrack];
   Float_t chi2[kMaxTrack];
   Double_t sumstat;
   Float_t *hugebranch0 = new Float_t[500000];
   Float_t *hugebranch1 = new Float_t[500000];
   Float_t *hugebranch2 = new Float_t[500000];
   Float_t *hugebranch3 = new Float_t[500000];
   Float_t *hugebranch4 = new Float_t[500000];

   memset(hugebranch0, 0, sizeof(Float_t)*500000);
   memset(hugebranch1, 0, sizeof(Float_t)*500000);
   memset(hugebranch2, 0, sizeof(Float_t)*500000);
   memset(hugebranch3, 0, sizeof(Float_t)*500000);
   memset(hugebranch4, 0, sizeof(Float_t)*500000);

   TFile f("tree3.root","recreate");
   TTree *t3 = new TTree("t3","Reconst ntuple");
   t3->Branch("ntrack",&ntrack,"ntrack/I");
   t3->Branch("stat",stat,"stat[ntrack]/I");
   t3->Branch("sign",sign,"sign[ntrack]/I");
   t3->Branch("px",px,"px[ntrack]/F");
   t3->Branch("py",py,"py[ntrack]/F");
   t3->Branch("pz",pz,"pz[ntrack]/F");
   t3->Branch("zv",zv,"zv[ntrack]/F");
   t3->Branch("chi2",chi2,"chi2[ntrack]/F");
   if (nhugebranch >= 1 || fillzero) {
      t3->Branch("hugebranch0",hugebranch0,"hugebranch0[500000]/F");
   }
   if (nhugebranch >= 2) {
      t3->Branch("hugebranch1",hugebranch1,"hugebranch1[500000]/F");
   }
   if (nhugebranch >= 3) {
      t3->Branch("hugebranch2",hugebranch2,"hugebranch2[500000]/F");
   }
   if (nhugebranch >= 4) {
      t3->Branch("hugebranch3",hugebranch3,"hugebranch3[500000]/F");
   }
   if (nhugebranch >= 5) {
      t3->Branch("hugebranch4",hugebranch4,"hugebranch4[500000]/F");
   }

   for (Int_t i=0;i<(fillzero?3000:30);i++) {
      cout<<i<<endl;
      Int_t nt = (Int_t)(gRandom->Rndm()*(kMaxTrack-1));
      ntrack = nt;
      sumstat = 0;
      for (Int_t n=0;n<nt;n++) {
         stat[n] = n%3;
         sign[n] = i%2;
         px[n]   = gRandom->Gaus(0,1);
         py[n]   = gRandom->Gaus(0,2);
         pz[n]   = gRandom->Gaus(10,5);
         zv[n]   = gRandom->Gaus(100,2);
         chi2[n] = gRandom->Gaus(0,.01);
         sumstat += chi2[n];
         pt[n]   = TMath::Sqrt(px[n]*px[n] + py[n]*py[n]);
      }
      if (!fillzero) {
         for (Int_t j = 0; j < 500000; j++) {
            hugebranch0[j] = gRandom->Gaus(100,2);
            hugebranch1[j] = gRandom->Gaus(100,2);
            hugebranch2[j] = gRandom->Gaus(100,2);
            hugebranch3[j] = gRandom->Gaus(100,2);
            hugebranch4[j] = gRandom->Gaus(100,2);
         }
      }
      t3->Fill();
   }
   t3->Print();
   cout<<"filling is done. sleeping 10 sec."<<endl;
   gSystem->Sleep(10000); // just to see memory by 'top'
   f.cd();
   t3->Write();

   delete [] hugebranch0;
   delete [] hugebranch1;
   delete [] hugebranch2;
   delete [] hugebranch3;
   delete [] hugebranch4;
}

void tree3(Int_t nhugebranch, Bool_t fillzero)
{
   tree3w(nhugebranch, fillzero);
}