#include <TTree.h>
#include <TFile.h>
#include <TString.h>
#include <TList.h>
#include <TROOT.h>
#include "AliLog.h"
#include "AliAODHandler.h"
#include "AliAODEvent.h"
#include "AliAODExtension.h"
#include "AliAODTracklets.h"
#include "AliStack.h"
#include "AliAODMCParticle.h"
#include "AliAODMCHeader.h"
#include "AliMCEventHandler.h"
#include "AliMCEvent.h"
#include "AliGenEventHeader.h"
#include "AliGenHijingEventHeader.h"
#include "AliGenDPMjetEventHeader.h"
#include "AliGenPythiaEventHeader.h"
#include "AliGenCocktailEventHeader.h"
#include "AliCodeTimer.h"
#include "AliAODBranchReplicator.h"
#include "Riostream.h"
using std::endl;
using std::cout;
ClassImp(AliAODHandler)
AliAODHandler::AliAODHandler() :
AliVEventHandler(),
fIsStandard(kTRUE),
fFillAOD(kTRUE),
fFillAODRun(kTRUE),
fFillExtension(kTRUE),
fNeedsHeaderReplication(kFALSE),
fNeedsTOFHeaderReplication(kFALSE),
fNeedsVZEROReplication(kFALSE),
fNeedsTracksBranchReplication(kFALSE),
fNeedsVerticesBranchReplication(kFALSE),
fNeedsV0sBranchReplication(kFALSE),
fNeedsCascadesBranchReplication(kFALSE),
fNeedsTrackletsBranchReplication(kFALSE),
fNeedsPMDClustersBranchReplication(kFALSE),
fNeedsJetsBranchReplication(kFALSE),
fNeedsFMDClustersBranchReplication(kFALSE),
fNeedsCaloClustersBranchReplication(kFALSE),
fNeedsCaloTriggerBranchReplication(kFALSE),
fNeedsMCParticlesBranchReplication(kFALSE),
fNeedsDimuonsBranchReplication(kFALSE),
fNeedsHMPIDBranchReplication(kFALSE),
fAODIsReplicated(kFALSE),
fTreeBuffSize(30000000),
fMemCountAOD(0),
fAODEvent(NULL),
fMCEventH(NULL),
fTreeA(NULL),
fFileA(NULL),
fFileName(""),
fExtensions(NULL),
fFilters(NULL)
{
}
AliAODHandler::AliAODHandler(const char* name, const char* title):
AliVEventHandler(name, title),
fIsStandard(kTRUE),
fFillAOD(kTRUE),
fFillAODRun(kTRUE),
fFillExtension(kTRUE),
fNeedsHeaderReplication(kFALSE),
fNeedsTOFHeaderReplication(kFALSE),
fNeedsVZEROReplication(kFALSE),
fNeedsTracksBranchReplication(kFALSE),
fNeedsVerticesBranchReplication(kFALSE),
fNeedsV0sBranchReplication(kFALSE),
fNeedsCascadesBranchReplication(kFALSE),
fNeedsTrackletsBranchReplication(kFALSE),
fNeedsPMDClustersBranchReplication(kFALSE),
fNeedsJetsBranchReplication(kFALSE),
fNeedsFMDClustersBranchReplication(kFALSE),
fNeedsCaloClustersBranchReplication(kFALSE),
fNeedsCaloTriggerBranchReplication(kFALSE),
fNeedsMCParticlesBranchReplication(kFALSE),
fNeedsDimuonsBranchReplication(kFALSE),
fNeedsHMPIDBranchReplication(kFALSE),
fAODIsReplicated(kFALSE),
fTreeBuffSize(30000000),
fMemCountAOD(0),
fAODEvent(NULL),
fMCEventH(NULL),
fTreeA(NULL),
fFileA(NULL),
fFileName(""),
fExtensions(NULL),
fFilters(NULL)
{
}
AliAODHandler::~AliAODHandler()
{
delete fAODEvent;
if (fFileA) fFileA->Close();
delete fFileA;
delete fTreeA;
delete fExtensions;
delete fFilters;
}
Bool_t AliAODHandler::Init(Option_t* opt)
{
Bool_t createStdAOD = fIsStandard || fFillAOD;
if(!fAODEvent && createStdAOD){
fAODEvent = new AliAODEvent();
if (fIsStandard)
fAODEvent->CreateStdContent();
}
TString option(opt);
option.ToLower();
if (createStdAOD) {
TDirectory *owd = gDirectory;
if (option.Contains("proof")) {
gROOT->ProcessLine(Form("AliAnalysisDataContainer *c_common_out = AliAnalysisManager::GetAnalysisManager()->GetCommonOutputContainer();"));
gROOT->ProcessLine(Form("AliAnalysisManager::GetAnalysisManager()->OpenProofFile(c_common_out, \"RECREATE\");"));
fFileA = gFile;
} else {
fFileA = new TFile(fFileName.Data(), "RECREATE");
}
CreateTree(1);
owd->cd();
}
if (fExtensions) {
TIter next(fExtensions);
AliAODExtension *ext;
while ((ext=(AliAODExtension*)next())) ext->Init(option);
}
if (fFilters) {
TIter nextf(fFilters);
AliAODExtension *filteredAOD;
while ((filteredAOD=(AliAODExtension*)nextf())) {
filteredAOD->SetEvent(fAODEvent);
filteredAOD->Init(option);
}
}
return kTRUE;
}
void AliAODHandler::Print(Option_t* opt) const
{
cout << opt << Form("IsStandard %d filename=%s",fIsStandard,fFileName.Data()) << endl;
if ( fExtensions )
{
cout << opt << fExtensions->GetEntries() << " extensions :" << endl;
PrintExtensions(*fExtensions);
}
if ( fFilters )
{
cout << opt << fFilters->GetEntries() << " filters :" << endl;
PrintExtensions(*fFilters);
}
}
void AliAODHandler::PrintExtensions(const TObjArray& array) const
{
TIter next(&array);
AliAODExtension* ext(0x0);
while ( ( ext = static_cast<AliAODExtension*>(next()) ) )
{
ext->Print(" ");
}
}
void AliAODHandler::StoreMCParticles(){
if (!fAODEvent) return;
TClonesArray *mcarray = (TClonesArray*)fAODEvent->FindListObject(AliAODMCParticle::StdBranchName());
if(!mcarray)return;
AliAODMCHeader *mcHeader = (AliAODMCHeader*)fAODEvent->FindListObject(AliAODMCHeader::StdBranchName());
if(!mcHeader)return;
if(!fMCEventH)return;
if(!fMCEventH->MCEvent())return;
AliStack *pStack = fMCEventH->MCEvent()->Stack();
if(!pStack)return;
fMCEventH->CreateLabelMap();
AliHeader* header = fMCEventH->MCEvent()->Header();
AliGenEventHeader* genHeader = 0;
if (header) genHeader = header->GenEventHeader();
if (genHeader) {
TArrayF vtxMC(3);
genHeader->PrimaryVertex(vtxMC);
mcHeader->SetVertex(vtxMC[0],vtxMC[1],vtxMC[2]);
AliGenCocktailEventHeader* genCocktailHeader = dynamic_cast<AliGenCocktailEventHeader*>(genHeader);
if(genCocktailHeader){
mcHeader->AddGeneratorName(genHeader->GetName());
TList* headerList = genCocktailHeader->GetHeaders();
AliGenEventHeader *headerEntry = dynamic_cast<AliGenEventHeader*>(headerList->At(0));
if (!headerEntry) {
AliFatal("AliGenEventHeader entry not found in the header list");
} else {
SetMCHeaderInfo(mcHeader,headerEntry);
}
}
else{
SetMCHeaderInfo(mcHeader,genHeader);
}
mcHeader->AddCocktailHeaders(genHeader);
}
AliMCEvent* mcEvent = fMCEventH->MCEvent();
Int_t np = mcEvent->GetNumberOfTracks();
Int_t nprim = mcEvent->GetNumberOfPrimaries();
Int_t j = 0;
TClonesArray& l = *mcarray;
for(int i = 0; i < np; ++i){
if(fMCEventH->IsParticleSelected(i)){
Int_t flag = 0;
AliMCParticle* mcpart = (AliMCParticle*) mcEvent->GetTrack(i);
if(i<nprim)flag |= AliAODMCParticle::kPrimary;
if(mcEvent->IsPhysicalPrimary(i))flag |= AliAODMCParticle::kPhysicalPrim;
if(mcEvent->IsSecondaryFromWeakDecay(i))flag |= AliAODMCParticle::kSecondaryFromWeakDecay;
if(mcEvent->IsSecondaryFromMaterial(i))flag |= AliAODMCParticle::kSecondaryFromMaterial;
if(fMCEventH->GetNewLabel(i)!=j){
AliError(Form("MISMATCH New label %d j: %d",fMCEventH->GetNewLabel(i),j));
}
AliAODMCParticle mcpartTmp(mcpart,i,flag);
mcpartTmp.SetStatus(mcpart->Particle()->GetStatusCode());
mcpartTmp.SetMCProcessCode(mcpart->Particle()->GetUniqueID());
Int_t d0 = mcpartTmp.GetDaughter(0);
Int_t d1 = mcpartTmp.GetDaughter(1);
Int_t m = mcpartTmp.GetMother();
if(d0<0 && d1<0){
mcpartTmp.SetDaughter(0,d0);
mcpartTmp.SetDaughter(1,d1);
} else if(d1 < 0 && d0 >= 0) {
if(fMCEventH->IsParticleSelected(d0)){
mcpartTmp.SetDaughter(0,fMCEventH->GetNewLabel(d0));
} else {
mcpartTmp.SetDaughter(0,-1);
}
mcpartTmp.SetDaughter(1,d1);
}
else if (d0 > 0 && d1 > 0 ){
Int_t d0Tmp = -1;
Int_t d1Tmp = -1;
for(int id = d0; id<=d1;++id){
if(fMCEventH->IsParticleSelected(id)){
if(d0Tmp==-1){
d0Tmp = fMCEventH->GetNewLabel(id);
d1Tmp = d0Tmp;
}
else d1Tmp = fMCEventH->GetNewLabel(id);
}
}
mcpartTmp.SetDaughter(0,d0Tmp);
mcpartTmp.SetDaughter(1,d1Tmp);
} else {
AliError(Form("Unxpected indices %d %d",d0,d1));
}
if(m<0){
mcpartTmp.SetMother(m);
} else {
if(fMCEventH->IsParticleSelected(m))mcpartTmp.SetMother(fMCEventH->GetNewLabel(m));
else AliError(Form("PROBLEM Mother not selected %d \n", m));
}
new (l[j++]) AliAODMCParticle(mcpartTmp);
}
}
AliInfo(Form("AliAODHandler::StoreMCParticles: Selected %d (Primaries %d / total %d) after validation \n",
j,nprim,np));
TClonesArray* tracks = fAODEvent->GetTracks();
Int_t tofLabel[3];
if(tracks){
for(int it = 0; it < fAODEvent->GetNumberOfTracks();++it){
AliAODTrack *track = dynamic_cast<AliAODTrack*>(fAODEvent->GetTrack(it));
if(!track) AliFatal("Not a standard AOD");
Int_t sign = 1;
Int_t label = track->GetLabel();
if(label<0){
label *= -1;
sign = -1;
}
if (label >= AliMCEvent::BgLabelOffset()) label = mcEvent->BgLabelToIndex(label);
if(label > np || track->GetLabel() == 0){
AliWarning(Form("Wrong ESD track label %5d (%5d)",track->GetLabel(), label));
}
if(fMCEventH->GetNewLabel(label) == 0) {
AliWarning(Form("New label not found for %5d (%5d)",track->GetLabel(), label));
}
track->SetLabel(sign*fMCEventH->GetNewLabel(label));
track->GetTOFLabel(tofLabel);
for (Int_t i = 0; i < 3; i++) {
label = tofLabel[i];
Int_t nlabel = label;
if (label < 0) continue;
if (label >= AliMCEvent::BgLabelOffset()) nlabel = mcEvent->BgLabelToIndex(label);
if(nlabel > np || label == 0) {
AliWarning(Form("Wrong TOF label %5d (%5d)", label, nlabel));
}
if(fMCEventH->GetNewLabel(label) == 0){
AliWarning(Form("New TOF label not found for %5d %5d",i, label ));
tofLabel[i] = -label;
} else {
tofLabel[i] = fMCEventH->GetNewLabel(label);
}
}
track->SetTOFLabel(tofLabel);
}
}
TClonesArray *clusters = fAODEvent->GetCaloClusters();
if(clusters){
for (Int_t iClust = 0;iClust < fAODEvent->GetNumberOfCaloClusters(); ++iClust) {
AliAODCaloCluster * cluster = fAODEvent->GetCaloCluster(iClust);
UInt_t nLabel = cluster->GetNLabels();
Int_t* labels = const_cast<Int_t*>(cluster->GetLabels());
if (labels){
for(UInt_t i = 0;i < nLabel;++i){
labels[i] = fMCEventH->GetNewLabel(cluster->GetLabelAt(i));
}
}
}
}
Int_t iCell, nCell, cellMCLabel, cellMCLabelNew;;
Short_t cellAbsId;
Double_t cellE, cellT, cellEFrac;
AliAODCaloCells *cells;
cells = fAODEvent->GetEMCALCells();
if( cells ){
nCell = cells->GetNumberOfCells() ;
for( iCell = 0; iCell < nCell; iCell++ ){
cells->GetCell( iCell, cellAbsId, cellE, cellT, cellMCLabel, cellEFrac );
if( cellMCLabel < 0 )
cellMCLabelNew = cellMCLabel;
else
cellMCLabelNew = fMCEventH->GetNewLabel( cellMCLabel );
cells->SetCell( iCell, cellAbsId, cellE, cellT, cellMCLabelNew, cellEFrac );
}
}
cells = fAODEvent->GetPHOSCells();
if( cells ){
nCell = cells->GetNumberOfCells() ;
for( iCell = 0; iCell < nCell; iCell++ ){
cells->GetCell( iCell, cellAbsId, cellE, cellT, cellMCLabel, cellEFrac );
if( cellMCLabel < 0 )
cellMCLabelNew = cellMCLabel;
else
cellMCLabelNew = fMCEventH->GetNewLabel( cellMCLabel );
cells->SetCell( iCell, cellAbsId, cellE, cellT, cellMCLabelNew, cellEFrac );
}
}
AliAODTracklets *tracklets = fAODEvent->GetTracklets();
if(tracklets){
for(int it = 0;it < tracklets->GetNumberOfTracklets();++it){
int label0 = tracklets->GetLabel(it,0);
int label1 = tracklets->GetLabel(it,1);
if(label0>=0)label0 = fMCEventH->GetNewLabel(label0);
if(label1>=0)label1 = fMCEventH->GetNewLabel(label1);
tracklets->SetLabel(it,0,label0);
tracklets->SetLabel(it,1,label1);
}
}
}
Bool_t AliAODHandler::FinishEvent()
{
if(fFillAOD && fFillAODRun && fAODEvent){
fAODEvent->MakeEntriesReferencable();
fTreeA->BranchRef();
FillTree();
}
if ((fFillAOD && fFillAODRun) || fFillExtension) {
if (fExtensions && fFillExtension) {
TIter next(fExtensions);
AliAODExtension *ext;
while ((ext=(AliAODExtension*)next())) ext->FinishEvent();
}
if (fFilters && fFillAOD && fFillAODRun) {
TIter nextf(fFilters);
AliAODExtension *ext;
while ((ext=(AliAODExtension*)nextf())) {
ext->FinishEvent();
}
}
}
if (fIsStandard && fAODEvent)
{
fAODEvent->ResetStd();
}
if (fAODEvent)
{
TClonesArray *mcarray = static_cast<TClonesArray*>(fAODEvent->FindListObject(AliAODMCParticle::StdBranchName()));
if(mcarray) mcarray->Delete();
AliAODMCHeader *mcHeader = static_cast<AliAODMCHeader*>(fAODEvent->FindListObject(AliAODMCHeader::StdBranchName()));
if(mcHeader) mcHeader->Reset();
}
fAODIsReplicated = kFALSE;
return kTRUE;
}
Bool_t AliAODHandler::Terminate()
{
AddAODtoTreeUserInfo();
TIter nextF(fFilters);
AliAODExtension *ext;
while ((ext=static_cast<AliAODExtension*>(nextF())))
{
ext->AddAODtoTreeUserInfo();
}
TIter nextE(fExtensions);
while ((ext=static_cast<AliAODExtension*>(nextE())))
{
ext->AddAODtoTreeUserInfo();
}
return kTRUE;
}
Bool_t AliAODHandler::TerminateIO()
{
if (fFileA) {
fFileA->Write();
fFileA->Close();
delete fFileA;
fFileA = 0;
fTreeA = 0;
}
TIter nextF(fFilters);
AliAODExtension *ext;
while ((ext=static_cast<AliAODExtension*>(nextF())))
{
ext->TerminateIO();
}
TIter nextE(fExtensions);
while ((ext=static_cast<AliAODExtension*>(nextE())))
{
ext->TerminateIO();
}
return kTRUE;
}
void AliAODHandler::CreateTree(Int_t flag)
{
fTreeA = new TTree("aodTree", "AliAOD tree");
fTreeA->Branch(fAODEvent->GetList());
if (flag == 0) fTreeA->SetDirectory(0);
fMemCountAOD = 0;
}
void AliAODHandler::FillTree()
{
Long64_t nbf = fTreeA->Fill();
if (fTreeBuffSize>0 && fTreeA->GetAutoFlush()<0 && (fMemCountAOD += nbf)>fTreeBuffSize ) {
nbf = fTreeA->GetZipBytes();
if (nbf>0) nbf = -nbf;
else nbf = fTreeA->GetEntries();
fTreeA->SetAutoFlush(nbf);
AliInfo(Form("Calling fTreeA->SetAutoFlush(%lld) | W:%lld T:%lld Z:%lld",
nbf,fMemCountAOD,fTreeA->GetTotBytes(),fTreeA->GetZipBytes()));
}
}
void AliAODHandler::AddAODtoTreeUserInfo()
{
if (fTreeA) fTreeA->GetUserInfo()->Add(fAODEvent);
fAODEvent = 0;
}
void AliAODHandler::AddBranch(const char* cname, void* addobj, const char* filename)
{
if (strlen(filename))
{
AliAODExtension *ext = AddExtension(filename);
ext->AddBranch(cname, addobj);
return;
}
if (fFilters) {
TIter next(fFilters);
AliAODExtension *ext;
while ((ext=(AliAODExtension*)next())) ext->AddBranch(cname, addobj);
}
TDirectory *owd = gDirectory;
if (fFileA)
{
fFileA->cd();
}
char** apointer = (char**) addobj;
TObject* obj = (TObject*) *apointer;
fAODEvent->AddObject(obj);
const Int_t kSplitlevel = 99;
const Int_t kBufsize = 32000;
if (!fTreeA->FindBranch(obj->GetName()))
{
fTreeA->Bronch(obj->GetName(), cname, fAODEvent->GetList()->GetObjectRef(obj),
kBufsize, kSplitlevel - 1);
}
owd->cd();
}
AliAODExtension *AliAODHandler::AddExtension(const char *filename, const char *title, Bool_t tomerge)
{
TString fname(filename);
if (!fname.EndsWith(".root")) fname += ".root";
if (!fExtensions) {
fExtensions = new TObjArray();
fExtensions->SetOwner();
}
AliAODExtension *ext = (AliAODExtension*)fExtensions->FindObject(fname);
if (!ext) {
ext = new AliAODExtension(fname, title);
fExtensions->Add(ext);
}
ext->SetToMerge(tomerge);
return ext;
}
AliAODExtension *AliAODHandler::GetExtension(const char *filename) const
{
if (!fExtensions) return NULL;
return (AliAODExtension*)fExtensions->FindObject(filename);
}
AliAODExtension *AliAODHandler::AddFilteredAOD(const char *filename, const char *filtername, Bool_t tomerge)
{
if (!fFilters) {
fFilters = new TObjArray();
fFilters->SetOwner();
}
AliAODExtension *filter = (AliAODExtension*)fFilters->FindObject(filename);
if (!filter) {
filter = new AliAODExtension(filename, filtername, kTRUE);
fFilters->Add(filter);
}
filter->SetToMerge(tomerge);
return filter;
}
AliAODExtension *AliAODHandler::GetFilteredAOD(const char *filename) const
{
if (!fFilters) return NULL;
return (AliAODExtension*)fFilters->FindObject(filename);
}
void AliAODHandler::SetOutputFileName(const char* fname)
{
fFileName = fname;
}
const char *AliAODHandler::GetOutputFileName() const
{
return fFileName.Data();
}
const char *AliAODHandler::GetExtraOutputs(Bool_t merge) const
{
static TString eoutputs;
eoutputs = "";
AliAODExtension *obj;
if (fExtensions) {
TIter next1(fExtensions);
while ((obj=(AliAODExtension*)next1())) {
if (merge && !obj->IsToMerge()) continue;
if (!eoutputs.IsNull()) eoutputs += ",";
eoutputs += obj->GetName();
}
}
if (fFilters) {
TIter next2(fFilters);
while ((obj=(AliAODExtension*)next2())) {
if (merge && !obj->IsToMerge()) continue;
if (!eoutputs.IsNull()) eoutputs += ",";
eoutputs += obj->GetName();
}
}
return eoutputs.Data();
}
Bool_t AliAODHandler::HasExtensions() const
{
if ( fExtensions && fExtensions->GetEntries()>0 ) return kTRUE;
return kFALSE;
}
void AliAODHandler::SetMCHeaderInfo(AliAODMCHeader *mcHeader,AliGenEventHeader *genHeader){
if(!genHeader)return;
AliGenPythiaEventHeader *pythiaGenHeader = dynamic_cast<AliGenPythiaEventHeader*>(genHeader);
if (pythiaGenHeader) {
mcHeader->SetEventType(pythiaGenHeader->ProcessType());
mcHeader->SetPtHard(pythiaGenHeader->GetPtHard());
return;
}
AliGenDPMjetEventHeader* dpmJetGenHeader = dynamic_cast<AliGenDPMjetEventHeader*>(genHeader);
if (dpmJetGenHeader){
mcHeader->SetEventType(dpmJetGenHeader->ProcessType());
return;
}
AliGenHijingEventHeader* hijingGenHeader = dynamic_cast<AliGenHijingEventHeader*>(genHeader);
if(hijingGenHeader){
mcHeader->SetImpactParameter(hijingGenHeader->ImpactParameter());
return;
}
}