// ************************************************************************
//
// 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 "PndScrutAnaTask.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   -------------------------------------------
PndScrutAnaTask::PndScrutAnaTask(double pbarmom, TString outname) :
  FairTask("Panda Scrutiny Analysis Task") 
{ 
	double mp=0.938272;
	fIni.SetXYZT(0,0,pbarmom, sqrt(pbarmom*pbarmom+mp*mp)+mp);

	fOutName = outname;
}
// -------------------------------------------------------------------------


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


// -----   Public method Init   --------------------------------------------
InitStatus PndScrutAnaTask::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", "pi0 analysis");
	ntp2 = new RhoTuple("ntp2", "pi+pi- pair analysis");
	ntp3 = new RhoTuple("ntp3", "h_c analysis");
	ntp4 = new RhoTuple("ntp4", "h_c 4C fit");
	nmc  = new RhoTuple("nmc",  "mctruth info");
	hchi2_4C = new TH1D("hchi2_4C",";#chi^{2}",2e3,0,1e3);

	// assign RhoTuples to outfile
	if (ntp1) ntp1->GetInternalTree()->SetDirectory(gDirectory);
	if (ntp2) ntp2->GetInternalTree()->SetDirectory(gDirectory);
	if (ntp3) ntp3->GetInternalTree()->SetDirectory(gDirectory);
	if (ntp4) ntp4->GetInternalTree()->SetDirectory(gDirectory);
	if (nmc)  nmc->GetInternalTree()->SetDirectory(gDirectory);
	if(hchi2_4C) hchi2_4C->SetDirectory(gDirectory);
	// *** restore original gDirectory
	dir->cd();
	
	return kSUCCESS;
}

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

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


// -----   Public method Exec   --------------------------------------------
void PndScrutAnaTask::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%100)) 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 gamma, pi0, piplus, piminus,pichpair,hcfit, hc, allCharged, allNeutral, mclist;
	

	
	// *** Select with PID info pidalg and ('All'); type and mass are set 		
	fAnalysis->FillList(gamma,  "Neutral");
	fAnalysis->FillList(piplus,  "PionAllPlus",  pidalg);
	fAnalysis->FillList(piminus, "PionAllMinus",  pidalg);
	int ngamma = gamma.GetNumberOfTracks();
	int npipl = piplus.GetNumberOfTracks();
	int npimn = piminus.GetNumberOfTracks();

	if(ngamma==2 && npipl==1 && npimn==1){
	  // cout<<"gamma size = "<<gamma.GetNumberOfTracks()<<
	  //   " pi+ size =  "<<piplus.GetNumberOfTracks()<<
	  //   " pi- size =  "<<piminus.GetNumberOfTracks()<<endl;
	  // // *** Setup event shape object
	// fAnalysis->FillList(allCharged,"Charged", pidalg);
	// fAnalysis->FillList(allNeutral, "Neutral");
	//	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();
	
	// *** combinatorics for pi0 -> gamma gamma


	pi0.Combine(gamma, gamma);
// 	// for (j=0;j<gamma.GetLength();++j){
// 	//   for (k=1;k<gamma.GetLength();++k)
// 	//     {
// 	      RhoCandidate *combCand_pi0 = gamma[0]->Combine(gamma[1]);
// 	      // ... some computation or checks ...
// 	      //pt vs pz mean:
// 	      // p0                        =     0.878754   +/-   0.000508978 
// 	      // p1                        =     0.779844   +/-   0.00144043  
// 	      // p2                        =    -0.288569   +/-   0.00128464  
// 	      // p3                        =    0.0536217   +/-   0.000418099 
// 	      // p4                        =  -0.00504796   +/-   4.4216e-05
// 	      // width: < 0.2
	      
	      
// 	      double wid_pi0 = 0.2;
// 	      double ptcur_pi0 =  combCand_pi0->Pt();
// 	      double pzcur_pi0 = combCand_pi0->Pz();
// 	      double ptcalc_pi0 = 0.878754+0.779844*pzcur_pi0 -0.288569*pzcur_pi0*pzcur_pi0 +0.0536217*pzcur_pi0*pzcur_pi0*pzcur_pi0
// 		-0.00504796*pzcur_pi0*pzcur_pi0*pzcur_pi0*pzcur_pi0;
// 	      //  cout<<"fabs(ptcalc-ptcur) = "<<fabs(ptcalc-ptcur)<<endl;
// 	      if(fabs(ptcalc_pi0-ptcur_pi0)<wid_pi0)
// 		pi0.Append(combCand_pi0);
// 	      //     }
// 	      // }
	      pi0.SetType(111);

	// *** some rough mass selection on pi0
	// *** Mass selector for the pi0 cands
        double m0_pi0 = fPdg->GetParticle("pi0")->Mass();   // Get nominal PDG mass of the h_c
	double pi0MassWind = 0.05;//TEST
	//	double pi0MassWind = 5e6;//TEST
	RhoMassParticleSelector *pi0MassSel=new RhoMassParticleSelector("pi0",m0_pi0,pi0MassWind);
	pi0.Select(pi0MassSel);
	
	int njmct = fAnalysis->McTruthMatch(pi0); // match the whole list to count #matches (should be only 1)
			
	// *** write ntuple for the pi0 reconstruction
	for (j=0;j<pi0.GetLength();++j) 
	{
		// some general info about event (actually written for each candidate!)
		ntp1->Column("ev",		(Float_t) i);
		ntp1->Column("cand",	(Float_t) j);
		ntp1->Column("ncand",   (Float_t) pi0.GetLength());
		ntp1->Column("nmct",    (Float_t) njmct);
	
		// store info about initial 4-vector
		qa.qaP4("beam", fIni, ntp1);
			
		// dump information about composite candidate tree recursively (see PndTools/AnalysisTools/PndRhoTupleQA)
		qa.qaComp("pi0", pi0[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 = pi0[j]->GetMcTruth();
		TLorentzVector lv;
		if (truth) lv = truth->P4();
		qa.qaP4("trpi0", lv, ntp1);
		
		ntp1->DumpData();
	}
	
	
	
	// // clean pion pairs list and fill it with pi+ pi- combinations
	// // pichpair.Cleanup();
	// // //	cout<<"Number of pi+ :"<<piplus.GetLength()<<endl;
	// // //	cout<<"Number of pi- :"<<piminus.GetLength()<<endl;
	// for (j=0;j<piplus.GetLength();++j){
	//   for (k=0;k<piminus.GetLength();++k)
	//     {
	//       RhoCandidate *combCand = piplus[j]->Combine(piminus[k]);
	//       // ... some computation or checks ...
	//       //pt vs pz mean:
	//       // p0                        =     0.661097   +/-   0.00177429  
	//       // p1                        =     0.919447   +/-   0.00376505  
	//       // p2                        =    -0.314186   +/-   0.00261582  
	//       // p3                        =     0.053605   +/-   0.000713217 
	//       // p4                        =  -0.00455441   +/-   6.6113e-05  
	//       //width: < 0.1 GeV/c

	//       double wid = 0.1;
	//       double ptcur =  combCand->Pt();
	//       double pzcur = combCand->Pz();
	//       double ptcalc = 0.661097+0.919447*pzcur -0.314186*pzcur*pzcur +0.053605*pzcur*pzcur*pzcur
	// 	-0.00455441*pzcur*pzcur*pzcur*pzcur;
	//       //  cout<<"fabs(ptcalc-ptcur) = "<<fabs(ptcalc-ptcur)<<endl;
	//       if(fabs(ptcalc-ptcur)<wid)
	// 	pichpair.Append(combCand);
	//     }
	// }

	// // *** combinatorics for rho0 -> pi+ pi-
	pichpair.Combine(piplus, piminus);
	pichpair.SetType(113);
	// *** some rough mass selection on rho0
	// *** Mass selector for the pi0 cands
        double m0_rho0 = fPdg->GetParticle("rho0")->Mass();   // Get nominal PDG mass of the h_c
		double rhoMassWind = 0.3;//TEST
	//	double rhoMassWind = 3e6;
	RhoMassParticleSelector *rho0MassSel=new RhoMassParticleSelector("rho0",m0_rho0,rhoMassWind);
	pichpair.Select(rho0MassSel);


	// *** write ntuple for the hc reconstruction
	for (j=0;j<pichpair.GetLength();++j) 
	{
	  // double wid = 0.1;
	  // double ptcur =  pichpair[j]->Pt();
	  // double pzcur = pichpair[j]->Pz();
	  // double ptcalc = 0.661097+0.919447*pzcur -0.314186*pzcur*pzcur +0.053605*pzcur*pzcur*pzcur
	  //   -0.00455441*pzcur*pzcur*pzcur*pzcur;
	  // //  cout<<"fabs(ptcalc-ptcur) = "<<fabs(ptcalc-ptcur)<<endl;
	  // if(fabs(ptcalc-ptcur)<wid){
		// some general info about event (actually written for each candidate!)
		ntp2->Column("ev",		(Float_t) i);
		ntp2->Column("cand",	(Float_t) j);
		ntp2->Column("ncand",   (Float_t) pichpair.GetLength());
		
		// store info about initial 4-vector
		qa.qaP4("beam", fIni, ntp2);
		
		// dump information about composite candidate tree recursively (see PndTools/AnalysisTools/PndRhoTupleQA)
		qa.qaComp("pichpair", pichpair[j], ntp2);

		// // dump info about event shapes
		// 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 = pichpair[j]->GetMcTruth();
		TLorentzVector lv;
		if (truth) lv = truth->P4();
		qa.qaP4("trpichpair", lv, ntp2);
		
		ntp2->DumpData();
	}
	//	}			


	// // *** write ntuple for the hc reconstruction
	// for (j=0;j<pichpairfit.GetLength();++j) 
	// {
	// 	// some general info about event (actually written for each candidate!)
	// 	ntp2->Column("ev",		(Float_t) i);
	// 	ntp2->Column("cand",	(Float_t) j);
	// 	ntp2->Column("ncand",   (Float_t) pichpairfit.GetLength());
		
	// 	// store info about initial 4-vector
	// 	qa.qaP4("beam", fIni, ntp2);
		
	// 	// dump information about composite candidate tree recursively (see PndTools/AnalysisTools/PndRhoTupleQA)
	// 	qa.qaComp("pichpairfit", pichpairfit[j], ntp2);

	// 	// // dump info about event shapes
	// 	// 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 = pichpairfit[j]->GetMcTruth();
	// 	TLorentzVector lv;
	// 	if (truth) lv = truth->P4();
	// 	qa.qaP4("trpichpairfit", lv, ntp2);
		
	// 	ntp2->DumpData();
	// }			



	// *** combinatorics for hc_c -> (pi+ pi-)pi0
	hc.Combine(pichpair, pi0);		
	//	hc.Combine(pichpairfit, pi0);		
	hc.SetType(10443);

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

	// *** write ntuple for the hc reconstruction
	for (j=0;j<hc.GetLength();++j) 
	{
		// some general info about event (actually written for each candidate!)
		ntp3->Column("ev",		(Float_t) i);
		ntp3->Column("cand",	(Float_t) j);
		ntp3->Column("ncand",   (Float_t) hc.GetLength());
		ntp3->Column("nmct",    (Float_t) npsimct);
		
		// store info about initial 4-vector
		qa.qaP4("beam", fIni, ntp3);
		
		// dump information about composite candidate tree recursively (see PndTools/AnalysisTools/PndRhoTupleQA)
		qa.qaComp("hc", hc[j], ntp3);

		// // dump info about event shapes
		// 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 = hc[j]->GetMcTruth();
		TLorentzVector lv;
		if (truth) lv = truth->P4();
		qa.qaP4("trhc", lv, ntp3);
		
		ntp3->DumpData();
	}			

	// *** the lorentz vector of the initial h_c; important for the 4C-fit
	TLorentzVector ini(0, 0, 5.6088485, 6.6250557);
	// ... in event loop ...
	// *** do 4c fit
	for (j=0;j<hc.GetLength();++j) 
	  {
	    Pnd4CFitter fitter(hc[j],ini);     // *** instantiate the 4C fitter
	    fitter.FitConserveMasses();           // *** perform the fit
	    //   fitter.Fit();           // *** perform the fit with floating of the masses of the electrons and pions
	    double chi2_4c = fitter.GetChi2();    // *** access the chi2 of the fit
	    //	    cout<<"chi2_4c = "<<chi2_4c<<endl;
	    RhoCandidate *tmpfit = hc[j]->GetFit(); // *** access the fitted candidate            
	    //    if (chi2_4c<20 && tmpfit)               // *** when good enough, fill some histos
	    if (tmpfit){               // *** when good enough, fill some histos
	      hchi2_4C->Fill(chi2_4c);
	      double chi2Wind = 1e2;//TEST
	      //   double chi2Wind = 1e9;//TEST
	      if(chi2_4c<chi2Wind)
		hcfit.Add(tmpfit);
	    }
	  }      

	// *** write ntuple for the hc fit
	for (j=0;j<hcfit.GetLength();++j) 
	{
		// some general info about event (actually written for each candidate!)
		ntp4->Column("ev",		(Float_t) i);
		ntp4->Column("cand",	(Float_t) j);
		ntp4->Column("ncand",   (Float_t) hcfit.GetLength());
		ntp4->Column("nmct",    (Float_t) npsimct);
		
		// store info about initial 4-vector
		qa.qaP4("beam", fIni, ntp4);
		
		// dump information about composite candidate tree recursively (see PndTools/AnalysisTools/PndRhoTupleQA)
		qa.qaComp("hcfit", hcfit[j], ntp4);

		// // dump info about event shapes
		// 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 = hcfit[j]->GetMcTruth();
		TLorentzVector lv;
		if (truth) lv = truth->P4();
		qa.qaP4("trhc", lv, ntp4);
		ntp4->DumpData();
	}			
	}


	
}


void PndScrutAnaTask::Finish()
{
	
	// *******
	// ******* STORE YOUR HISTOS AND TUPLES
	// *******
	
	fFile->Write();
	fFile->Close();	

}

ClassImp(PndScrutAnaTask)