//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      TPC-CDC Matching Routine
//
//
// Environment:
//      Software developed for the Prototype Detector at FOPI
//
// Author List:
//      Felix Boehmer    (TUM)
//
//-----------------------------------------------------------



// This Class' Header ------------------
#include "TpcCdcMatchingTask.h"

// C++ headers
#include <iostream>
#include <vector>
#include <cmath>
#include <algorithm>


// Collaborating Class Headers --------
#include "FairRootManager.h"
#include "FairRun.h"
#include "FairRunAna.h"
#include "FairRuntimeDb.h"
#include "FairField.h"
#include "TpcDigiPar.h" 
#include "GFAbsTrackRep.h"
#include "RecoHits/GFAbsRecoHit.h"
#include "GFDetPlane.h"
#include "GFTrack.h"
#include "GFTools.h"
#include "CdcTrack.h"
#include "CdcHit.h"
#include "RKTrackRep.h"
#include "CdcCircle.h"

#include "MatchingPair.h"
#include "MatchingCriterion.h"

#include "TVector3.h"
#include "TMatrixT.h"
#include "TMath.h"


TpcCdcMatchingTask::TpcCdcMatchingTask()
: fTpcTrackBranchName("TpcTrackPreFit"),
  fCdcTrackBranchName("FopiCdcTrack"),
  fGFCdcTrackBranchName("GFCdcTrack"),
  fCdcHitBranchName("FopiCdcHit"),
  fOutBranchName("TpcCdcMatchingPair"),
  fPersistence(kTRUE),
  fUseGFCdc(kFALSE),
  //fNumReps(1),
  fWriteFailedMatches(kFALSE)
{;}

TpcCdcMatchingTask::~TpcCdcMatchingTask() {
	;
}


InitStatus
TpcCdcMatchingTask::Init() {

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

	// Get input ------------------------
	fTpcTrackArray=(TClonesArray*) ioman->GetObject(fTpcTrackBranchName);
	if(fTpcTrackArray==0) {
		Error("TpcCdcMatchingTask::Init","GFTrack array not found!");
		return kERROR;
	}
	fCdcHitArray=(TClonesArray*) ioman->GetObject(fCdcHitBranchName);
	if(fCdcHitArray==0) {
		Error("TpcCdcMatchingTask::Init","CDC hits array not found!");
		return kERROR;
	}
	fCdcTrackArray=(TClonesArray*) ioman->GetObject(fCdcTrackBranchName);
	if(fCdcTrackArray==0) {
		Error("TpcCdcMatchingTask::Init","CDC track array not found!");
		return kERROR;
	}
	if(fUseGFCdc) {
		fGFCdcTrackArray = (TClonesArray*) ioman->GetObject(fGFCdcTrackBranchName);
		if(fGFCdcTrackArray==0) {
			Error("TpcCdcMatchingTask::Init","GF CDC track array not found!");
			return kERROR;
		}
	}

	fMatchingOutArray = new TClonesArray("MatchingPair");
	ioman->Register(fOutBranchName,"Tpc",fMatchingOutArray,fPersistence);

	//fTpcId = FairRootManager::Instance()->GetBranchId(fTpcClusterBranchName);
	fCdcId = FairRootManager::Instance()->GetBranchId(fCdcHitBranchName);

	return kSUCCESS;
}

void 
TpcCdcMatchingTask::SetParContainers() {

	std::cout<<"TpcCdcMatchingTask::SetParContainers"<<std::endl;
	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");
}

TString genError(const MatchingCriterion& mc) {
	TString err("Criterion with name ");
	err += mc.getName();
	err += " already exists";
	return err;
}

bool radialCdcHitSort(CdcHit* cdc1, CdcHit* cdc2) {
	return cdc1->GetHitPos().Perp() < cdc2->GetHitPos().Perp();
}

bool radialGFRecoHitSort(GFAbsRecoHit* hit1, GFAbsRecoHit* hit2) {
	TVectorD hitPos(3);
	hitPos = hit1->getRawHitCoord();
	TVector3 hitVec1(hitPos[0], hitPos[1], hitPos[2]);
	hitPos = hit2->getRawHitCoord();
	TVector3 hitVec2(hitPos[0], hitPos[1], hitPos[2]);
	return hitVec1.Perp() < hitVec2.Perp();
}

void
TpcCdcMatchingTask::Exec(Option_t* opt) {

	if(fMatchingOutArray==NULL)
		Fatal("TpcCdcMatchingTask::Exec()","No output array");
	fMatchingOutArray->Clear();
	if(fVerbose){
		std::cout<<"TpcCdcMatchingTask::Exec() ... "<<std::endl;
	}
	std::vector<MatchingPair*> pairs;
	//avoid rebuilding Cdc hit-vectors for evert TPC track:
	std::map<unsigned int, std::vector<CdcHit*>*> cdcHitMap;

	unsigned int nCdcTracks = fCdcTrackArray->GetEntriesFast();
	unsigned int nTpcTracks = fTpcTrackArray->GetEntriesFast();

	//loop over TPC tracks:
	for(unsigned int iTpcTr=0; iTpcTr<nTpcTracks; iTpcTr++) {
		GFTrack* iTpc = (GFTrack*) (*fTpcTrackArray)[iTpcTr];
		if(iTpc->getCardinalRep()->getStatusFlag())
			continue;
		TVector3 tpcStartPos = iTpc->getCardinalRep()->getPos();
		TVector3 tpcStartMom = iTpc->getCardinalRep()->getMom();
		double tpcTheta = tpcStartMom.Theta()*180/TMath::Pi();
		std::vector<GFAbsRecoHit*> tpcHits = iTpc->getHits();
		//radial sorting, could be needed for some track geometries
		std::sort(tpcHits.begin(),tpcHits.end(),radialGFRecoHitSort);
		TVector3 lastMom;
		if(iTpc->getSmoothing()) {
			unsigned int cardRepID = iTpc->getCardinalRepID();
			std::map<GFAbsRecoHit*, unsigned int> hitMap;
			iTpc->getHitMap(hitMap);
			unsigned int lastHitID = hitMap[(tpcHits[tpcHits.size()-1])];
			try{
				lastMom = GFTools::getSmoothedMomXYZ(iTpc,cardRepID,lastHitID);
			}catch(GFException& exp){
				std::cerr<<"TpcCdcMatchingTask::Exec() error getting getSmoothedMomXYZ"<<std::endl;
				std::cerr<<exp.what()<<std::endl;
			}
		}
		else {
			GFDetPlane lP = iTpc->getCardinalRep()->getLastPlane();
			TVector3 dummy;
			TMatrixDSym dummy2;
			try {
				iTpc->getCardinalRep()->getPosMomCov(lP,dummy,lastMom,dummy2);
			}
			catch (GFException& ex) {
				if(ex.isFatal()) {
					std::cout<<ex.what()<<std::endl;
					continue;
				}
			}
		}


		//loop over Cdc tracks:
		for(unsigned int iCdcTr=0; iCdcTr<nCdcTracks; iCdcTr++) {
			CdcTrack* iCdc = (CdcTrack*) (*fCdcTrackArray)[iCdcTr];
			if(iCdc->GetCdcIdHits().size()<1)
				continue;
			TVector3 circCent(iCdc->GetMx(),iCdc->GetMy(),0.);
			double circRad = iCdc->GetRadius();

			//create MatchingPair object for this combination:
			MatchingPair* thePair = new MatchingPair(iTpcTr, fTpcTrackBranchName,
					iCdcTr, fCdcTrackBranchName);
			pairs.push_back(thePair);
			double cdcMom = iCdc->GetMom();




			//loop over TPC hits and calculate MEAN TRACK HIT RESIDUAL to CDC circle (in X-Y) ----
			//requires complete fit
			double sumRes=0;
			for(unsigned int iTpcHit=0; iTpcHit<tpcHits.size(); iTpcHit++) {
				GFAbsRecoHit* iHit = tpcHits[iTpcHit];
				TVectorD hitPos(3);
				hitPos = iHit->getRawHitCoord();
				TVector3 hitVec(hitPos[0], hitPos[1], 0.);
				TVector3 tpcRadVec = circCent - hitVec;
				sumRes += tpcRadVec.Mag() - circRad;
			}

			MatchingCriterion meanRes("meanResidualXY", sumRes/((double)tpcHits.size()));
			//meanRes.setCuts(-1.,1.);
			if(thePair->addCriterion(meanRes))
				Fatal("TpcCdcMatchingTask::Exec", genError(meanRes));
			// -----------------------------------------------------------------------------------


			// TRACK TRANSITION PHI RESIDUAL at border TPC-CDC ---------------------------------------
			std::vector<CdcHit*>* cdcHits;
			if(!cdcHitMap.count(iCdcTr)) { //extract CdcHits from TCA and store
				std::vector<unsigned int> cdcHitIDs = iCdc->GetCdcIdHits();
				cdcHitMap[iCdcTr] = new std::vector<CdcHit*>();
				for(unsigned int iId=0; iId<cdcHitIDs.size(); iId++)
					(cdcHitMap[iCdcTr])->push_back((CdcHit*)fCdcHitArray->At(cdcHitIDs[iId]));
				//radial sorting (no clue how the CdcTracks are ordered to begin with ...)
				std::sort((cdcHitMap[iCdcTr])->begin(),(cdcHitMap[iCdcTr])->end(),radialCdcHitSort);
			}
			cdcHits = cdcHitMap[iCdcTr];
			TVectorD lastTpcHitPos(3);
			lastTpcHitPos = (tpcHits[tpcHits.size()-1])->getRawHitCoord();
			TVector3 lastTpcHitVec(lastTpcHitPos[0], lastTpcHitPos[1], lastTpcHitPos[2]);
			TVector3 firstCdcHitVec = cdcHits->at(0)->GetHitPos();
			double dPhiHits = (lastTpcHitVec.Phi() - firstCdcHitVec.Phi())*180./TMath::Pi();
			if(dPhiHits > 180.)
				dPhiHits -= 360.;
			if(dPhiHits < -180.)
				dPhiHits += 360.;
			MatchingCriterion dPhi("dPhiHitTransition", dPhiHits);
			dPhi.setCuts(-15.,15.);
			if(thePair->addCriterion(dPhi))
				Fatal("TpcCdcMatchingTask::Exec", genError(dPhi));

			// Z-RESIDUAL (HITS)
			// remove theta dependence introduced by distance between CDC and TPC
			// xy-distance of TPC and CDC point
			TVector3 xyd = lastTpcHitVec - firstCdcHitVec;
			xyd.SetZ(0.);
			//along phi:
			TVector3 testvec(1.,0.,0.);
			testvec.SetPhi(lastTpcHitVec.Phi());
			testvec.SetMag(1.);
			double scale = fabs(xyd*testvec);
			//TVector3 xyds = scale * xyd;
			double corr = scale / TMath::Tan(lastMom.Theta());

			double dz = lastTpcHitVec.Z() - firstCdcHitVec.Z() + corr;

			MatchingCriterion dZ("dZHits", dz);
			if(thePair->addCriterion(dZ))
				Fatal("TpcCdcMatchingTask::Exec", genError(dZ));
			// -----------------------------------------------------------------------------------


			// TRACK TANGENTIAL RESIDUAL at border TPC-CDC ---------------------------------------
			TVector3 radVec =  circCent - lastTpcHitVec;
			radVec.SetZ(0.);
			radVec.RotateZ(90*TMath::Pi()/180.);
			//check if the sign is correct:
			TVector3 approxTrack = firstCdcHitVec - lastTpcHitVec;
			approxTrack.SetZ(0.);
			if(approxTrack*radVec < 0)
				radVec.RotateZ(TMath::Pi());
			double dPhiTang = (lastMom.Phi()-radVec.Phi())*180./TMath::Pi();
			if(dPhiTang > 180.)
				dPhiTang -= 360.;
			if(dPhiTang < -180.)
				dPhiTang += 360.;
			MatchingCriterion dPhiTrack("dPhiTrackTransition", dPhiTang);
			dPhiTrack.setCuts(-10.,10.);
			if(thePair->addCriterion(dPhiTrack))
				Fatal("TpcCdcMatchingTask::Exec", genError(dPhiTrack));

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


			// THETA RESIDUAL --------------------------------------------------------------------
			double cdcTheta = iCdc->GetTheta()*180./TMath::Pi();
			MatchingCriterion dThetaTrack("dThetaTrack", tpcTheta-cdcTheta);
			//dThetaTrack.setCuts(-20.,20.);
			if(thePair->addCriterion(dThetaTrack))
				Fatal("TpcCdcMatchingTask::Exec", genError(dThetaTrack));
			// -----------------------------------------------------------------------------------

			// CHARGE RESIDUAL -------------------------------------------------------------------
			double tpcCharge = iTpc->getCardinalRep()->getCharge();
			double cdcCharge = iCdc->GetCharge();
			//cdcCharge *= 1./std::fabs(cdcCharge);
			MatchingCriterion dCharge("dCharge", tpcCharge - cdcCharge);
			//dCharge.setCuts(-0.5,0.5);
			if(thePair->addCriterion(dCharge))
				Fatal("TpcCdcMatchingTask::Exec", genError(dCharge));
			// -----------------------------------------------------------------------------------



			//MOMENTUM RESIDUAL ------------------------------------------------------------------
			MatchingCriterion dMom("dMom", tpcStartMom.Mag() - cdcMom);
			if(thePair->addCriterion(dMom))
				Fatal("TpcCdcMatchingTask::Exec", genError(dMom));

			MatchingCriterion dMomBorder("dMomBorder", lastMom.Mag() - cdcMom);
			if(thePair->addCriterion(dMomBorder))
				Fatal("TpcCdcMatchingTask::Exec", genError(dMomBorder));


			//TRANSVERSE MOMENTUM RESIDUAL
			//just an approximation: use circle center
			double bZ = FairRunAna::Instance()->GetField()->GetBz(circCent.X(),
					circCent.Y(),
					circCent.Z());
			double ptCdc = std::fabs(bZ/10.*0.3*circRad/100.);
			TVector3 tpcPt = tpcStartMom;
			tpcPt.SetZ(0.);
			MatchingCriterion dMomT("dMomT", tpcPt.Mag()-ptCdc);
			if(thePair->addCriterion(dMomT))
				Fatal("TpcCdcMatchingTask::Exec", genError(dMomT));

			TVector3 tpcPtB = lastMom;
			tpcPtB.SetZ(0.);
			MatchingCriterion dMomTBorder("dMomTBorder", tpcPtB.Mag() - ptCdc);
			if(thePair->addCriterion(dMomTBorder))
				Fatal("TpcCdcMatchingTask::Exec", genError(dMomTBorder));


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

			//Define what is *suggested* to be matched
			bool matched = true;
			unsigned int nCuts = 0;
			std::vector<TString> crits = thePair->getListOfCritNames();
			for(unsigned int iCr=0; iCr<crits.size(); iCr++) {
				TString iName = crits[iCr];
				if(!(thePair->getCriterion(iName).isConstrained()))
					continue;
				nCuts++;
				std::pair<double,double> iCuts = thePair->getCriterion(iName).getCuts();
				double val = thePair->getCriterion(iName).getValue();
				if(val < iCuts.first || val > iCuts.second)
					matched = false;
			}
			if(nCuts == 0)
				matched = false; //default is false, even with no cuts set
			thePair->setMatched(matched);

		} //end loop over CDC tracks

	}  //end loop over TPC tracks

	//clean up trackMap:
	std::map<unsigned int, std::vector<CdcHit*>*>::iterator it;
	for(it=cdcHitMap.begin(); it!=cdcHitMap.end(); it++)
		delete it->second;  //ownership of the hits themselves still lies with the RunManager


	//decide which pairs to write out
	unsigned int countW = 0;
	for(unsigned int iPair=0; iPair<pairs.size(); iPair++) {
		MatchingPair* ip = pairs[iPair];
		if(ip->isMatched() || fWriteFailedMatches) {
			int nOut = fMatchingOutArray->GetEntriesFast();
			new ((*fMatchingOutArray)[nOut]) MatchingPair(*ip);
			countW++;
		}
		delete ip;
	}
	if(fVerbose)
	{
		std::cout<<"    ... suggested "<<countW<<" matches for "
				<<nTpcTracks<<" TPC tracks."<<std::endl;
	}
}

ClassImp(TpcCdcMatchingTask)