/*
* PndLLbarAnaTask_DecayTreeFit.cxx
*
*  Created on: July 7, 2015
*  Author: Walter Ikegami Andersson
*/

// The header file
#include "PndLLbarAnaTask_DecayTreeFit.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 "TLorentzVector.h"
#include "TVector3.h"
#include "TH1F.h"
#include "TH2F.h"
#include "TParticlePDG.h"

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

// PandaRoot headers
#include "PndAnalysis.h"
#include "Pnd4CFitter.h"
#include "PndKinVtxFitter.h"
#include "PndKinFitter.h"
#include "PndVtxPoca.h"
#include "PndPidCandidate.h"
		
using std::cout;
using std::endl;


// -----   Default constructor   -------------------------------------------
PndLLbarAnaTask_DecayTreeFit::PndLLbarAnaTask_DecayTreeFit() :
  FairTask("PandaLLbarAnaTask"), fIni(1.64){ 
}
// -------------------------------------------------------------------------


// -----   Destructor   ----------------------------------------------------
PndLLbarAnaTask_DecayTreeFit::~PndLLbarAnaTask_DecayTreeFit() { }
// -------------------------------------------------------------------------


// ----- Method to select  PID candidates
int PndLLbarAnaTask_DecayTreeFit::SelectTruePid(PndAnalysis *ana, RhoCandList &l)
{
	int removed = 0;
	
	for (int ii=l.GetLength()-1;ii>=0;--ii)
	{
		if ( !(ana->McTruthMatch(l[ii])) )
		{
			l.Remove(l[ii]);
			removed++;
		}
	}
	
	return removed;
}
// -------------------------------------------------------------------------


// -----   Public method Init   --------------------------------------------
InitStatus PndLLbarAnaTask_DecayTreeFit::Init()
{		
	// initialize analysis object
	theAnalysis = new PndAnalysis();

	// reset the event counter
	nevts = 0;
	
	//PDG database
	pdg = new TDatabasePDG();

	//Particle masses
	m0_pi = ((TParticlePDG*)pdg->GetParticle(-211))->Mass();
	m0_p = ((TParticlePDG*)pdg->GetParticle(-2212))->Mass();
	m0_lam = ((TParticlePDG*)pdg->GetParticle(-3122))->Mass();

	//Mass selector
	lambdaMassSelector = new RhoMassParticleSelector("lambda0", m0_lam, 0.3);

	// tuples
	fntpPiMinus_P4 = new RhoTuple("ntpPiMinus_P4", "PiMinus info");
	fntpPiPlus_P4 = new RhoTuple("ntpPiPlus_P4", "PiPlus info");
	fntpProton_P4 = new RhoTuple("ntpProton_P4", "Proton info");
	fntpAntiProton_P4 = new RhoTuple("ntpAntiProton_P4", "AntiProton info");

	fntpLambda_P4 = new RhoTuple("ntpLambda_P4", "Lambda info");
	fntpLambdaBar_P4 = new RhoTuple("ntpLambdaBar_P4", "LambdaBar info");

	fntpPbarp = new RhoTuple("ntpPbarp", "Pbarp system info");
	fntpBestPbarp = new RhoTuple("ntpBestPbarp", "BestPbarp system info");

	// *** ------------------------- ***
	// *** Create some 1D histograms ***
	// *** ------------------------- ***

	//Piminus histos
	hpiminus_P = new TH1F("hpiminus_P","#pi^{-} P",300,0.,3.);

	//Piplus histos
	hpiplus_P = new TH1F("hpiplus_P","#pi^{+} P",300,0.,3.);

	//Proton histos
	hproton_P = new TH1F("hproton_P"," p P",300,0.,3.);

	//AntiProton histos
	hantiproton_P = new TH1F("hantiproton_P","p Pbar",300,0.,3.);

	//Lambda histos
	hlam0_M_all = new TH1F("hlam0_M_all","#Lambda_{0} mass",200,0,5);

	//Lambdabar histos
	hlam0bar_M_all = new TH1F("hlam0bar_M_all","#Lambda_{0} bar mass",200,0,5);
	
	//Pbarp system histos
	hpbarp_M_all = new TH1F("hppbar_M_all","pp bar mass",200,0,10);

	// *** --------------------- ***
	// *** True MC 1D histograms ***
	// *** --------------------- ***

	// *** ------------------------- ***
	// *** Create some 2D histograms ***
	// *** ------------------------- ***

	return kSUCCESS;
	cout<<"PndLLbarAnaTask_DecayTreeFit - Initalization successful!"<<endl;
}//End PndLLbarAnaTask_DecayTreeFit::Init()

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

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


// -----   Public method Exec   --------------------------------------------
void PndLLbarAnaTask_DecayTreeFit::Exec(Option_t* opt)
{
	// some variables
	int i=0,j=0, k=0, l=0;
	
	// necessary to read the next event
	theAnalysis->GetEvent();
	
	if (!(++nevts%100)) cout << "evt "<<nevts<<endl;
	
	// *** ---------------------------- ***
	// *** Now the analysis stuff comes ***
	// *** ---------------------------- ***
		
	// *** RhoCandLists for the analysis
	RhoCandList p, pbar, piplus, piminus, lam0, lam0bar, pbarpsystem, mclist;

	double beamres = 1e-4* fIni;
	RhoError inicov(4);
	inicov[0][0] = 1e-6*1e-6;
	inicov[1][1] = 1e-6*1e-6;
	inicov[2][2] = beamres*beamres;
	inicov[3][3] = beamres*beamres /(m0_p*m0_p + fIni*fIni );

	TLorentzVector targ(0,0,0,m0_p);
	TLorentzVector beam(0,0,fIni,sqrt(m0_p*m0_p + fIni*fIni));

	PndRhoTupleQA qa(theAnalysis, fIni);

	// Select with ideal PID info
	//theAnalysis->FillList(p,"ProtonBestPlus","PidAlgoIdealCharged");
	//theAnalysis->FillList(pbar,"ProtonBestMinus","PidAlgoIdealCharged");
	//theAnalysis->FillList(piplus,"PionBestPlus","PidAlgoIdealCharged");
	//theAnalysis->FillList(piminus,"PionBestMinus","PidAlgoIdealCharged");
	//theAnalysis->FillList(mclist,"McTruth");

	// Select with no PID info ('All'); type and mass are set
	theAnalysis->FillList(p,"ProtonAllPlus");
	theAnalysis->FillList(pbar,"ProtonAllMinus");
	theAnalysis->FillList(piplus,"PionAllPlus");
	theAnalysis->FillList(piminus,"PionAllMinus");
	theAnalysis->FillList(mclist,"McTruth");

	// *** ---------------------------------- ***
	// *** All candidates, realistic analysis ***
	// *** ---------------------------------- ***

	// *** combinatorics for lam0 -> p pi
	lam0.Combine(p,piminus);
	lam0bar.Combine(pbar,piplus);

	lam0.Select(lambdaMassSelector);
	lam0bar.Select(lambdaMassSelector);

	lam0.SetType(3122);
	lam0bar.SetType(-3122);

	pbarpsystem.Combine(lam0,lam0bar);
	pbarpsystem.SetType(88888);

	// *** ---------- ***
	// *** Pion- loop ***
	// *** ---------- ***
	for (j=0;j<piminus.GetLength();++j)
	{
		TLorentzVector piminus4=piminus[j]->P4();

		//Fill Rho tuples
		qa.qaP4("piminus_", piminus4, fntpPiMinus_P4);
		qa.qaCand("piminus_", piminus[j], fntpPiMinus_P4);

		fntpPiMinus_P4->DumpData();

		//1D histos
		hpiminus_P->Fill(piminus4.P());

	}

	// *** ---------- ***
	// *** Pion+ loop ***
	// *** ---------- ***
	for (j=0;j<piplus.GetLength();++j)
	{
		TLorentzVector piplus4=piplus[j]->P4();	//Construct ROOT TLorentzVectors

		//Fill Rho tuples
		qa.qaP4("piplus_", piplus4, fntpPiPlus_P4);
		qa.qaCand("piplus_", piplus[j], fntpPiPlus_P4);

		fntpPiPlus_P4->DumpData();
		//1D histos
		hpiplus_P->Fill(piplus4.P());

	}

	// *** ----------- ***
	// *** Proton loop ***
	// *** ----------- ***
	for (j=0;j<p.GetLength();++j)
	{
		TLorentzVector p4=p[j]->P4();

		//Fill Rho tuples
		qa.qaP4("proton_", p4, fntpProton_P4);
		qa.qaCand("proton_", p[j], fntpProton_P4);

		fntpProton_P4->DumpData();

		//1D histos
		hproton_P->Fill(p4.P());

	}

	// *** --------------- ***
	// *** AntiProton loop ***
	// *** --------------- ***
	for (j=0;j<pbar.GetLength();++j)
	{
		TLorentzVector pbar4=pbar[j]->P4();

		//Fill Rho tuples
		qa.qaP4("antiproton_", pbar4, fntpAntiProton_P4);
		qa.qaCand("antiproton_", pbar[j], fntpAntiProton_P4);

		fntpAntiProton_P4->DumpData();

		//1D histos
		hantiproton_P->Fill(pbar4.P());

	}

	// *** ----------------------------- ***
	// *** Lambda (-> Proton Pion-) loop ***
	// *** ----------------------------- ***
	for (j=0;j<lam0.GetLength();++j)
	{
		TLorentzVector lam04 = lam0[j]->P4();

		//Fill Rho tuples
		qa.qaP4("lambda_", lam04, fntpLambda_P4);
		qa.qaCand("lambda_", lam0[j], fntpLambda_P4);

		fntpLambda_P4->DumpData();

		hlam0_M_all->Fill(lam0[j]->M());
	}

	// *** ------------------------------------ ***
	// *** LambdaBar (-> AntiProton Pion+) loop ***
	// *** ------------------------------------ ***
	for (j=0;j<lam0bar.GetLength();++j)
	{
		TLorentzVector lam0bar4 = lam0bar[j]->P4();

		//Fill Rho tuples
		qa.qaP4("lambdabar_", lam0bar4, fntpLambdaBar_P4);
		qa.qaCand("lambdabar_", lam0bar[j], fntpLambdaBar_P4);

		fntpLambdaBar_P4->DumpData();

		hlam0bar_M_all->Fill(lam0bar[j]->M());

	}

	// *** --------------------------------------- ***
	// *** PbarP system (-> Lambda LambdaBar) loop ***
	// *** --------------------------------------- ***
	if (pbarpsystem.GetLength() > 1) cout<<"More than one pbarp combi! Total: "
			<<pbarpsystem.GetLength()<<" event: "<<nevts<<endl;
	//Simple way of choosing best chi2 pbarp system
	Float_t bestChi2 = 999999.;
	Float_t bestProb = 0.;
	Int_t bestIndex = -1;
	bool successfit = false;
	bool bestsuccessfit = false;
	for (j=0;j<pbarpsystem.GetLength();++j)
	{
		RhoLorentzVectorErr pbarperr(pbarpsystem[j]->P4(),inicov);
		RhoDecayTreeFitter treefit(pbarpsystem[j],pbarperr);

		treefit.SetToleranceZ(0.01);
		treefit.setMassConstraint(pbarpsystem[j]->Daughter(0),m0_lam);
		treefit.setMassConstraint(pbarpsystem[j]->Daughter(1),m0_lam);
		successfit = treefit.Fit();

		if (successfit) {
			qa.qaFitter("fitter_",&treefit,fntpPbarp,false);
			qa.qaComp("pbarpfit_",pbarpsystem[j]->GetFit(),fntpPbarp,true,true);
			fntpPbarp->DumpData();
			hpbarp_M_all->Fill(pbarpsystem[j]->M());

			bestsuccessfit = true;
			Float_t dummyChi2 = treefit.GetChi2();
			Float_t dummyProb = treefit.GetProb();
			Int_t dummyIndex = j;

			if (dummyChi2 < bestChi2) {
				bestChi2 = dummyChi2;
				bestIndex = dummyIndex;
				bestProb = dummyProb;
			}

		}

	}

	//Save best pbarpsystem fitted data
	if (bestsuccessfit) {
		fntpBestPbarp->Column("fitter_chisq", (Float_t) bestChi2);
		fntpBestPbarp->Column("fitter_prob", (Float_t) bestProb);
		qa.qaComp("pbarpfit_",pbarpsystem[bestIndex]->GetFit(),fntpBestPbarp,true,true);
		fntpBestPbarp->DumpData();
	}

	p.Cleanup();
	pbar.Cleanup();
	piplus.Cleanup();
	piminus.Cleanup();
	lam0.Cleanup();
	lam0bar.Cleanup();
	pbarpsystem.Cleanup();
	mclist.Cleanup();



}	//End PndLLbarAnaTask_DecayTreeFit::Exec()

void PndLLbarAnaTask_DecayTreeFit::Finish()
{	
	//Write RhoTuple data
	fntpPiMinus_P4->GetInternalTree()->Write();
	fntpPiPlus_P4->GetInternalTree()->Write();
	fntpProton_P4->GetInternalTree()->Write();
	fntpAntiProton_P4->GetInternalTree()->Write();

	fntpLambda_P4->GetInternalTree()->Write();
	fntpLambdaBar_P4->GetInternalTree()->Write();

	fntpPbarp->GetInternalTree()->Write();
	fntpBestPbarp->GetInternalTree()->Write();

	//Write Histos

	//Pi-
	hpiminus_P->Write();

	//Pi+
	hpiplus_P->Write();

	//Proton
	hproton_P->Write();

	//Pbar
	hantiproton_P->Write();

	//Lam0
	hlam0_M_all->Write();
	
	//Lam0bar
	hlam0bar_M_all->Write();

	//Pbarp system
	hpbarp_M_all->Write();
	
}	//End PndLLbarAnaTask_DecayTreeFit::Finish()

ClassImp(PndLLbarAnaTask_DecayTreeFit)