#include "TChain.h"
#include "TList.h"
#include "TCanvas.h"
#include "TLorentzVector.h"
#include "TGraphErrors.h"
#include "TH1F.h"
#include "TH2F.h"

#include "AliAnalysisTaskSE.h"
#include "AliAnalysisManager.h"

#include "AliMCParticle.h"
#include "AliMCEventHandler.h"
#include "AliVEvent.h"
#include "AliMCEvent.h"
#include "AliVVertex.h"

#include "AliAnalysisTaskMC.h"


// Analysis task for the BF code
// Authors: Panos Cristakoglou@cern.ch

ClassImp(AliAnalysisTaskMC)

//________________________________________________________________________
AliAnalysisTaskMC::AliAnalysisTaskMC(const char *name) 
: AliAnalysisTaskSE(name), 
  fList(0),
  fHistEventStats(0),
  fHistTrackStats(0),
  fHistVx(0),
  fHistVy(0),
  fHistVz(0),
  fHistMultiplicity(0),
  fVxMax(0.3),
  fVyMax(0.3),
  fVzMax(10.),
  fPtMin(0.1),
  fPtMax(10.5),
  fEtaMin(-0.8),
  fEtaMax(-0.8) {
// Constructor

  // Define input and output slots here
  // Input slot #0 works with a TChain
  DefineInput(0, TChain::Class());
  // Output slot #0 writes into a TH1 container
  DefineOutput(1, TList::Class());
}

//________________________________________________________________________
AliAnalysisTaskMC::~AliAnalysisTaskMC() {
  //
}

//________________________________________________________________________
void AliAnalysisTaskMC::UserCreateOutputObjects() {
  // Create histograms
  // Called once
  //QA list
  fList = new TList();
  fList->SetName("listQA");
  fList->SetOwner();

  //Event stats.
  TString gCutName[4] = {"Total","Offline trigger",
                         "Vertex","Analyzed"};
  fHistEventStats = new TH1F("fHistEventStats",
                             "Event statistics;;N_{events}",
                             4,0.5,4.5);
  for(Int_t i = 1; i <= 4; i++)
    fHistEventStats->GetXaxis()->SetBinLabel(i,gCutName[i-1].Data());
  fList->Add(fHistEventStats);

  fHistTrackStats = new TH1F("fHistTrackStats","Event statistics;TriggerBit;N_{events}",130,0,130);
  fList->Add(fHistTrackStats);

  // Vertex distributions
  fHistVx = new TH1F("fHistVx","Primary vertex distribution - x coordinate;V_{x} (cm);Entries",100,-0.5,0.5);
  fList->Add(fHistVx);
  fHistVy = new TH1F("fHistVy","Primary vertex distribution - y coordinate;V_{y} (cm);Entries",100,-0.5,0.5);
  fList->Add(fHistVy);
  fHistVz = new TH1F("fHistVz","Primary vertex distribution - z coordinate;V_{z} (cm);Entries",100,-20.,20.);
  fList->Add(fHistVz);

  // QA histograms
  fHistMultiplicity = new TH1F("fHistMultiplicity",";N_{acc.};Counts",500,-0.5,499.5);
  fList->Add(fHistMultiplicity);

  // Post output data.
  PostData(1, fList);
  
}

//________________________________________________________________________
void AliAnalysisTaskMC::UserExec(Option_t *) {
  // Main loop
  // Called for each event
  //AOD analysis (vertex and track cuts also here!!!!)
  //AliMCEvent* gMCEvent = dynamic_cast<AliMCEvent*>(InputEvent()); // from TaskSE
  AliMCEvent* gMCEvent = dynamic_cast<AliMCEvent*>(MCEvent());
  if(!gMCEvent) {
    Printf("ERROR: gMC not available");
    return;
  }

  Int_t nAcceptedParticles = 0;
  
  // event selection done in AliAnalysisTaskSE::Exec() --> this is not used
  fHistEventStats->Fill(1); //all events
  fHistEventStats->Fill(2); //triggered events
  
  const AliVVertex *vertex = gMCEvent->GetPrimaryVertex();
  if(vertex) {
    fHistEventStats->Fill(3); //events with a proper vertex
    if(TMath::Abs(vertex->GetX()) < fVxMax) {
      if(TMath::Abs(vertex->GetY()) < fVyMax) {
	if(TMath::Abs(vertex->GetZ()) < fVzMax) {
	  fHistEventStats->Fill(4); //analyzed events
	  fHistVx->Fill(vertex->GetX());
	  fHistVy->Fill(vertex->GetY());
	  fHistVz->Fill(vertex->GetZ());
	  	  
	  //Printf("There are %d tracks in this event", gMC->GetNumberOfTracks());
	  for (Int_t iParticle = 0; iParticle < gMCEvent->GetNumberOfTracks(); iParticle++) {
	    AliMCParticle* gParticle = dynamic_cast<AliMCParticle *>(gMCEvent->GetTrack(iParticle));
	    if (!gParticle) {
	      Printf(Form("Could not receive particle %d", iParticle));
	      continue;
	    }
	    
	    //exclude non stable particles
	    if(!(gMCEvent->IsPhysicalPrimary(iParticle))) continue;
	    
	    //kinematic cuts cuts
	    Float_t pX  = gParticle->Px();
	    Float_t pY  = gParticle->Py();
	    Float_t pT  = gParticle->Pt();
	    Float_t eta = gParticle->Eta();
	    
	    // Kinematics cuts from ESD track cuts
	    if( pT < fPtMin || pT > fPtMax)      continue;
	    if( eta < fEtaMin || eta > fEtaMax)  continue;
	    
	    nAcceptedParticles += 1;
	  } //particle loop
	  
	  fHistMultiplicity->Fill(nAcceptedParticles);
	}//Vz cut
      }//Vy cut
    }//Vx cut
  }//vertex object valid
}      

//________________________________________________________________________
void  AliAnalysisTaskMC::FinishTaskOutput(){
  //
}

//________________________________________________________________________
void AliAnalysisTaskMC::Terminate(Option_t *) {
  // Draw result to the screen
  // Called once at the end of the query

  // not implemented ...

}