//This class header
#include "KaonTask.h"
#include "Kaon.h"
#include "Pion.h"

#include "FairRootManager.h"
#include "FairRun.h"
#include "FairRuntimeDb.h"
#include "TpcDigiPar.h"

#include "GFException.h"
#include "GFTrack.h"
#include "GFTrackProximity.h"
#include "RKTrackRep.h"
#include "TLorentzVector.h"
#include "TVector3.h"
#include "TDatabasePDG.h"
#include "PndMCTrack.h"

#include <iostream>

KaonTask::~KaonTask()
{;}

KaonTask::KaonTask()
:fOutBranchName("TpcKaons"),
 fTpcHitBranchName("TpcSPHit"),
 fTpcClusterBranchName("TpcCluster"),
 fTpcTrackBranchName("TrackPostFit"),
 fParticleBranchName("Particles"),
 fMCTrackBranchName("MCTrack"),
 fPersistence(kFALSE)
{

}

InitStatus KaonTask::Init()
{

	//Get ROOT Manager ------------------
	FairRootManager* ioman= FairRootManager::Instance();
	if(ioman==0)
	{
		Error("Init","RootManager not instantiated!");
		return kERROR;
	}

	fTpcTrackArray=(TClonesArray*) ioman->GetObject(fTpcTrackBranchName);
	if(fTpcTrackArray==0)
	{
		std::cout<<"KaonTask::Init TpcTrackArray not found "<<fTpcTrackBranchName<<"\n";
		return kERROR;
	}


	fTpcHitArray=(TClonesArray*) ioman->GetObject(fTpcHitBranchName);
	if(fTpcHitArray==0)
	{
		std::cout<<"KaonTask::Init TpcHitArray not found ("<<fTpcHitBranchName<<")\n";
		return kERROR;
	}

	fTpcClusterArray=(TClonesArray*) ioman->GetObject(fTpcClusterBranchName);
	if(fTpcClusterArray==0)
	{
		std::cout<<"KaonTask::Init TpcClusterArray not found ("<<fTpcClusterBranchName<<")\n";
		return kERROR;
	}

	fParticleArray=(TClonesArray*) ioman->GetObject(fParticleBranchName);
	if(fParticleArray==0)
	{
		std::cout<<"KaonTask::Init ParticleArray not found ("<<fParticleBranchName<<")\n";
		return kERROR;
	}

	fMCTrackArray=(TClonesArray*) ioman->GetObject(fMCTrackBranchName);
	if(fMCTrackArray==0)
	{
		std::cout<<"KaonTask::Init MCTrackArray not found ("<<fMCTrackBranchName<<")\n";
		return kERROR;
	}

	fOutArray = new TClonesArray("Kaon");
	TString title = fOutBranchName;
	ioman->Register(title,"Kaon",fOutArray,fPersistence);

	std::cout<<"INIT OUT ARRAYS\n";
	return kSUCCESS;
}

void KaonTask::SetParContainers()
{

	std::cout<<"CollectParticlesTask::SetParContainers\n";
	std::cout.flush();

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

	FairRuntimeDb* db = run->GetRuntimeDb();
	if ( ! db ) Fatal("SetParContainers", "No runtime database");

	// Get Tpc digitisation parameter container
	fPar= (TpcDigiPar*) db->getContainer("TpcDigiPar");
	if (!fPar ) Fatal("SetParContainers", "TpcDigiPar not found");

	fPDG=TDatabasePDG::Instance();
}

void KaonTask::Exec(Option_t* opt)
{
	if(fVerbose)
	{
		std::cout<<"KaonTask::Exec: starting kaon task exec\n";
	}
	FairRootManager* ioman= FairRootManager::Instance();

	fOutArray->Delete();

	Int_t npart=fParticleArray->GetEntriesFast();
	Int_t nMCTracks=fMCTrackArray->GetEntriesFast();
	if(fVerbose)
		std::cout<<"there are "<<npart<<" pions\n";
	if (npart!=2)
	{return;}
	if(nMCTracks!=3)
	{return;}

	Kaon* kaon=new Kaon();
	Pion pion1(*((Particle*)(fParticleArray->At(0))));
	Pion pion2(*((Particle*)(fParticleArray->At(1))));

	kaon->appendPion(pion1);
	kaon->appendPion(pion2);

	if( pion1.getNumHits()==0 or pion2.getNumHits()==0)
	{
		if(fVerbose){std::cout<<"One or two pions has 0 hits, skipping\n";}
		delete kaon;
		return;
	}

	TLorentzVector l_pion1, l_pion2;
	TVector3 pocaTr, pion1_poca, pion2_poca;

	Bool_t gotPoca = Poca((GFTrack*)fTpcTrackArray->At(pion1.getTrackID()),
			              (GFTrack*)fTpcTrackArray->At(pion2.getTrackID()),
			              l_pion1, l_pion2,
			              pocaTr, pion1_poca, pion2_poca);

	if (gotPoca)
	{
		kaon->appendPionPoca(pion1_poca);
		kaon->appendPionPoca(pion2_poca);
		kaon->setTrackPoca(pocaTr);
	}


	TVector3 pionVertp, pionVertm;
	bool foundPim=false;
	bool foundPip=false;
	bool foundK0s=false;
	for(int imctr=0;imctr<nMCTracks;imctr++)
	{
		PndMCTrack* mcTrack=(PndMCTrack*)fMCTrackArray->At(imctr);
		if(fVerbose){std::cout<<"KaonTask::Exec: checking MC particle ("<<mcTrack->GetPdgCode()<<")\n";}
		if(mcTrack->GetPdgCode()==211)
		{
			pionVertp=mcTrack->GetStartVertex();
			if(fVerbose){std::cout<<"KaonTask::Exec: found a pion at ("<<pionVertp.X()<<","<<pionVertp.Y()<<","<<pionVertp.Z()<<")\n";}
			foundPip=true;
		}
		if(mcTrack->GetPdgCode()==-211)
		{
			pionVertm=mcTrack->GetStartVertex();
			if(fVerbose){std::cout<<"KaonTask::Exec: found a pion at ("<<pionVertm.X()<<","<<pionVertm.Y()<<","<<pionVertm.Z()<<")\n";}
			foundPim=true;
		}
		if(mcTrack->GetPdgCode()==310)
		{
			foundK0s=true;
		}
	}
	if( not(foundPip and foundPim and foundK0s))
	{
		delete kaon;
		return;
	}

	kaon->setMCTrackPoca(0.5*(pionVertp+pionVertm));
	unsigned int nFS = fOutArray->GetEntriesFast();
	new((*fOutArray)[nFS]) Kaon(*kaon);
	delete kaon;
}

void KaonTask::getTrackPosMom(GFAbsTrackRep* _rep, TVector3& _pos, TVector3& _poserr, TVector3& _mom, TVector3& _momerr)
{
	TVector3 poca;
	TVector3 mom;
	try
	{
		_rep->extrapolateToPoint(_pos,poca,mom);
	}
	catch(GFException& e)
	{
		std::cerr<<"CollectParticlesTask::getTrackPosMom():could not extrapolate to point\n";
		std::cerr << e.what();
		return;
	}
	return;
}

bool KaonTask::Poca(GFTrack* pionp,
		GFTrack* pionm,
		TLorentzVector &l_pionp,
		TLorentzVector &l_pionm,
		TVector3& pocaTr,
		TVector3& pocaPionP,
		TVector3& pocaPionM)
{
	if(fVerbose) std::cout<<"KaonTask::Poca\n";
	RKTrackRep*  pionp_rep   =new RKTrackRep(*((RKTrackRep* )pionp->getCardinalRep()));
	RKTrackRep*  pionm_rep   =new RKTrackRep(*((RKTrackRep* )pionm->getCardinalRep()));
	TVector3 poca1,norm1;
	TVector3 poca2,norm2;
	GFDetPlane fPlane1,fPlane2;
	bool passed =false;
	try{
		TVector3 startPoint1=pionp_rep->GFAbsTrackRep::getPos();
		TVector3 startPoint2=pionm_rep->GFAbsTrackRep::getPos();

		TVector3 dist1=startPoint2-startPoint1;

		//vtxResX_1->Fill(dist1.X());
		//vtxResY_1->Fill(dist1.Y());
		//vtxResZ_1->Fill(dist1.Z());
		if(fVerbose)
		{
			startPoint1.Print();
			startPoint2.Print();
		}
		pocaTr=trackProximity((GFAbsTrackRep*)pionp_rep,(GFAbsTrackRep*)pionm_rep);

		pionp_rep->extrapolateToPoint(pocaTr, poca1, norm1);
		fPlane1.setO(pocaTr);
		fPlane1.setNormal(norm1);
		pionm_rep->extrapolateToPoint(pocaTr, poca2, norm2);
		fPlane2.setO(pocaTr);
		fPlane2.setNormal(norm2);

		TVector3 mom1=pionp_rep->getMom(fPlane1);
		TVector3 mom2=pionm_rep->getMom(fPlane2);
		l_pionp.SetVectM(mom1,fPDG->GetParticle(2212)->Mass());
		l_pionm.SetVectM(mom2,fPDG->GetParticle(211)->Mass());

		startPoint1=pionp_rep->GFAbsTrackRep::getPos();
		startPoint2=pionm_rep->GFAbsTrackRep::getPos();

		//vtxXY->Fill(pocaTr.X(),pocaTr.Y());
		//vtxZX->Fill(pocaTr.Z(),pocaTr.X());
		//vtxZY->Fill(pocaTr.Z(),pocaTr.Y());

		dist1=startPoint2-startPoint1;
		pocaPionP=poca1;
		pocaPionM=poca2;
		//vtxResX_2->Fill(dist1.X());
		//vtxResY_2->Fill(dist1.Y());
		//vtxResZ_2->Fill(dist1.Z());

		passed=true;
		delete pionp_rep;
		delete pionm_rep  ;
	}
	catch(GFException &exp){
		std::cout<<exp.what()<<std::endl;
		std::cout<<"KaonTask::POCA could not find a POCA\n";
	}
	catch(...){
		std::cout<<"KaonTask::POCA could not do anything\n";
	}
	return passed;
}


ClassImp(KaonTask)