// ************************************************************************
//
// Template for an Analysis Task, 
// 
// for the December 2013 CM Tutorial 
//
// For further info see also
//
// http://panda-wiki.gsi.de/cgi-bin/viewauth/Computing/PandaRootRhoTutorial
// http://panda-wiki.gsi.de/cgi-bin/view/Computing/PandaRootAnalysisJuly13
//
// K.Goetzen 11/2013
//
// ************************************************************************


// The header file
#include "PndLambdacAnaTask.h"

// C++ headers
#include <string>
#include <iostream>

// FAIR headers
#include "FairRootManager.h"
#include "FairRunAna.h"
#include "FairRuntimeDb.h"
#include "FairRun.h"
#include "FairRuntimeDb.h"

// ROOT headers
#include "TClonesArray.h"
#include "TVector3.h"
#include "TH1F.h"
#include "TH2F.h"
#include "TDatabasePDG.h"

// RHO headers
#include "RhoCandidate.h"
#include "RhoHistogram/RhoTuple.h"
#include "RhoFactory.h"
#include "RhoMassParticleSelector.h"
#include "RhoTuple.h"

// Analysis headers
#include "PndAnalysis.h"
#include "Pnd4CFitter.h"
#include "PndKinVtxFitter.h"
#include "PndKinFitter.h"
#include "PndVtxPoca.h"
#include "PndRhoTupleQA.h"
#include "PndEventShape.h"
		
		
using std::cout;
using std::endl;


// -----   Default constructor   -------------------------------------------
PndLambdacAnaTask::PndLambdacAnaTask(double pbarmom) :
  FairTask("Panda Lambdac Analysis Task") 
{ 
	double mp=0.938272;
	fIni.SetXYZT(0,0,pbarmom, sqrt(pbarmom*pbarmom+mp*mp)+mp);
}
// -------------------------------------------------------------------------


// -----   Destructor   ----------------------------------------------------
PndLambdacAnaTask::~PndLambdacAnaTask() 
{ 
  ;
}
// -------------------------------------------------------------------------


// -----   Public method Init   --------------------------------------------
InitStatus PndLambdacAnaTask::Init() 
{		
	// *** initialize PndAnalysis object
	fAnalysis = new PndAnalysis();
	
	// *** reset the event counter
	fEvtCount = 0;

	// *******
	// ******* PREPARE/CREATE THE STUFF YOU NEED
	// *******
	
	// *** Mass selector for the jpsi cands

        double m0_Kplus = TDatabasePDG::Instance()->GetParticle("K+")->Mass();   
        double m0_Kminus = TDatabasePDG::Instance()->GetParticle("K-")->Mass();   
        double m0_piplus = TDatabasePDG::Instance()->GetParticle("pi+")->Mass();   
        double m0_piminus = TDatabasePDG::Instance()->GetParticle("pi-")->Mass();   
        double m0_pplus = TDatabasePDG::Instance()->GetParticle("proton")->Mass();   
        double m0_pminus = TDatabasePDG::Instance()->GetParticle("antiproton")->Mass();   
        double m0_lplus = TDatabasePDG::Instance()->GetParticle("Lambda_c+")->Mass();   
        double m0_lminus = TDatabasePDG::Instance()->GetParticle("Lambda_c-")->Mass();  

	fMassLambdac = m0_lplus;

        lpMassSel=new RhoMassParticleSelector("lambdac",m0_lplus,0.2);
        lmMassSel=new RhoMassParticleSelector("lambdacb",m0_lminus,1.0);

	// *** create some ntuples etc.

	nmc  = new RhoTuple("nmc",  "mctruth info");
	nlr  = new RhoTuple("nlr",  "Lambda_c reconstruction, inclusive");
	nlr2  = new RhoTuple("nlr2",  "Lambda_c reconstruction, exclusive");

        hlm = new TH2F("hlm","",200,1.9,2.7,200,1.9,2.7);
        hlms = new TH2F("hlms","",200,1.9,2.7,200,1.9,2.7);
        hb = new TH1F("hb","",200,4,6);
        hlmf = new TH2F("hlmf","",200,1.9,2.7,200,1.9,2.7);
	hbf = new TH1F("hbf","",200,4,6);
        hchisq = new TH1F("hchisq","",200,0,100.);
	hlpm = new TH1F("hlpm", "", 200, 1.9,2.7);
	hlpmf = new TH1F("hlpmf", "", 200, 1.9,2.7);
	hlmm = new TH1F("hlmm", "", 200, 1.9,2.7);
	hlmmf = new TH1F("hlmmf", "", 200, 1.9,2.7);
	hmpi0 = new TH1F("hmpi0","",500,0.,1.);
	hmpi0e = new TH1F("hmpi0e","",500,0.,1.);
	
	return kSUCCESS;
}

// -------------------------------------------------------------------------
	
void PndLambdacAnaTask::SetParContainers() 
{
  // *** Get run and runtime database
  FairRun* run = FairRun::Instance();
  if ( ! run ) Fatal("SetParContainers", "No analysis run");
}

// -------------------------------------------------------------------------

// -----   Public method Exec   --------------------------------------------
void PndLambdacAnaTask::Exec(Option_t* opt)
{
	// *** some variables
        int i=0,j=0,k=0;
	
	// *** necessary to read the next event
	fAnalysis->GetEvent();
	
	// *** print event counter
	if (!(++fEvtCount%1000)) cout << "evt "<<fEvtCount<<endl;
	
	// *******
	// ******* PUT ANALYSIS CODE HERE			
	// *******
	
	TString pidalg = "PidChargedProbability";
	
	// *** QA tool for simple dumping of analysis results in RhoRuple
	// *** WIKI: https://panda-wiki.gsi.de/foswiki/bin/view/Computing/PandaRootAnalysisJuly13#PndRhoTupleQA
	PndRhoTupleQA qa(fAnalysis,fIni.P());

	// *** RhoCandLists for the analysis
	RhoCandList pplus, pminus, piplus, piminus, kaon, kplus, kminus, all, mclist, lambdac, lambdacb, gamma, beam, pi0;
	
	
	// *** Select with PID info pidalg and ('All'); type and mass are set 		

	fAnalysis->FillList(kplus, "KaonLoosePlus", "DrcBarrelProbability;DrcDiscProbability");
	fAnalysis->FillList(kminus, "KaonLooseMinus", "DrcBarrelProbability;DrcDiscProbability");

	fAnalysis->FillList(piminus, "PionLooseMinus", "DrcBarrelProbability;DrcDiscProbability");
	fAnalysis->FillList(piplus, "PionLoosePlus", "DrcBarrelProbability;DrcDiscProbability");
	
	fAnalysis->FillList(pminus, "ProtonLooseMinus", "DrcBarrelProbability;DrcDiscProbability");
	fAnalysis->FillList(pplus, "ProtonLoosePlus", "DrcBarrelProbability;DrcDiscProbability");
	fAnalysis->FillList(gamma, "Neutral");
	
	// *** Setup event shape object
	fAnalysis->FillList(all,   "All", pidalg);
	PndEventShape evsh(all, fIni, 0.05, 0.1);	
		
	// *** store MC info in ntuple
	fAnalysis->FillList(mclist, "McTruth");
	
	nmc->Column("ev", (Int_t) fEvtCount);
	qa.qaMcList("",   mclist, nmc);
	nmc->DumpData();
	
	pi0.Combine(gamma,gamma);

	int pi0cand=-1;
	double diffmass=-1;

	for (j=0; j<pi0.GetLength(); j++)
	  {
	    hmpi0->Fill(pi0[j]->M());
	    if (0==j) 
	      {
		pi0cand = 0;
		diffmass = fabs(pi0[0]->M()-0.135);
	      }
	    if (fabs(pi0[j]->M()-0.135)<diffmass)
	      {
		pi0cand = j;
		diffmass = fabs(pi0[j]->M()-0.135);
	      }
	  }

	lambdac.Combine(pplus,kminus,piplus);
	lambdac.SetType(TDatabasePDG::Instance()->GetParticle("Lambda_c+")->PdgCode());

	int nlcpmct = fAnalysis->McTruthMatch(lambdac);

	for (j=0; j<lambdac.GetLength(); j++)
	  {
	    TLorentzVector lv(lambdac[j]->GetMomentum(),lambdac[j]->GetEnergy());
	    TLorentzVector lmissing = fIni - lv;
	    hlmm->Fill(lmissing.M());
	    hlpm->Fill(lambdac[j]->M());
	    PndKinFitter mfitter(lambdac[j]);
	    mfitter.AddMassConstraint(fMassLambdac);  
	    mfitter.Fit();

	    TLorentzVector lvf(lambdac[j]->GetFit()->GetMomentum(),lambdac[j]->GetFit()->GetEnergy());
	    TLorentzVector lmissingf = fIni - lvf;

	    if (mfitter.GetProb()>0.01)
	      {
		hlmmf->Fill(lmissingf.M());
		hlpmf->Fill(lambdac[j]->GetFit()->M());
	      }

	    nlr->Column("ev",      (Float_t) fEvtCount);
	    nlr->Column("cand",    (Float_t) j);
	    nlr->Column("ncand",   (Float_t) lambdac.GetLength());
	    nlr->Column("nmctp",    (Float_t) nlcpmct);
	    nlr->Column("npp",    (Float_t) pplus.GetLength() );
	    nlr->Column("npm",    (Float_t) pminus.GetLength() );
	    nlr->Column("nkp",    (Float_t) kplus.GetLength() );
	    nlr->Column("nkm",    (Float_t) kminus.GetLength() );
	    nlr->Column("npip",    (Float_t) piplus.GetLength() );
	    nlr->Column("npim",    (Float_t) piminus.GetLength() );
	    nlr->Column("ngam",    (Float_t) gamma.GetLength() );
	    if (pi0cand>0)
	      {
		nlr->Column("mpi0",    (Float_t) pi0[pi0cand]->M() );
	      }
	    else
	      {
		nlr->Column("mpi0", -1.);
	      }

	    qa.qaP4("beam", fIni, nlr);

	    qa.qaP4("lp", lv, nlr);
	    qa.qaP4("lpf", lvf, nlr);

	    qa.qaP4("lm", lmissing, nlr);
	    qa.qaP4("lmf", lmissingf, nlr);

	    nlr->Column("chi2", (Float_t) mfitter.GetChi2());

	    TLorentzVector pp(lambdac[j]->Daughter(0)->GetMomentum(),
			      lambdac[j]->Daughter(0)->GetEnergy());

	    TLorentzVector km(lambdac[j]->Daughter(1)->GetMomentum(),
			      lambdac[j]->Daughter(1)->GetEnergy());

	    TLorentzVector pip(lambdac[j]->Daughter(2)->GetMomentum(),
			      lambdac[j]->Daughter(2)->GetEnergy());

	    TLorentzVector ppf(lambdac[j]->Daughter(0)->GetFit()->GetMomentum(),
			       lambdac[j]->Daughter(0)->GetFit()->GetEnergy());

	    TLorentzVector kmf(lambdac[j]->Daughter(1)->GetFit()->GetMomentum(),
			       lambdac[j]->Daughter(1)->GetFit()->GetEnergy());

	    TLorentzVector pipf(lambdac[j]->Daughter(2)->GetFit()->GetMomentum(),
				lambdac[j]->Daughter(2)->GetFit()->GetEnergy());

	    qa.qaP4("pp", pp, nlr);
	    qa.qaP4("ppf", ppf, nlr);

	    qa.qaP4("km", km, nlr);
	    qa.qaP4("kmf", kmf, nlr);

	    qa.qaP4("pip", pip, nlr);
	    qa.qaP4("pipf", pipf, nlr);

	    RhoCandidate *truth = lambdac[j]->GetMcTruth();
	    TLorentzVector lt;
	    if (truth) lt = truth->P4();
	    qa.qaP4("trlp", lt, nlr);

	    //	    qa.qaEventShapeShort("es",&evsh, nlr);
	 
	    nlr->DumpData();
	  }

	
	lambdacb.Combine(pminus,gamma);
	lambdacb.SetType(TDatabasePDG::Instance()->GetParticle("Lambda_c-")->PdgCode());	

	int nlcmmct = fAnalysis->McTruthMatch(lambdacb);

	lambdac.Select(lpMassSel);
	lambdacb.Select(lmMassSel);
	
	beam.Combine(lambdac,lambdacb);
	beam.SetType(88880);
	
	for (j=0;j<beam.GetLength();j++)
	  {
	    Pnd4CFitter fitter(beam[j],fIni);
	    fitter.FitConserveMasses();
	    hchisq->Fill(fitter.GetChi2());
	    
	    hb->Fill(beam[j]->M());
	    hlm->Fill(beam[j]->Daughter(0)->M(),
		      beam[j]->Daughter(1)->M());
	    
	    if (fitter.GetChi2()<20.)
	      {
		hbf->Fill(beam[j]->GetFit()->M());
		hlms->Fill(beam[j]->Daughter(0)->M(),
			   beam[j]->Daughter(1)->M());
		hlmf->Fill(beam[j]->Daughter(0)->GetFit()->M(),
			   beam[j]->Daughter(1)->GetFit()->M());

		for (k=0; k<pi0.GetLength(); k++)
		  {
		    hmpi0e->Fill(pi0[k]->M());
		  }
	      }

	    nlr2->Column("ev",      (Float_t) fEvtCount);
	    nlr2->Column("cand",    (Float_t) j);
	    nlr2->Column("ncand",   (Float_t) beam.GetLength());
	    nlr2->Column("nmctp",    (Float_t) nlcpmct);
	    nlr2->Column("nmctm",    (Float_t) nlcmmct);
	    nlr2->Column("npp",    (Float_t) pplus.GetLength() );
	    nlr2->Column("npm",    (Float_t) pminus.GetLength() );
	    nlr2->Column("nkp",    (Float_t) kplus.GetLength() );
	    nlr2->Column("nkm",    (Float_t) kminus.GetLength() );
	    nlr2->Column("npip",    (Float_t) piplus.GetLength() );
	    nlr2->Column("npim",    (Float_t) piminus.GetLength() );
	    nlr2->Column("ngam",    (Float_t) gamma.GetLength() );
	    if (pi0cand>0)
	      {
		nlr2->Column("mpi0",    (Float_t) pi0[pi0cand]->M() );
	      }
	    else
	      {
		nlr2->Column("mpi0", -1.);
	      }

	    qa.qaP4("beam", fIni, nlr2);

	    TLorentzVector lp(beam[j]->Daughter(0)->GetMomentum(),
			      beam[j]->Daughter(0)->GetEnergy());
	    TLorentzVector lpf(beam[j]->Daughter(0)->GetFit()->GetMomentum(),
			       beam[j]->Daughter(0)->GetFit()->GetEnergy());

	    TLorentzVector lm(beam[j]->Daughter(1)->GetMomentum(),
			      beam[j]->Daughter(1)->GetEnergy());
	    TLorentzVector lmf(beam[j]->Daughter(1)->GetFit()->GetMomentum(),
			       beam[j]->Daughter(1)->GetFit()->GetEnergy());

	    TLorentzVector pp(beam[j]->Daughter(0)->Daughter(0)->GetMomentum(),
			      beam[j]->Daughter(0)->Daughter(0)->GetEnergy());

	    TLorentzVector km(beam[j]->Daughter(0)->Daughter(1)->GetMomentum(),
			      beam[j]->Daughter(0)->Daughter(1)->GetEnergy());

	    TLorentzVector pip(beam[j]->Daughter(0)->Daughter(2)->GetMomentum(),
			       beam[j]->Daughter(0)->Daughter(2)->GetEnergy());

	    TLorentzVector ppf(beam[j]->Daughter(0)->Daughter(0)->GetFit()->GetMomentum(),
			       beam[j]->Daughter(0)->Daughter(0)->GetFit()->GetEnergy());

	    TLorentzVector kmf(beam[j]->Daughter(0)->Daughter(1)->GetFit()->GetMomentum(),
			       beam[j]->Daughter(0)->Daughter(1)->GetFit()->GetEnergy());

	    TLorentzVector pipf(beam[j]->Daughter(0)->Daughter(2)->GetFit()->GetMomentum(),
				beam[j]->Daughter(0)->Daughter(2)->GetFit()->GetEnergy());

	    TLorentzVector pm(beam[j]->Daughter(1)->Daughter(0)->GetMomentum(),
			      beam[j]->Daughter(1)->Daughter(0)->GetEnergy());

	    TLorentzVector pmf(beam[j]->Daughter(1)->Daughter(0)->GetFit()->GetMomentum(),
			       beam[j]->Daughter(1)->Daughter(0)->GetFit()->GetEnergy());

	    TLorentzVector g(beam[j]->Daughter(1)->Daughter(1)->GetMomentum(),
			     beam[j]->Daughter(1)->Daughter(1)->GetEnergy());

	    TLorentzVector gf(beam[j]->Daughter(1)->Daughter(1)->GetFit()->GetMomentum(),
			      beam[j]->Daughter(1)->Daughter(1)->GetFit()->GetEnergy());

	    TLorentzVector lppb(beam[j]->GetMomentum(),
				  beam[j]->GetEnergy());

	    TLorentzVector lppbf(beam[j]->GetFit()->GetMomentum(),
				   beam[j]->GetFit()->GetEnergy());

	    qa.qaP4("pp", pp, nlr2);
	    qa.qaP4("ppf", ppf, nlr2);

	    qa.qaP4("km", km, nlr2);
	    qa.qaP4("kmf", kmf, nlr2);

	    qa.qaP4("pip", pip, nlr2);
	    qa.qaP4("pipf", pipf, nlr2);

	    qa.qaP4("pm", pm, nlr2);
	    qa.qaP4("pmf", pmf, nlr2);

	    qa.qaP4("g", g, nlr2);
	    qa.qaP4("gf", gf, nlr2);

	    qa.qaP4("lp", lp, nlr2);
	    qa.qaP4("lpf", lpf, nlr2);

	    qa.qaP4("lm", lm, nlr2);
	    qa.qaP4("lmf", lmf, nlr2);

	    qa.qaP4("ppb", lppb, nlr2);
	    qa.qaP4("ppbf", lppbf, nlr2);

	    nlr2->Column("chi2", (Float_t) fitter.GetChi2());

	    RhoCandidate *truth0 = beam[j]->Daughter(0)->GetMcTruth();
	    RhoCandidate *truth1 = beam[j]->Daughter(1)->GetMcTruth();
	    TLorentzVector lt0, lt1;
	    if (truth0) lt0 = truth0->P4();
	    if (truth1) lt1 = truth1->P4();
	    qa.qaP4("trlp", lt0, nlr2);
	    qa.qaP4("trlm", lt1, nlr2);

	    // qa.qaEventShapeShort("es",&evsh, nlr2);
	 
	    nlr2->DumpData();
	  }
	
}


void PndLambdacAnaTask::Finish()
{
	
	// *******
	// ******* STORE YOUR HISTOS AND TUPLES
	// *******

  nmc->GetInternalTree()->Write();
  nlr->GetInternalTree()->Write();
  nlr2->GetInternalTree()->Write();

  hlm->Write();
  hlms->Write();
  hb->Write();
  hlmf->Write();
  hbf->Write();
  hchisq->Write();
  hlpm->Write();
  hlpmf->Write();
  hlmm->Write();
  hlmmf->Write();
  hmpi0->Write();  
  hmpi0e->Write();  
}

ClassImp(PndLambdacAnaTask)