#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", "ppbar analysis");
//	nmc  = new RhoTuple("nmc",  "mctruth info");

	// assign RhoTuples to outfile
	if (ntp1) ntp1->GetInternalTree()->SetDirectory(gDirectory);
//	if (nmc)  nmc->GetInternalTree()->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;
	
	// *** 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 kplusList, kminusList, phiList, gammaList, etaList, all;
	
	
	// *** Select with PID info pidalg and ('All'); type and mass are set 		
	fAnalysis->FillList(kplusList,  "KaonAllPlus",  pidalg);
	fAnalysis->FillList(kminusList, "KaonAllMinus", pidalg);
	fAnalysis->FillList(gammaList, "Neutral", pidalg);

	
	// *** 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();
	
	phiList.Combine(kplusList, kminusList);
	etaList.Combine(gammaList, gammaList);
	all.Combine(phiList, phiList, etaList);

	if(0 == all.GetLength()) return;
			
	// *** combinatorics for FULL event: pbar p -> Phi Phi Eta
	double bestChi2 = 9e6;
	int bestCand = -1;
	double bestPull = 9e6;

	for (i=0;i<all.GetLength();++i) 
	{
		PndKinFitter Fit4c(all[i]);
		Fit4c.Add4MomConstraint(fIni);
		Fit4c.Fit();

		if(Fit4c.GetChi2()<bestChi2) {
			bestChi2 = Fit4c.GetChi2();
			bestPull = Fit4c.GetPull();
			bestCand = i;
		} else continue;
	}			
	
	// some general info about event - only written for best candidate!
	ntp1->Column("ev", fEvtCount);
	ntp1->Column("cand", bestCand);
	ntp1->Column("ncand", all.GetLength());
//	ntp1->Column("nmct", npsimct);
	
	RhoCandidate* etaCand = all[bestCand]->Daughter(2);
	PndKinFitter etaFit(etaCand);
	etaFit.AddMassConstraint(0.547853);
	etaFit.Fit();

	RhoCandidate *etaFitted=all[bestCand]->Daughter(2)->GetFit();
	RhoCandidate *allFitted=all[bestCand]->GetFit();

	// store info about initial 4-vector
	qa.qaP4("beam", fIni, ntp1);
		
	// dump information about composite candidate tree recursively (see PndTools/AnalysisTools/PndRhoTupleQA)
	qa.qaComp("ppb", all[bestCand], ntp1);
	qa.qaComp("fEta", etaFitted, ntp1);
	qa.qaComp("f4C", allFitted, ntp1);
	ntp1->Column("fchi2Eta", etaFit.GetChi2());
	ntp1->Column("fchi24c",  bestChi2);
	ntp1->Column("fpull4c",  bestPull);
	ntp1->Column("fpullEta",  etaFit.GetPull());

	// 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 = psi2s[j]->GetMcTruth();
//	TLorentzVector lv;
//	if (truth) lv = truth->P4();
//	qa.qaP4("trpsi", lv, ntp2);
		
	ntp1->DumpData();
}


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

}

ClassImp(PndScrutAnaTask)