#include "AliConversionSelection.h"
#include "AliAODHeader.h"
#include "AliVVZERO.h"
#include "AliMultiplicity.h"
#include <iostream>


// Author Daniel Lohner (Daniel.Lohner@cern.ch)

using namespace std;

ClassImp(AliConversionSelection)


//________________________________________________________________________
AliConversionSelection::AliConversionSelection(AliConvEventCuts *evtCut, AliConversionPhotonCuts *convCut, AliConversionMesonCuts *mesonCut) : TObject(),
	fInputEvent(NULL),
	fMCEvent(NULL),
	fEventCut(evtCut),
	fConversionCut(convCut),
	fMesonCut(mesonCut),
	fESDTrackCuts(NULL),
	fGoodGammas(NULL),
	fPi0Candidates(NULL),
	fBGPi0s(NULL),
	fRandomizer(NULL),
	fBGHandler(NULL),
	fCurrentEventNumber(-1),
	fIsOwner(kFALSE)
{
	// Default Values
	fInvMassRange[0]=0.05;
	fInvMassRange[1]=0.3;
}

//________________________________________________________________________
AliConversionSelection::AliConversionSelection(TString evtCut, TString convCut, TString mesonCut) : TObject(),
	fInputEvent(NULL),
	fMCEvent(NULL),
	fEventCut(NULL),
	fConversionCut(NULL),
	fMesonCut(NULL),
	fESDTrackCuts(NULL),
	fGoodGammas(NULL),
	fPi0Candidates(NULL),
	fBGPi0s(NULL),
	fRandomizer(NULL),
	fBGHandler(NULL),
	fCurrentEventNumber(-1),
	fIsOwner(kTRUE)
{
	// Default Values
	fInvMassRange[0]=0.05;
	fInvMassRange[1]=0.3;

	fEventCut = new AliConvEventCuts();
	fEventCut -> InitializeCutsFromCutString(evtCut.Data());
	fConversionCut = new AliConversionPhotonCuts();
	fConversionCut -> InitializeCutsFromCutString(convCut.Data());
	fMesonCut = new AliConversionMesonCuts();
	fMesonCut -> InitializeCutsFromCutString(mesonCut.Data());

}


//________________________________________________________________________
AliConversionSelection::AliConversionSelection(const AliConversionSelection &ref) : TObject(ref),
	fInputEvent(NULL),
	fMCEvent(NULL),
	fEventCut(NULL),
	fConversionCut(NULL),
	fMesonCut(NULL),
	fESDTrackCuts(NULL),
	fGoodGammas(NULL),
	fPi0Candidates(NULL),
	fBGPi0s(NULL),
	fRandomizer(NULL),
	fBGHandler(NULL),
	fCurrentEventNumber(-1),
	fIsOwner(kTRUE)
{
	// Copy Constructor
	fEventCut=new AliConvEventCuts(*ref.fEventCut);
	fConversionCut=new AliConversionPhotonCuts(*ref.fConversionCut);
	fMesonCut=new AliConversionMesonCuts(*ref.fMesonCut);

	fInvMassRange[0]=ref.fInvMassRange[0];
	fInvMassRange[1]=ref.fInvMassRange[1];
}

//________________________________________________________________________
AliConversionSelection::~AliConversionSelection(){

	if(fBGHandler){
		delete fBGHandler;
		fBGHandler=NULL;
	}
	if(fRandomizer){
		delete fRandomizer;
		fRandomizer=NULL;
	}
	if(fPi0Candidates){
		delete fPi0Candidates;
		fPi0Candidates=NULL;
	}
	if(fBGPi0s){
		delete fBGPi0s;
		fBGPi0s=NULL;
	}
	if(fESDTrackCuts){
		delete fESDTrackCuts;
		fESDTrackCuts=NULL;
	}
	if(fIsOwner){
		if(fEventCut){
			delete fEventCut;
			fEventCut=NULL;
		}
		if(fConversionCut){
			delete fConversionCut;
			fConversionCut=NULL;
		}
		if(fMesonCut){
			delete fMesonCut;
			fMesonCut=NULL;
		}
	}
}

//________________________________________________________________________
Bool_t AliConversionSelection::ProcessEvent(TClonesArray *photons,AliVEvent *inputEvent,AliMCEvent *mcEvent){
	fInputEvent=inputEvent;
	fMCEvent=mcEvent;

	// Protection
	Int_t eventnumber=GetEventNumber(inputEvent);
	if(eventnumber==fCurrentEventNumber){
		AliWarning("Event already analyzed! Return.");
		return kFALSE;
	}
	else{
		fCurrentEventNumber=eventnumber;
	}

	// Initialize and Reset Arrays
	if(fGoodGammas == NULL){
		fGoodGammas=new TObjArray(30);
	}
	fGoodGammas->Clear();

	if(fPi0Candidates == NULL){
		fPi0Candidates = new TClonesArray("AliAODConversionMother",100);
	}
	fPi0Candidates->Delete();

	if(fBGPi0s == NULL){
		fBGPi0s = new TClonesArray("AliAODConversionMother",100);
	}
	fBGPi0s->Delete();


	if(!photons||!fInputEvent)return kFALSE;
	
	if(!fEventCut->EventIsSelected(fInputEvent,fMCEvent))return kFALSE;

	// Select photons
	for(Int_t i = 0; i < photons->GetEntriesFast(); i++) {
		AliAODConversionPhoton* gamma =dynamic_cast<AliAODConversionPhoton*>(photons->At(i));
		if(!gamma) continue;
		if(!fConversionCut->PhotonIsSelected(gamma,fInputEvent))continue;
		fGoodGammas->Add(gamma);
	}

	// Do MC Smearing
	Double_t *fUnsmearedPx=NULL;
	Double_t *fUnsmearedPy=NULL;
	Double_t *fUnsmearedPz=NULL;
	Double_t *fUnsmearedE=NULL;

	if(fMesonCut->UseMCPSmearing() && fMCEvent){
		fUnsmearedPx = new Double_t[fGoodGammas->GetEntries()]; // Store unsmeared Momenta
		fUnsmearedPy = new Double_t[fGoodGammas->GetEntries()];
		fUnsmearedPz = new Double_t[fGoodGammas->GetEntries()];
		fUnsmearedE =  new Double_t[fGoodGammas->GetEntries()];

		for(Int_t gamma=0;gamma<fGoodGammas->GetEntries();gamma++){ // Smear the AODPhotons in MC
			fUnsmearedPx[gamma] = ((AliAODConversionPhoton*)fGoodGammas->At(gamma))->Px();
			fUnsmearedPy[gamma] = ((AliAODConversionPhoton*)fGoodGammas->At(gamma))->Py();
			fUnsmearedPz[gamma] = ((AliAODConversionPhoton*)fGoodGammas->At(gamma))->Pz();
			fUnsmearedE[gamma] =  ((AliAODConversionPhoton*)fGoodGammas->At(gamma))->E();
			fMesonCut->SmearParticle(dynamic_cast<AliAODConversionPhoton*>(fGoodGammas->At(gamma)));
		}
	}

	// Reconstruct Pi0 and BG
	CalculatePi0Candidates();
	CalculateBackground();
	if(fBGHandler)fBGHandler->AddEvent(fGoodGammas,fInputEvent);

	// Undo MC Smearing
	if(fMesonCut->UseMCPSmearing() && fMCEvent){
		for(Int_t gamma=0;gamma<fGoodGammas->GetEntries();gamma++){ // Smear the AODPhotons in MC
			((AliAODConversionPhoton*)fGoodGammas->At(gamma))->SetPx(fUnsmearedPx[gamma]); // Reset Unsmeared Momenta
			((AliAODConversionPhoton*)fGoodGammas->At(gamma))->SetPy(fUnsmearedPy[gamma]);
			((AliAODConversionPhoton*)fGoodGammas->At(gamma))->SetPz(fUnsmearedPz[gamma]);
			((AliAODConversionPhoton*)fGoodGammas->At(gamma))->SetE(fUnsmearedE[gamma]);
		}
		delete[] fUnsmearedPx; fUnsmearedPx = 0x0;
		delete[] fUnsmearedPy; fUnsmearedPy = 0x0;
		delete[] fUnsmearedPz; fUnsmearedPz = 0x0;
		delete[] fUnsmearedE;  fUnsmearedE  = 0x0;
	}
	return kTRUE;
}

//________________________________________________________________________
AliAODConversionMother* AliConversionSelection::GetPi0(Int_t index){

	if(index>=0&&index<GetNumberOfPi0s()){
		return dynamic_cast<AliAODConversionMother*>(fPi0Candidates->At(index));
	}
	return NULL;
}

//________________________________________________________________________
AliAODConversionMother* AliConversionSelection::GetBG(Int_t index){

	if(index>=0&&index<GetNumberOfBGs()){
		return dynamic_cast<AliAODConversionMother*>(fBGPi0s->At(index));
	}
	return NULL;
}

//________________________________________________________________________
AliAODConversionPhoton* AliConversionSelection::GetPhoton(Int_t index){

	if(index>=0&&index<GetNumberOfPhotons()){
		return dynamic_cast<AliAODConversionPhoton*>(fGoodGammas->At(index));
	}
	return NULL;
}

//________________________________________________________________________
void AliConversionSelection::CalculatePi0Candidates(){
	// Conversion Gammas
	if(fGoodGammas->GetEntriesFast()>1){
		for(Int_t firstGammaIndex=0;firstGammaIndex<fGoodGammas->GetEntriesFast()-1;firstGammaIndex++){
			AliAODConversionPhoton *gamma0=dynamic_cast<AliAODConversionPhoton*>(fGoodGammas->At(firstGammaIndex));
			if (gamma0==NULL) continue;
			// Combine Photons

			for(Int_t secondGammaIndex=firstGammaIndex+1;secondGammaIndex<fGoodGammas->GetEntriesFast();secondGammaIndex++){

				AliAODConversionPhoton *gamma1=dynamic_cast<AliAODConversionPhoton*>(fGoodGammas->At(secondGammaIndex));
				if (gamma1==NULL) continue;
				//Check for same Electron ID
				if(gamma0->GetTrackLabelPositive()==gamma1->GetTrackLabelPositive()||gamma0->GetTrackLabelNegative()==gamma1->GetTrackLabelNegative()
				||gamma0->GetTrackLabelNegative()==gamma1->GetTrackLabelPositive()||gamma0->GetTrackLabelPositive()==gamma1->GetTrackLabelNegative())continue;

				AliAODConversionMother pi0cand(gamma0,gamma1);
				pi0cand.SetLabels(firstGammaIndex,secondGammaIndex);

				// Set MC Label

				if(fMCEvent){

					TParticle *mcgam0=gamma0->GetMCParticle(fMCEvent->Stack());
					TParticle *mcgam1=gamma1->GetMCParticle(fMCEvent->Stack());

					if(mcgam0&&mcgam1){
					// Have same Mother?

					if(mcgam0->GetMother(0)==mcgam1->GetMother(0)){

						pi0cand.SetMCLabel(mcgam0->GetMother(0));
					}
					}
				}

				if((fMesonCut->MesonIsSelected(&pi0cand))){
					if(MesonInMassWindow(&pi0cand)){

					// Add Pi0 to Stack
					new((*fPi0Candidates)[fPi0Candidates->GetEntriesFast()]) AliAODConversionMother(pi0cand);
					}
				}
			}
		}
	}
}

//________________________________________________________________________
Bool_t AliConversionSelection::MesonInMassWindow(AliAODConversionMother *pi0cand)
{
	if (pi0cand->M() > fInvMassRange[0] && pi0cand->M() < fInvMassRange[1] ){
		return kTRUE;
	}
	return kFALSE;
}

//________________________________________________________________________
void AliConversionSelection::RotateParticle(AliAODConversionPhoton *gamma,Int_t nDegreesPMBackground){

	if(!fRandomizer){
		fRandomizer=new TRandom3();
		fRandomizer->SetSeed(0);
	}
	Double_t nRadiansPM = nDegreesPMBackground*TMath::Pi()/180;
	Double_t rotationValue = fRandomizer->Rndm()*2*nRadiansPM + TMath::Pi()-nRadiansPM;
	gamma->RotateZ(rotationValue);
}

//________________________________________________________________________

void AliConversionSelection::CalculateBackground(){

	//Rotation Method
	if(fMesonCut->UseRotationMethod()){
		// Correct for the number of rotations
		// BG is for rotation the same, except for factor NRotations
		Double_t weight=1./Double_t(fMesonCut->GetNumberOfBGEvents());
		for(Int_t firstGammaIndex=0;firstGammaIndex<fGoodGammas->GetEntriesFast();firstGammaIndex++){
			AliAODConversionPhoton *gamma0=dynamic_cast<AliAODConversionPhoton*>(fGoodGammas->At(firstGammaIndex));
			if (gamma0 ==NULL) continue;
			for(Int_t secondGammaIndex=firstGammaIndex+1;secondGammaIndex<fGoodGammas->GetEntriesFast();secondGammaIndex++){
				AliAODConversionPhoton *gamma1=dynamic_cast<AliAODConversionPhoton*>(fGoodGammas->At(secondGammaIndex));
				if (gamma1==NULL) continue;
				if(!fConversionCut->PhotonIsSelected(gamma1,fInputEvent))continue;
				for(Int_t nRandom=0;nRandom<fMesonCut->GetNumberOfBGEvents();nRandom++){
					RotateParticle(gamma1,fMesonCut->NDegreesRotation());
					AliAODConversionMother BGcandidate(gamma0,gamma1);
					if(fMesonCut->MesonIsSelected(&BGcandidate,kFALSE)){
						if(MesonInMassWindow(&BGcandidate)){
							new((*fBGPi0s)[fBGPi0s->GetEntriesFast()]) AliAODConversionMother(BGcandidate);
							dynamic_cast<AliAODConversionMother*>(fBGPi0s->At(fBGPi0s->GetEntriesFast()-1))->SetWeight(weight);
						}
					}
				}
			}
		}
	} else {
		// Do Event Mixing
		if(fBGHandler==NULL){
			fBGHandler=new AliConversionAODBGHandlerRP(fEventCut->IsHeavyIon(),fMesonCut->UseTrackMultiplicity());
		}

		for(Int_t nEventsInBG=0;nEventsInBG <fBGHandler->GetNBGEvents(fGoodGammas,fInputEvent);nEventsInBG++){
			AliGammaConversionPhotonVector *previousEventGammas = fBGHandler->GetBGGoodGammas(fGoodGammas,fInputEvent,nEventsInBG);
			if(previousEventGammas){
				// test weighted background
				Double_t weight=1.0;
				// Correct for the number of eventmixing:
				// N gammas -> (N-1) + (N-2) +(N-3) ...+ (N-(N-1))  using sum formula sum(i)=N*(N-1)/2  -> N*(N-1)/2
				// real combinations (since you cannot combine a photon with its own)
				// but BG leads to N_{a}*N_{b} combinations
				weight*=0.5*(Double_t(fGoodGammas->GetEntriesFast()-1))/Double_t(previousEventGammas->size());
				for(Int_t iCurrent=0;iCurrent<fGoodGammas->GetEntriesFast();iCurrent++){
					AliAODConversionPhoton *gamma0 = (AliAODConversionPhoton*)(fGoodGammas->At(iCurrent));
					for(UInt_t iPrevious=0;iPrevious<previousEventGammas->size();iPrevious++){
						AliAODConversionPhoton *gamma1 = (AliAODConversionPhoton*)(previousEventGammas->at(iPrevious));
						AliAODConversionMother BGcandidate(gamma0,gamma1);
						if(fMesonCut->MesonIsSelected(&BGcandidate,kFALSE)){
							if(MesonInMassWindow(&BGcandidate)){
								new((*fBGPi0s)[fBGPi0s->GetEntriesFast()]) AliAODConversionMother(BGcandidate);
								dynamic_cast<AliAODConversionMother*>(fBGPi0s->At(fBGPi0s->GetEntriesFast()-1))->SetWeight(weight);
							}
						}
					}
				}
			}
		}
    }
}
//________________________________________________________________________
Double_t AliConversionSelection::GetMultiplicity(AliVEvent *inputEvent){

    switch(fEventCut->GetMultiplicityMethod()){
		case 0:
			return Double_t(GetNumberOfPhotons());
		case 1:
			return Double_t(GetNumberOfChargedTracks(inputEvent));
		case 2:
			return GetVZEROMult(inputEvent);
		case 3:
			return GetSPDMult(inputEvent);
		case 9:
			return 1; // if mult is used as a weight, this number can be used to switch off weighting
		default:
		return 0;
    }
}

//________________________________________________________________________
Int_t AliConversionSelection::GetNumberOfChargedTracks(AliVEvent *inputEvent){

	Int_t ntracks = 0;

	AliESDEvent * esdEvent = dynamic_cast<AliESDEvent*>(inputEvent);
	if(esdEvent) {
		if(!fESDTrackCuts){
			fESDTrackCuts= AliESDtrackCuts::GetStandardITSTPCTrackCuts2010(kTRUE);
			fESDTrackCuts->SetMaxDCAToVertexZ(2);
			Double_t etamax=fConversionCut->GetEtaCut();
			fESDTrackCuts->SetEtaRange(-etamax, etamax);
			fESDTrackCuts->SetPtRange(0.15);
		}
		for(Int_t iTracks = 0; iTracks < inputEvent->GetNumberOfTracks(); iTracks++){
			AliESDtrack* currentTrack = esdEvent->GetTrack(iTracks);
			if(!currentTrack) continue;
			if(fESDTrackCuts->AcceptTrack(currentTrack))ntracks++;
		}
	} else {
		for(Int_t ii=0; ii<inputEvent->GetNumberOfTracks(); ii++) {
			AliVTrack * track = dynamic_cast<AliVTrack*>(inputEvent->GetTrack(ii));
			if (track==NULL) continue;
			if(TMath::Abs(track->Eta())>fConversionCut->GetEtaCut())continue;
			if(track)ntracks++;
		}
	}
	return ntracks;
}

//________________________________________________________________________
Double_t AliConversionSelection::GetVZEROMult(AliVEvent *inputEvent){

	AliVVZERO *vzero=inputEvent->GetVZEROData();
	Double_t multV0A=vzero->GetMTotV0A();
	Double_t multV0C=vzero->GetMTotV0C();
	Double_t mult=multV0A+multV0C;

	return mult;
}

//________________________________________________________________________
Double_t AliConversionSelection::GetSPDMult(AliVEvent *inputEvent){

	AliESDEvent * esdEvent = dynamic_cast<AliESDEvent*>(inputEvent);
	if(esdEvent) {
		const AliMultiplicity *esdmult=esdEvent->GetMultiplicity();
		return esdmult->GetNumberOfITSClusters(1);
	} else {
		// AOD implementation
		AliAODHeader *header=(AliAODHeader*)inputEvent->GetHeader();
		return header->GetNumberOfITSClusters(1);
	}
	return 0;
}

//________________________________________________________________________
Int_t AliConversionSelection::GetEventNumber(AliVEvent *inputEvent){

	AliESDEvent * esdEvent = dynamic_cast<AliESDEvent*>(inputEvent);
	if(esdEvent) {
		return esdEvent->GetEventNumberInFile();
	} else{
		AliAODHeader *header=(AliAODHeader*)inputEvent->GetHeader();
		return header->GetEventNumberESDFile();
	}
	return 0;
}

//________________________________________________________________________
TString AliConversionSelection::GetCutString(){
	TString a= Form("%s_%s_%s",fEventCut->GetCutNumber().Data(), fConversionCut->GetCutNumber().Data(),fMesonCut->GetCutNumber().Data());
	return a;
}