//*-- Author : Jochen Markert 28.03.2011

//_HADES_CLASS_DESCRIPTION
////////////////////////////////////////////////////////////////////////////
//
// HLoop is a helper class to allow for fast looping of HADES dsts.
// The categories are maped directly from the Tree and allow partical
// reading to speed up. If Hades object exists, the current event structure
// will be replaced by the one defined by HLoop. If Hades not exists, it
// can be created by HLoop constructor. The Hades eventstructure is important
// if one wants to access the data via gHades->getCurrentEvent()->getCategory(cattype)
// or implicit by Classes like HParticleTrackSorter. See the example below for
// further explanation.
//
//
//
//    HEventHeader* getEventHeader()  // retrieve the pointer to the HADES HEventHeader
//    TTree*        getTree       ()  // retrieve the pointer to the HADES tree
//    TChain*       getChain      ()  // retrieve TChain pointer
//    Long64_t      getEntries    ()  // get Number of events in the chain
//
//--------------------------------------------------------
// example:
//
// #include "hloop.h"
// #include "hcategory.h"
// #include "heventheader.h"
// #include "hparticlecand.h"
// #include "hparticletracksorter.h"
//
// #include "TIterator.h"
//
// #include <iostream>
// using namespace std;
// void myLoop()
// {
//
//    //---------------LOOP CONFIG----------------------------------------------------------------
//    Bool_t createHades = kTRUE;  // kTRUE = create HADES new
//    HLoop* loop = new HLoop(createHades); // create HADES (needed if one wants to use gHades)
//    // add files : Input sources may be combined
//    loop->addFile ("myFile1.root");
//    loop->addFile ("myFile2.root");  // ....
//    loop->addFilesList("list.txt");  // add all files in list. list contains 1 root file per line (including path)
//    loop->addFiles("myFiles*.root"); // add all files matching this expression
//
//    if(!loop->setInput("-*,+HParticleCand")) // global disbale "-*" has to be first in the list
//    {                                        // read only one category from file (HEventHeader is allways on)
//        exit(1);                             // Correct name for real data / sim data have to be used here
//    }
//    loop->printCategories(); // print status from all categories in event
//    //------------------------------------------------------------------------------------------
//
//    TIterator* iterCand = 0;
//    if (loop->getCategory("HParticleCand")) iterCand = loop->getCategory("HParticleCand")->MakeIterator();
//    HEventHeader* header = loop->getEventHeader();
//
//    //--------------------------CONFIGURATION---------------------------------------------------
//    //At begin of the program (outside the event loop)
//
//    HParticleTrackSorter sorter;
//    //sorter.setDebug();                                            // for debug
//    //sorter.setPrintLevel(3);                                      // max prints
//    //sorter.setRICHMatching(HParticleTrackSorter::kUseRKRICHWindow,4.); // select matching RICH-MDC for selectLeptons() function
//    //sorter.setIgnoreInnerMDC();                                   // do not reject Double_t inner MDC hits
//    //sorter.setIgnoreOuterMDC();                                   // do not reject Double_t outer MDC hits
//    //sorter.setIgnoreMETA();                                       // do not reject Double_t META hits
//    //sorter.setIgnorePreviousIndex();                              // do not reject indices from previous selctions
//    sorter.init();                                                  // get catgegory pointers etc...
//    //--------------------------------------------------------------------------------------------
//
//    for (Int_t i=0; i < 10; i++) {
//        if(loop->nextEvent(i) <= 0) break;  // get next event. categories will be cleared before
//                                            // if 0 (entry not found) or -1 (Io error) is returned stop the loop
//
//        cout<<i<<"----------------------------------- "<<endl;
//        cout<<"sequence Nr = "<<header->getEventSeqNumber()<<endl; // retrieve full header infos
//
//
//        sorter.cleanUp();
//        sorter.resetFlags(kTRUE,kTRUE,kTRUE,kTRUE);     // reset all flags for flags (0-28) ,reject,used,lepton
//        Int_t nCandLep     = sorter.fill(HParticleTrackSorter::selectLeptons);   // fill only good leptons
//        Int_t nCandLepBest = sorter.selectBest(HParticleTrackSorter::kIsBestRKRKMETA,HParticleTrackSorter::kIsLepton);
//        Int_t nCandHad     = sorter.fill(HParticleTrackSorter::selectHadrons);   // fill only good hadrons (already marked good leptons will be skipped)
//        Int_t nCandHadBest = sorter.selectBest(HParticleTrackSorter::kIsBestRKRKMETA,HParticleTrackSorter::kIsHadron);
//
//        if(iterCand){
//           iterCand->Reset();
//
//           HParticleCand* cand;
//           while ( (cand = (HParticleCand*)iterCand->Next()) != 0 ){
//              // do some work  ....
//          }
//       }
//    }
//    sorter.finalize(); // clean up stuff
// }
////////////////////////////////////////////////////////////////////////////


#include "hloop.h"

#include "hades.h"
#include "htree.h"
#include "heventheader.h"
#include "hcategory.h"
#include "htool.h"

#include "TDirectory.h"
#include "TROOT.h"
#include "TList.h"
#include "TKey.h"
#include "TObject.h"
#include "TString.h"
#include "TObjString.h"
#include "TObjArray.h"
#include "TChainElement.h"
#include "TSystem.h"



#include <iostream>
#include <fstream>
#include <iomanip>

ClassImp(HLoop)

HLoop::HLoop(Bool_t createHades)
{
    // if createHades == kTRUE (default=kFALSE) Hades will be created and eventstructure set
    //                == kFALSE Hades will not be created, but if exists
    //                   the event structure will be replaced by the one defined in HLoop
    // if Hades does not exists and is not created new, the eventsructure will be
    // available only via HLoop and not via gHades

    fChain        = new TChain("T");
    fHead         = NULL;
    fTree         = NULL;
    fFileCurrent  = NULL;
    fMaxEntries   = 0;
    fCurrentEntry = 0;
    fCurrentName  ="unset";

    fHasCreatedHades = kFALSE;
    fRecEvent        = new HRecEvent();
    if(createHades){
	if(gHades == 0) {
            fHasCreatedHades = kTRUE;
	    Hades* hades = new Hades();
            hades->setEvent(fRecEvent);
            Info("HLoop()","HADES created and event structure set !" );
	} else {
	    Error("HLoop()","HADES already exists, but should be created new! Check your macro ..." );
            exit(1);
	}

    } else {
	if(gHades != 0) {
	    gHades->setEvent(fRecEvent);
            Info("HLoop()","HADES already exists, event structure replaced !" );
	} else {
            Info("HLoop()","HADES does not exists! No functions via gHades will be available!" );
	}
    }
}

HLoop::~HLoop()
{
    delete fChain;
    if(fHasCreatedHades && gHades) delete gHades;
}

Bool_t HLoop::addFile (TString infile)
{
    // add a single root file to the chain.
    // It will checked if the file exists.


    TObjArray* elements = fChain->GetListOfFiles();
    Int_t nfiles        = elements->GetEntries();

    if(gSystem->AccessPathName(infile.Data()) == 0) {
	cout<<setw(4)<<nfiles<<" add file : "<<infile.Data()<<endl;
	fChain->AddFile(infile.Data());
        return kTRUE;
    } else {
        Error("addFile()","File = %s not found! Will be skipped .... ",infile.Data());
        return kFALSE;
    }

}
Bool_t HLoop::addFiles(TString expression){

    // add all files matching the expression to the chain

    TObjArray* arr = HTool::glob(expression);
    Int_t n = arr->GetEntries();
    if(n == 0) {
	Warning("addFiles()","No file matching expression '%s' found found !",expression.Data());
    } else {
	for(Int_t i = 0; i < n; i ++){
	    TString name = ((TObjString*)arr->At(i))->GetString();
	    addFile(name);
	}
    }
    delete arr;
    return kTRUE;

}
Bool_t HLoop::addFilesList(TString filelist)
{
    // add all files in list filelist. The list should contain
    // 1 root file per line (including path). Return kFALSE if
    // list files does not exist or and error occured while
    // reading the file

    if(gSystem->AccessPathName(filelist.Data()) == 0 )
    {
	ifstream in;
	in.open(filelist.Data());
        TString name;
	while(!in.eof()){
	    in>>name;
	    if(!in.fail()) {
		addFile(name);
	    }
	    else {
               Error("addFilesList()","Error detected during reading the files list ! Will skipped all other files if any.... ");
	       in.close();
               return kFALSE;
	    }
	}
        in.close();

       return kTRUE;
    } else {
       Error("addFilesList()","File = %s not found! Will be skipped .... ",filelist.Data());
       return kFALSE;
    }
}

Bool_t HLoop::setInput(TString readCategories)
{
    // read in the given file structure. Categories can be
    // disabled/enabled by the string readCategories. This
    // string should contain a comma separated list. By default
    // all categories are enabled. The format should be
    // like (example 1) "-*,+HMdcCal1,+HParticleCand"
    // (switch off all categories and enable just the to given
    // ones. The global disable -* has to be the first in the
    // list. Or (example 2) "-HMdcRaw,-HRichHit" (enable all
    // categories, but not the give ones).  For simulation
    // or real data one has to use the correct coresponing
    // names (expample : HParticleCandSim / HParticleCand)


    TObjArray* fList  = fChain->GetListOfFiles();

    if(fList->GetEntries() == 0){
        Error("setInput()","No File in List!");
	return kFALSE;
    }



    TChainElement* el = (TChainElement*)fList->At(0);
    TString file = el->GetTitle();

    TFile* fIn = (TFile*)gROOT->GetListOfFiles()->FindObject(file.Data());
    if (!fIn) {
	fIn = new TFile(file.Data());
	if (!fIn) {
	    Error("setFile()","Could not open infile %s",file.Data());
	    return kFALSE;
	}
    }
    fIn->cd();

    fChain->LoadTree(0);

    fTree = (HTree*) fChain->GetTree();
    if(!fTree) {
	Error("setInput()","Could not find Tree in infile %s",file.Data());
	return kFALSE;
    }

    fHead = fRecEvent->getHeader();

    TBranch* brhead = fTree->GetBranch("EventHeader.");
    if(!brhead) {
	Error("setFile()","Eventheader Branch not found !");
	return kFALSE;
    } else {
	fChain->SetBranchAddress("EventHeader.",&fHead);
	fChain->SetBranchStatus("EventHeader.",1);
	fChain->SetBranchStatus("EventHeader.*",1);
    }

    fChain->SetBranchStatus("Event.",0);
    fChain->SetBranchStatus("Event.*",0);

    TList* dirkeys = gDirectory->GetListOfKeys();

    for(Int_t d = 0; d < dirkeys->GetEntries(); d ++){

	TKey* kdir = (TKey*)dirkeys->At(d);
	TString dirname = kdir->GetName();

	if(dirname.BeginsWith("dir"))
	{

	    TString partialEvent = dirname;
	    partialEvent.ReplaceAll("dir","");

	    if(fIn->cd(dirname.Data()) )
	    {
		TList* catkeys = gDirectory->GetListOfKeys();

		for(Int_t c = 0; c <catkeys->GetEntries(); c ++)
		{
		    fIn->cd(dirname.Data());
		    TKey*   kcat      = (TKey*) catkeys->At(c);
		    TString catname   = kcat->GetName();
		    TString classname = kcat->GetClassName();

		    if(classname == "HPartialEvent") {

			fChain ->SetBranchStatus(Form("%s.",catname.Data()),0);
			fChain ->SetBranchStatus(Form("%s.*",catname.Data()),0);
		        continue;
		    }

		    //------------------------------------------------------------
		    if(classname == "HLinearCategory" || classname == "HMatrixCategory")
		    {
			fEvent[catname] = (HCategory*) gDirectory->Get(catname.Data());
			if(!fEvent[catname]) {
			    Error("setInput()","NULL pointer for category %s retrieved!",catname.Data()) ;
			    return kFALSE;
			}
			fStatus [catname] = 0;
                        fPartial[catname] = partialEvent;

			TBranch* br = fTree->GetBranch(Form("%s.",catname.Data()));
			if(!br) {
			    Error("setInput()","Branch not found : %s !",Form("%s.",catname.Data()));
			    return kFALSE;
			} else {
			    fChain->SetBranchAddress(Form("%s.",catname.Data()),&fEvent[catname]);
			    fChain->SetBranchStatus (Form("%s.",catname.Data()),1);
			    fChain->SetBranchStatus (Form("%s.*",catname.Data()),1);
			    fStatus[catname] = 1;
			}
		    }
		    //------------------------------------------------------------

		} // end cat keys
	    } else {
		Error("setInput()","Could no enter dir %s !",dirname.Data());
		return kFALSE;
	    }
	}
    }
    fIn->cd();

    fMaxEntries = fTree->GetEntries();
    TString readCat = readCategories;
    readCat.ReplaceAll(" ","");

    TObjArray* catNames = readCat.Tokenize(",");
    Int_t n = catNames->GetEntries();

    //------------------------------------------------------------
    if(n == 0) {
	// no categories defined in input.
	// by default switch all on
	map<TString,HCategory*>::iterator iter;
	for( iter = fEvent.begin(); iter != fEvent.end(); ++iter ) {
	    addCatName(iter->first,getCategory(iter->first)->getCategory());
	}
    }
    //------------------------------------------------------------

    for(Int_t i = 0; i < n; i ++){
	TString name = ((TObjString*)catNames->At(i))->GetString();

	//------------------------------------------------------------
	if(i == 0){
	    if(name.CompareTo("-*") == 0){
		setStatus("-*", 0);
		continue;
	    }
	} else {
	    if(name.Contains("*") != 0){
		Warning("setInput()","Wild cards are only supported for '-*' in the first position of the string! Will ignore this one!");
		continue;
	    }
	}
	//------------------------------------------------------------

	//------------------------------------------------------------
	if(name.BeginsWith("+")) { // switch on
	    name.ReplaceAll("+","");
	    if(getCategory(name)) { // exists in input
		setStatus(name, 1);
                addCatName(name,getCategory(name)->getCategory());
	    }
	}
	//------------------------------------------------------------

	//------------------------------------------------------------
	if(name.BeginsWith("-")) { // switch off
	    name.ReplaceAll("-","");
	    if(getCategory(name)) { // exists in input
		setStatus(name, 0);
	    }
	}
	//------------------------------------------------------------

    }
    fIn->Close();

    printChain();

    return kTRUE;
}

void HLoop::printCategories()
{
    // print all categories found in input + status

    cout<<"--------------------------------------"<<endl;
    cout<<"current event :"<<endl;
    map<TString,HCategory*>::iterator iter;
    for( iter = fEvent.begin(); iter != fEvent.end(); ++iter ) {
	cout << setw(20)<<iter->first.Data() << " = " << iter->second <<" status = "<<fStatus[iter->first]<< endl;
    }
    cout<<"--------------------------------------"<<endl;
}

void HLoop::printChain()
{
    // print all files in the chain


    TObjArray* elements = fChain->GetListOfFiles();
    cout<<"--------------------------------------"<<endl;
    cout<<" Analize Chain : .... Can take while  for N file = "<<elements->GetEntries()<<endl;
    fMaxEntries = fChain->GetEntries();
    cout<<"N total entries = "<<fMaxEntries<<endl;
    for(Int_t i = 0; i < elements->GetEntries(); ++i ) {
        TChainElement* element = (TChainElement*)elements->At(i);
	cout <<setw(4)<<i<< setw(80)<<element->GetTitle() << ", N entries = " << element->GetEntries()<< endl;
    }
    cout<<"--------------------------------------"<<endl;
}

void HLoop::printBranchStatus()
{
    // print the branch status of the chain

    cout<<"--------------------------------------"<<endl;
    cout<<"Branch status :"<<endl;

    TObjArray* branches = fChain->GetListOfBranches();
    for(Int_t i = 0; i < branches ->GetEntries(); ++i ) {
        TBranch* br = (TBranch*) branches->At(i);
	cout <<setw(4)<<i<< setw(30)<<br->GetName() << ", status = " <<(Int_t) fChain->GetBranchStatus(br->GetName())<< endl;
    }
    cout<<"--------------------------------------"<<endl;
}

void HLoop::clearCategories()
{
    //  clear all active catgegories of the event

    map<TString,HCategory*>::iterator iter;
    for( iter = fEvent.begin(); iter != fEvent.end(); ++iter ) {
	if(fStatus[iter->first]) iter->second->Clear();
    }
}

HCategory* HLoop::getCategory(TString catname,Bool_t silent)
{
    // return the pointer to given category. NULL if not found

    map<TString,HCategory*>::iterator iter = fEvent.find(catname);
    if( iter == fEvent.end() ){
	if(!silent)Error("getCategory()","Category \"%s\" not found!",catname.Data());
	return 0;
    }
    else return iter->second;
}

Bool_t HLoop::getCategoryStatus(TString catname,Bool_t silent)
{
    // return the status of a given category. kFALSE if not found

    map<TString,Int_t>::iterator iter = fStatus.find(catname);
    if( iter == fStatus.end() ){
	if(!silent)Error("getCategory()","Category \"%s\" not found!",catname.Data());
	return kFALSE;
    } else {
	return (Bool_t)fStatus[catname];
    }
}

Bool_t HLoop::setStatus(TString catname, Int_t stat )
{
    // set the status of a given category. return kFALSE
    // if something went wrong. if catname="-*" or "+*"
    // all categories of the event are disabled /enabled.

    if(catname.CompareTo("-*") == 0 || catname.CompareTo("+*") == 0 ){
	map<TString,HCategory*>::iterator iter;
	for( iter = fEvent.begin(); iter != fEvent.end(); ++iter ) {
	    if(catname.CompareTo("-*") == 0)  {
		fChain->SetBranchStatus(Form("%s.",iter->first.Data()),0);
		fChain->SetBranchStatus(Form("%s.*",iter->first.Data()),0);
		fStatus[iter->first] = 0;
	    }
	    if(catname.CompareTo("+*") == 0)  {
		fChain->SetBranchStatus(Form("%s.",iter->first.Data()),1);
		fChain->SetBranchStatus(Form("%s.*",iter->first.Data()),1);
		fStatus[iter->first] = 1;
	    }
	}
	return kTRUE;
    }
    if(!getCategory(catname)) {
	return kFALSE;
    }

    map<TString,Int_t>::iterator iter = fStatus.find(catname);
    if( iter == fStatus.end() ){
	Error("getCategory()","Category \"%s\" not found!",catname.Data());
	return kFALSE;
    }
    else {
	fChain->SetBranchStatus(Form("%s.",catname.Data()),stat);
	fChain->SetBranchStatus(Form("%s.*",catname.Data()),stat);
	iter->second = stat;
	return kTRUE;
    }
    return kFALSE;
}

Int_t HLoop::nextEvent(Int_t iev)
{
    // get next entry form treee. clears all active categories
    // before. Returns number of read bytes
    // (0 = entry not found , -1 = IO error)
    // prints the current file name if a new file
    // is entered in the chain

    if(!fHead) {
	Warning("nextEvent()","File not initialized ! Call setInput(....) before!");
	exit(1);
    }

    fCurrentEntry = iev;

    clearCategories();

    Long64_t centry = fChain->LoadTree(fCurrentEntry);
    if(centry < 0 ){
	Warning("nextEvent()","LoadEntry %i < 0 !",(Int_t)fCurrentEntry);
    }

    Int_t nbytes = fChain->GetEntry(fCurrentEntry);  // 0 = entry not found , -1 = IO error
    if(nbytes == 0) {
	Warning("nextEvent()","Entry %i not found!",(Int_t)fCurrentEntry);
    } else if(nbytes == -1) {
	Warning("nextEvent()","Entry %i read with Io error!",(Int_t)fCurrentEntry);
    }

    TString name = fChain->GetFile()->GetName();
    if(name != fCurrentName) {
        Int_t entries = 0;
	TObjArray* elements = fChain->GetListOfFiles();
	for(Int_t i = 0; i < elements->GetEntries(); ++i ) {
	    TChainElement* element = (TChainElement*)elements->At(i);
	    if(name.CompareTo(element->GetTitle() ) == 0) {
		entries = element->GetEntries();
                break;
	    }
	}
        cout<<"---------------------------------------------------------"<<endl;
	cout<<name.Data()<<endl;
	cout<<"at entry "<<setw(8)<<fCurrentEntry<<" with entries = "<<entries<<endl;
        cout<<"---------------------------------------------------------"<<endl;
	fCurrentName = fChain->GetFile()->GetName() ;
	fFileCurrent = fChain->GetFile() ;
	fTree        = fChain->GetTree() ;
    }

    return nbytes;
}

Bool_t HLoop::addCatName(TString catname,Short_t catNum)
{
    // add agiven category number to the HADES event structure.
    // this function has to be called after setInput() !!!!

    if(fHead == 0) {
        Error("addCatName()","You have to call setInput() before calling this function !");
	exit(1);
    }

    map<TString,Short_t>::iterator iter = fNameToCat.find(catname);
    if( iter == fNameToCat.end() ){
	if(getCategory(catname,kTRUE)){
	   return fRecEvent->addCategory((Cat_t)catNum,getCategory(catname),fPartial[catname].Data());
       } else {
	   Error("addCatName()","Category \"%s\" does not exists in input! Will skip ...",catname.Data());
           return kFALSE;
       }
    } else {
	Error("addCatName()","Category \"%s\" already exists with number = %i!",catname.Data(),fNameToCat[catname]);
        return kFALSE;
    }
    return kTRUE;
}