/******************************************************************************
* macro to test SetBranchAddress()                                            *
*                                                                             *
* Author: Dr. Christian Stratowa, Vienna, Austria.                            *
* Created: 21 Sep 2008                             Last modified: 21 Sep 2008 *
******************************************************************************/

///////////////////////////
//
// in new root session(s):
//   create file "MyTest.root"
//     > .L macroTestBranch.C+ 
//     > CreateFile("MyTest.root", 10000) 
//
// in new root session(s):
//   draw plot
//     > .L macroTestBranch.C+ 
//     > TestBranch("MyTest.root", "MyTreeXY", 1) 
//     > TestBranch("MyTest.root", "MyTreeXY", 0) 
//     > TestBranch("MyTest.root", "MyTreeXYZ", 1) 
//     > TestBranch("MyTest.root", "MyTreeXYZ", 0) 
//
///////////////////////////

#include "TBranch.h"
#include "TFile.h"
#include "TLeaf.h"
#include "TMath.h"
#include "TRandom.h"
#include "TTree.h"

#ifndef __CINT__
#include <Riostream.h>
#endif

#include "float.h"


class ClassXY {

   protected:
      Int_t     fN;   
      Double_t  fX;   
      Double_t  fY;  

   public :
      ClassXY()          {}
      virtual ~ClassXY() {}

      void  SetN(Int_t n)    {fN = n;}
      void  SetX(Double_t x) {fX = x;}
      void  SetY(Double_t y) {fY = y;}
      Int_t    GetN()  const {return fN;}
      Double_t GetX()  const {return fX;}
      Double_t GetY()  const {return fY;}

      ClassDef(ClassXY,1) //ClassXY
};

class ClassXYZ: public ClassXY {

   protected:
      Double_t  fZ;   

   public :
      ClassXYZ()          {}
      virtual ~ClassXYZ() {}

      void  SetZ(Double_t z) {fZ = z;}
      Double_t GetZ() const  {return fZ;}

      ClassDef(ClassXYZ,1) //ClassXYZ
};

ClassImp(ClassXY);
ClassImp(ClassXYZ);


//______________________________________________________________________________
void CreateFile(const char *filename="MyTest.root", Int_t nentries = 10000)
{
// Create "MyTest.root"
   TFile *file = new TFile(filename, "recreate");

// Fill treeXY
   TTree *treeXY = new TTree("MyTreeXY", "classXY");
   if (treeXY == 0) return;

   Int_t    split   = 99;
   ClassXY *classXY = new ClassXY();
   treeXY->Branch("MyBranch", "ClassXY", &classXY, 64000, split);

   Double_t px, py;
   for (Int_t j=0; j<nentries; j++) {
      gRandom->Rannor(px,py);

      classXY->SetN(j);
      classXY->SetX(px);
      classXY->SetY(py);

      treeXY->Fill();
   }//for_j

   treeXY->Write();

// Fill treeXYZ
   TTree *treeXYZ = new TTree("MyTreeXYZ", "classXYZ");
   if (treeXYZ == 0) return;

   ClassXYZ *classXYZ = new ClassXYZ();
   treeXYZ->Branch("MyBranch", "ClassXYZ", &classXYZ, 64000, split);

   Double_t pz;
   for (Int_t j=0; j<nentries; j++) {
      gRandom->Rannor(px,py);
      pz = px*px + py*py;

      classXYZ->SetN(j);
      classXYZ->SetX(px);
      classXYZ->SetY(py);
      classXYZ->SetZ(pz);

      treeXYZ->Fill();
   }//for_j

   treeXYZ->Write();

   delete file;
}//CreateFile

//______________________________________________________________________________
//TTree *InitMyTree(const char *treename, ClassXY *classXY, Int_t &nentries)
TTree *InitMyTree(const char *treename, void *classXY, Int_t &nentries)
{
   TTree *tree = (TTree*)gDirectory->Get(treename); 
   if (tree == 0) return 0;

   //if (strcmp(tree->GetTitle(), "classXY") == 0) { 
      //tree->SetBranchAddress("MyBranch", &classXY);
      tree->SetBranchAddress("MyBranch", classXY);
   //} else if (strcmp(tree->GetTitle(), "classXYZ") == 0) {
   //   classXY = ((ClassXYZ*)classXY);
   //   tree->SetBranchAddress("MyBranch", &classXY);
   //}//if

   nentries = tree->GetEntries();

   return tree;
}//InitMyTree

//______________________________________________________________________________
Int_t *FillArrays(TTree *tree, ClassXY *&classXY,
                  Int_t n, Int_t *arr, Double_t *x, Double_t *y)
{
   for (Int_t i=0; i<n; i++) {
      tree->GetEntry(i);

      arr[i] = classXY->GetN();
      x[i]   = classXY->GetX();
      y[i]   = classXY->GetY();
   }//for_i

   return arr;
}//FillArrays

//______________________________________________________________________________
Int_t *FillArrays(TTree *tree, ClassXYZ *&classXYZ,
                  Int_t n, Int_t *arr, Double_t *x, Double_t *y, Double_t *z)
{
   for (Int_t i=0; i<n; i++) {
      tree->GetEntry(i);

      arr[i] = classXYZ->GetN();
      x[i]   = classXYZ->GetX();
      y[i]   = classXYZ->GetY();
      z[i]   = classXYZ->GetZ();
   }//for_i

   return arr;
}//FillArrays

//______________________________________________________________________________
Int_t TestBranch(const char *filename = "MyTest.root",
                 const char *treename = "MyTreeXY",
                 Int_t setbr = 0)
{
// Open file
   TFile *file = TFile::Open(filename, "READ");
   if (!file) return 1;

   Int_t     err  = 0;
   Int_t     size = 0;
   Int_t    *arrN = 0;
   Double_t *arrX = 0;
   Double_t *arrY = 0;
   Double_t *arrZ = 0;

   TTree   *mytree  = 0; 
   ClassXY *classXY = 0;
//   ClassXY *classXY = new ClassXY();

cout << "*classXY = " << classXY << endl;

// Init mytree and set branch address
   mytree = InitMyTree(treename, &classXY, size); //<==========
   if (mytree == 0) return 1;

cout << "*classXY = " << classXY << endl;

   if (setbr) {
      //if (strcmp(mytree->GetTitle(), "classXYZ") == 0) classXY = ((ClassXYZ*)classXY);
      mytree->SetBranchAddress("MyBranch", &classXY);
   }

cout << "*classXY = " << classXY << endl;

// Init memory (What are you doing here with nothrow ?!!)
   //if (!(arrN = new (nothrow) Int_t[size]))    {err = 1; goto cleanup;}
   //if (!(arrX = new (nothrow) Double_t[size])) {err = 1; goto cleanup;}
   //if (!(arrY = new (nothrow) Double_t[size])) {err = 1; goto cleanup;}
   //if (!(arrZ = new (nothrow) Double_t[size])) {err = 1; goto cleanup;}
   arrN = new  Int_t[size];    
   arrX = new  Double_t[size]; 
   arrY = new  Double_t[size]; 
   arrZ = new  Double_t[size]; 

   for (Int_t i=0; i<10; i++) {
      arrN[i] = 0;
      arrX[i] = arrY[i] = arrZ[i] = 0.0;
   }//for_i

// Fill arrays
   if (strcmp(mytree->GetTitle(), "classXY") == 0) { 
      arrN = FillArrays(mytree, classXY, size, arrN, arrX, arrY);
   } else if (strcmp(mytree->GetTitle(), "classXYZ") == 0) {
      arrN = FillArrays(mytree, (ClassXYZ*&)classXY, size, arrN, arrX, arrY, arrZ);
   }//if

// Print arrays
   for (Int_t i=0; i<6; i++) {
      cout << "n= " << arrN[i] << "   x= " << arrX[i] << "   y= " << arrY[i] << "   z= " << arrZ[i] << endl;
   }//for_i

// Cleanup
//cleanup:
//?   SafeDelete(classXY);
   if (arrZ) {delete [] arrZ; arrZ = 0;}
   if (arrY) {delete [] arrY; arrY = 0;}
   if (arrX) {delete [] arrX; arrX = 0;}
   if (arrN) {delete [] arrN; arrN = 0;}

   return err;
}//TestBranch