// ************************************************************************
//
// 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
//
// ****************************************************************************************
//
// Analysis Task adopted for FastSim studies of J/Psi, Eta_c various decay and recoil modes
//
// F.Nerling 04/2013
// ****************************************************************************************

// The header file
#include "PndScrutJpsiEtacTask.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"
#include "TString.h"

// RHO headers
#include "RhoCandidate.h"
#include "RhoHistogram/RhoTuple.h"
#include "RhoFactory.h"
#include "RhoMassParticleSelector.h"
#include "RhoEnergyParticleSelector.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   -------------------------------------------
PndScrutJpsiEtacTask::PndScrutJpsiEtacTask(double pbarmom, TString outname, int decayMode, int recoilMode, int dataMode, TString pidQuality, Double_t chiqCut) :
  FairTask("Panda Scrutiny Analysis Task")
{
	double mp=0.938272;
	fIni.SetXYZT(0,0,pbarmom, sqrt(pbarmom*pbarmom+mp*mp)+mp);

	fOutName = outname;
	// *** Decay and Recoil mode

        dMode = decayMode;  // 0: etac  -> K+K-pi0, 1: etac -> KsK+/-pi-/+
	                    // 2: J/psi -> e+e-,    3: J/psi -> mu+mu-
	rMode = recoilMode; // 0: pi+pi-, 1: pi0pi0, 2: pi0, 3: gam, 4: eta 5: eta eta
        fMode = dataMode;   // 0: 1rModedMode zB 100 fuer  900: DPMbkgd

        // *** PID quality selector
	fPIDquality = pidQuality; // All, Loose, Tight, VeryTight, Best
	fChiqCut=chiqCut;
}
// -------------------------------------------------------------------------


// -----   Destructor   ----------------------------------------------------
PndScrutJpsiEtacTask::~PndScrutJpsiEtacTask()
{
	delete fFile;
}
// -------------------------------------------------------------------------


// -----   Public method Init   --------------------------------------------
InitStatus PndScrutJpsiEtacTask::Init()
{
	// *** initialize PndAnalysis object
	fAnalysis = new PndAnalysis();

	// *** reset the event counter
	fEvtCount = 0;

	// *******
	// ******* PREPARE/CREATE THE STUFF YOU NEED
	// *******

	fPdg = TDatabasePDG::Instance();

	// ***
	// *** Prepare RhoTuple output
	// ***
	TDirectory *dir = gDirectory;
	fFile=new TFile(fOutName,"RECREATE");
	fFile->cd();

	// *** create some ntuples
	//ntp1 = new RhoTuple("ntp1", "etac analysis");
	ntp2 = new RhoTuple("ntp2", "pbarp analysis");
	nmc  = new RhoTuple("nmc",  "mctruth info");


	// *** Mass selector for the jpsi cands
	double m0_eta  = fPdg->GetParticle("eta")->Mass();   // Get nominal PDG mass of the eta
	double m0_etac = fPdg->GetParticle("eta_c")->Mass();   // Get nominal PDG mass of the eta_c
	double m0_Ks   = fPdg->GetParticle("K_S0")->Mass();      // Get nominal PDG mass of the Ks0
	double m0_jpsi = fPdg->GetParticle("J/psi")->Mass();   // Get nominal PDG mass of the J/psi

	// create selectors
	jpsiMassSel  = new RhoMassParticleSelector("jpsi",m0_jpsi,1.2);
	etaMassSel   = new RhoMassParticleSelector("eta",m0_eta,0.050); // +- 5sigma!
	etacMassSel  = new RhoMassParticleSelector("etac",m0_etac,0.9);
	pi0MassSel   = new RhoMassParticleSelector("pi0", 0.135 ,0.050); // +- 5 sigma!
	gammaSel     = new RhoEnergyParticleSelector("gam", 50+0.100, 100); // >= 0.1
	KsMassSel    = new RhoMassParticleSelector("Ks", m0_Ks ,0.085); // +- 5 sigma
	pbarpMassSel = new RhoMassParticleSelector("pbarp", fIni.M() ,0.7); // +- 5 sigma


	// assign RhoTuples to outfile
	//if (ntp1) ntp1->GetInternalTree()->SetDirectory(gDirectory);
	if (ntp2) ntp2->GetInternalTree()->SetDirectory(gDirectory);
	if (nmc && fMode!=900)  nmc->GetInternalTree()->SetDirectory(gDirectory);

	// *** restore original gDirectory
	dir->cd();

	return kSUCCESS;
}

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

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

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


// -----   Public method Exec   --------------------------------------------
void PndScrutJpsiEtacTask::Exec(Option_t* opt)
{
	// *** some variables
	int j=0, k=0, kk=0;

	// *** necessary to read the next event
	fAnalysis->GetEvent();

	// *** print event counter
	if (!(++fEvtCount%100)) cout << "evt "<<fEvtCount<<endl;

	// *******
	// ******* PUT ANALYSIS CODE HERE
	// *******

	TString pidalg = "PidChargedProbability"; //alles oder nix, spaeter

	// *** 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 muplus, muminus, piplus, piminus, jpsi, psi2s, all, mclist;


	// dMode = decayMode;     // 0: etac  -> K+K-pi0, 1: etac -> KsK+/-pi-/+
	                          // 2: J/psi -> e+e-,    3: J/psi -> mu+mu-
	// rMode = recoilMode;    // 0: pi+pi-, 1: pi0pi0, 2: pi0, 3: gam, 4: eta 5: eta eta

	// needed reso name for tree
	TString resonName = "etac";
	if(dMode == 2 || dMode == 3) resonName = "jpsi";

	RhoCandList comp, eplus, eminus, muplus, muminus, piplus, piminus, pi0, gamma, Kplus, Kminus, Ks, eta, pbarp, all, mclist; //comp -> jpsi oder etac

	//evtl 	if(dMode==0 || rMode==1) { //-- etac decay 0: KKpi0
	fAnalysis->FillList(eplus,   "Electron"+fPIDquality+"Plus",  pidalg);
	fAnalysis->FillList(eminus,  "Electron"+fPIDquality+"Minus", pidalg);
	fAnalysis->FillList(muplus,  "Muon"+fPIDquality+"Plus",  pidalg);
	fAnalysis->FillList(muminus, "Muon"+fPIDquality+"Minus", pidalg);
	fAnalysis->FillList(piplus,  "Pion"+fPIDquality+"Plus",  pidalg);
	fAnalysis->FillList(piminus, "Pion"+fPIDquality+"Minus", pidalg);
	fAnalysis->FillList(Kplus,   "Kaon"+fPIDquality+"Plus",  pidalg);
	fAnalysis->FillList(Kminus,  "Kaon"+fPIDquality+"Minus", pidalg);
	fAnalysis->FillList(gamma,   "Neutral", pidalg);
	gamma.Select(gammaSel);

	// *** Select with PID info pidalg and ('All'); type and mass are set
	//-- complete lists needed
	if(dMode==0 || rMode==1 || rMode==2) { //-- etac decay 0: KKpi0 oder recoil case
	  pi0.Combine(gamma,gamma);
	  pi0.SetType(111);
	  pi0.Select(pi0MassSel);
	}

	if(rMode==4 || rMode==5) { //-- eta recoil case
	  eta.Combine(gamma,gamma);
	  eta.SetType(221);
	  eta.Select(etaMassSel);
	}


	// *** Setup event shape object
	fAnalysis->FillList(all,   "All", pidalg);
	PndEventShape evsh(all, fIni, 0.05, 0.1);

	if(fMode != 900){
	  // *** store MC info in ntuple
	  fAnalysis->FillList(mclist, "McTruth");
	  nmc->Column("ev", (Int_t) fEvtCount);
	  qa.qaMcList("",   mclist, nmc);
	  nmc->DumpData();
	}

	// *** combinatorics for etac -> KKpi0
	if(dMode==0){
	  comp.Combine(Kplus, Kminus, pi0);
	  comp.SetType(441);
	  comp.Select(etacMassSel);
	}

	// *** combinatorics for etac -> KsKpi
	if(dMode==1){
	  Ks.Combine(piplus,piminus);
	  Ks.SetType(310);
	  Ks.Select(KsMassSel);
	  comp.Combine(Ks, Kminus, piplus);
	  comp.CombineAndAppend(Ks, Kplus, piminus);
	  comp.SetType(441);
	  comp.Select(etacMassSel);
	}

	// *** combinatorics for J/Psi -> e+e-
	if(dMode==2){
	  comp.Combine(eplus,eminus);
	  comp.SetType(443);
	  comp.Select(jpsiMassSel);
	}

	// *** combinatorics for J/Psi -> mu+mu-
	if(dMode==3){
	  comp.Combine(muplus,muminus);
	  comp.SetType(443);
	  comp.Select(jpsiMassSel);
	}


	int njmct = fAnalysis->McTruthMatch(comp); // match the whole list to count #matches (should be only 1)

// 	// *** write ntuple for the ccbar reconstruction
// 	for (j=0;j<comp.GetLength();++j)
// 	{
// 		// some general info about event (actually written for each candidate!)
// 		ntp1->Column("ev",		(Float_t) fEvtCount);
// 		ntp1->Column("cand",	(Float_t) j);
// 		ntp1->Column("ncand",   (Float_t) comp.GetLength());
// 		ntp1->Column("nmct",    (Float_t) njmct);
// 		ntp1->Column("mode",    (Int_t) fMode);
// 		ntp1->Column("dmode",    (Int_t) dMode);
// 		ntp1->Column("rmode",    (Int_t) dMode);
//
// 		// store info about initial 4-vector
// 		qa.qaP4("beam", fIni, ntp1);
//
// 		// dump information about composite candidate tree recursively (see PndTools/AnalysisTools/PndRhoTupleQA)
// 		qa.qaComp(resonName+"_", comp[j], ntp1);
//
// 		// dump info about event shapes
// 		//qa.qaEventShapeShort("es",&evsh, ntp1);
//
// 		// *** store the 4-vector of the truth matched candidate (or a dummy, if not matched to keep ntuple consistent)
// 		RhoCandidate *truth = comp[j]->GetMcTruth();
// 		TLorentzVector lv;
// 		if (truth) lv = truth->P4();
// 		qa.qaP4("true_"+resonName+"_", lv, ntp1);
//
// 		ntp1->DumpData();
// 	}


	// *** some rough mass selection on pbarp
	//if(dMode == 0 || dMode == 1) comp.Select(etacMassSel);
	//if(dMode == 2 || dMode == 3) comp.Select(jpsiMassSel);

	// *** combinatorics for recoil (pi+ pi-)
	if(rMode == 0){
	  pbarp.Combine(comp, piplus, piminus);
	}
	// *** combinatorics for recoil (pi0 pi0)
	if(rMode == 1){
	  pbarp.Combine(comp, pi0, pi0);
	}
	// *** combinatorics for recoil (pi0)
	if(rMode == 2){
	  pbarp.Combine(comp, pi0);
	}
	// *** combinatorics for recoil (gam)
	if(rMode == 3){
	  pbarp.Combine(comp, gamma);
	}
	// *** combinatorics for recoil (eta)
	if(rMode == 4){
	  pbarp.Combine(comp, eta);
	}
	// *** combinatorics for recoil (eta eta)
	if(rMode == 5){
	  pbarp.Combine(comp, eta, eta);
	}
	if(rMode == 6){
	  pbarp.Combine(comp, Kplus, Kminus);
	}
    pbarp.SetType(88880); // pbarpSystem0

	pbarp.Select(pbarpMassSel);

	int npsimct = fAnalysis->McTruthMatch(pbarp); // match the whole list to count #matches (should be only 1)

	// *** write ntuple for the pbarp reconstruction
	for (j=0;j<pbarp.GetLength();++j)
	{
		//FN 4C fit introduced/used
		PndKinFitter kinfit(pbarp[j]);
		kinfit.Add4MomConstraint(fIni);
		kinfit.Fit();

        // now do a wide chi2 cut to save calculation time and disk space
        if (kinfit.GetChi2() > fChiqCut) continue;

        // some general info about event (actually written for each candidate!)
        ntp2->Column("ev",      (Float_t) fEvtCount);
        ntp2->Column("cand",    (Float_t) j);
        ntp2->Column("ncand",   (Float_t) pbarp.GetLength());
        ntp2->Column("nmct",    (Float_t) npsimct);
        ntp2->Column("mode",    (Int_t) fMode);
        ntp2->Column("dmode",    (Int_t) dMode);
        ntp2->Column("rmode",    (Int_t) rMode);

		RhoCandidate *pbarpfit=pbarp[j]->GetFit();

		// store info about initial 4-vector
		qa.qaP4("beam", fIni, ntp2);

		// dump information about composite candidate tree recursively (see PndTools/AnalysisTools/PndRhoTupleQA)
		qa.qaComp("pbarp_", pbarp[j], ntp2);
		//FN 4C fit introduced/used
		//qa.qaComp("fpbarp_",pbarpfit, ntp2);
		ntp2->Column("fpbarp_m",   pbarpfit->M() );
		ntp2->Column("fpbarp_p",   pbarpfit->P() );
		ntp2->Column("fpbarp_e",   pbarpfit->E() );
		for (k=0;k<pbarpfit->NDaughters();++k)
		{
			TString colname=TString::Format("fpbarp_d%d",k);
			ntp2->Column(colname+"m", pbarpfit->Daughter(k)->M() );
			ntp2->Column(colname+"p", pbarpfit->Daughter(k)->P() );
            for (kk=0;kk<pbarpfit->Daughter(k)->NDaughters();++kk)
            {
              TString colname2=TString::Format("fpbarp_d%dd%d",k,kk);
              ntp2->Column(colname2+"m", pbarpfit->Daughter(k)->Daughter(kk)->M() );
              ntp2->Column(colname2+"p", pbarpfit->Daughter(k)->Daughter(kk)->P() );
            }
		}
		ntp2->Column("fchi2", (Float_t) kinfit.GetChi2());

		// dump info about event shapes
		qa.qaESMult("es",  &evsh, ntp2);

		// PID multiplicities
		qa.qaESPidMult("esl",  &evsh, 0.25, 0.0, ntp2);

		// event vars like thrust, sphericity etc
		qa.qaESEventVars("es", &evsh, ntp2);

		// standard min, max values of p, pt in lab, cms
		qa.qaESMinMax("es", &evsh, ntp2);
		//qa.qaEventShapeShort("es",&evsh, ntp2);

		// *** store the 4-vector of the truth matched candidate (or a dummy, if not matched to keep ntuple consistent)
		RhoCandidate *truth = pbarp[j]->GetMcTruth();
		TLorentzVector lv;
		if (truth) lv = truth->P4();
		qa.qaP4("true_pbarp_", lv, ntp2);

		ntp2->DumpData();
	}

}


void PndScrutJpsiEtacTask::Finish()
{

	// *******
	// ******* STORE YOUR HISTOS AND TUPLES
	// *******

	fFile->Write();
	fFile->Close();

}

ClassImp(PndScrutJpsiEtacTask)