// ************************************************************************
//
// 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
	// *******
	
	// *** 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);

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

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

	// assign RhoTuples to outfile

	if (nmc)  nmc->GetInternalTree()->SetDirectory(gDirectory);
	if (hlm) hlm->SetDirectory(gDirectory);
	if (hlms) hlms->SetDirectory(gDirectory);
	if (hlmf) hlmf->SetDirectory(gDirectory);
	if (hb) hb->SetDirectory(gDirectory);
	//if (hbf) hbf->SetDirectory(gDirectory);
	if (hchisq) hchisq->SetDirectory(gDirectory);
	if (hlpm) hlpm->SetDirectory(gDirectory);
	if (hlpmf) hlpmf->SetDirectory(gDirectory);
	if (hlmm) hlpm->SetDirectory(gDirectory);
	if (hlmmf) hlpmf->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 pplus, pminus, piplus, piminus, kaon, kplus, kminus, all, mclist, lambdac, lambdacb, gamma, beam;
	
	
	// *** 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();
	
	lambdac.Combine(pplus,kminus,piplus);

	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();
	    if (mfitter.GetProb()>0.01)
	      {
		TLorentzVector lvf(lambdac[j]->GetFit()->GetMomentum(),lambdac[j]->GetFit()->GetEnergy());
		TLorentzVector lmissingf = fIni - lvf;
		hlmmf->Fill(lmissingf.M());
		hlpmf->Fill(lambdac[j]->GetFit()->M());
	      }
	  }
	
	lambdacb.Combine(pminus,gamma);
	
	lambdac.Select(lpMassSel);
	lambdacb.Select(lmMassSel);
	
	beam.Combine(lambdac,lambdacb);
	
	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());
	      }
	    //	    delete fitter;
	  }
	
}


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

}

ClassImp(PndScrutAnaTask)