//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Track-merging based on MatchingTuples
//      Sept. 2012
//
//
// Environment:
//      Software developed for the Prototype Detector at FOPI
//
// Author List:
//      Felix Boehmer     (TUM)
//
//
//-----------------------------------------------------------

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

#include <vector>
#include <cmath>

#include "TpcDigiPar.h"
#include "MatchingTuple.h"
#include "CdcTrack.h"
#include "CdcSpacepoint.h"
#include "PseudoSpacePoint.h"

#include "TClonesArray.h"
#include "GFTrack.h"
#include "RecoHits/GFAbsRecoHit.h"
#include "GFException.h"

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


FopiTrackMergingTask::FopiTrackMergingTask()
: fRealData(true),
  fTpcTrackBranchName("TpcTrack"),
  fCdcTrackBranchName("CdcTrack"),
  fCdcGFTrackBranchName("CdcTrackPostFit"),
  fOutBranchName("FopiTrackPreFit"),
  fTupleOutBranchName("FopiTrackTuples"),
  fPersistence(kTRUE),
  fPrefitPersistence(kTRUE)
{
	fTupleBranchMap.clear();
}

FopiTrackMergingTask::~FopiTrackMergingTask() { ; }

InitStatus FopiTrackMergingTask::Init() {

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

	// Get input ------------------------
	fTpcTrackArray = (TClonesArray *)ioman->GetObject(fTpcTrackBranchName);
	if (fTpcTrackArray == 0) {
		Error("FopiTrackMergingTask::Init", "GFTrack array not found!");
		return kERROR;
	}
	fCdcTrackArray = (TClonesArray *)ioman->GetObject(fCdcTrackBranchName);
	if (fCdcTrackArray == 0) {
		Error("FopiTrackMergingTask::Init", "CDC track array not found!");
		return kERROR;
	}
	fCdcGFTrackArray = (TClonesArray *)ioman->GetObject(fCdcGFTrackBranchName);
	if (fCdcGFTrackArray == 0) {
		Error("FopiTrackMergingTask::Init", "CDC GFTrack array not found!");
		return kERROR;
	}

	// attach tuple branches
	std::map<TString, TClonesArray *>::iterator it;
	for (it = fTupleBranchMap.begin(); it != fTupleBranchMap.end(); it++) {
		it->second = (TClonesArray *)ioman->GetObject(it->first);
		if (it->second == NULL) {
			TString err("Branch '");
			err += it->first;
			err += "' not existing in tree!";
			Fatal("FopiTrackMergingTask::Init", err);
		}
	}

	fTrackOutArray = new TClonesArray("GFTrack");
	ioman->Register(fOutBranchName, "Tpc", fTrackOutArray, fPrefitPersistence);

	fTupleOutArray = new TClonesArray("MatchingTuple");
	ioman->Register(fTupleOutBranchName, "Tpc", fTupleOutArray, fPersistence);

	return kSUCCESS;
}

void FopiTrackMergingTask::SetParContainers() {

	//  std::cout<<"TpcResidualTask::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");
}

void FopiTrackMergingTask::Exec(Option_t *opt) {
	if (fTrackOutArray == NULL)
		Fatal("FopiTrackMergingTask::Exec()", "No output array");
	fTrackOutArray->Delete();

	if (fTupleOutArray == NULL)
		Fatal("FopiTrackMergingTask::Exec()", "No tuple output array");
	fTupleOutArray->Delete();

	if(fVerbose)
	{
		std::cout << "FopiTrackMergingTask::Exec() ... \n";
	}
	// create map of MatchingTuples based on TPC trackID:
	std::map<unsigned int, std::vector<MatchingTuple *> > tupmapTpc;
	// memorize all tpc track IDs seen during looping:
	std::set<unsigned int> tpcIdsSeen;

	std::map<TString, TClonesArray *>::iterator it;
	// loop over all associated tuple branches
	for (it = fTupleBranchMap.begin(); it != fTupleBranchMap.end(); it++) {
		TClonesArray *tupArr = it->second;
		unsigned int nTuples = tupArr->GetEntriesFast();
		for (unsigned int iT = 0; iT < nTuples; iT++) {
			MatchingTuple *ituple = (MatchingTuple *)tupArr->At(iT);
			// decide if tuple belongs to TPC:
			if (ituple->hasBranch(fTpcTrackBranchName)) {
				tpcIdsSeen.insert(ituple->getIdFromBranch(fTpcTrackBranchName));
				if (ituple->isMatched()) // only consider pre-matched tuples for MERGING
					(tupmapTpc[ituple->getIdFromBranch(fTpcTrackBranchName)])
					.push_back(ituple);
			}
		}
	}

	// CDC-TPC matching part anchored on TPC, using FOPI CDC tracks (no GFTracks)
	// GFtracks are accessed ASSUMING SAME INDEXING as CdcTracks!
	std::set<unsigned int>::const_iterator itTpcSet;

	for (itTpcSet = tpcIdsSeen.begin(); itTpcSet != tpcIdsSeen.end();
			itTpcSet++) { // loop over TPC track ID

		unsigned int nTup = fTupleOutArray->GetEntriesFast();
		MatchingTuple *iTup = new ((*fTupleOutArray)[nTup]) MatchingTuple();
		iTup->addIdAndBranch(*itTpcSet, fTpcTrackBranchName);

		// make a copy of the existing track for merging
		GFTrack *iTpcTr = (GFTrack *)fTpcTrackArray->At(*itTpcSet);
		unsigned int nOut = fTrackOutArray->GetEntriesFast();
		GFTrack *outTrack=NULL;
		if(fRealData){
			outTrack= new ((*fTrackOutArray)[nOut]) GFTrack(*iTpcTr);
		}
		iTup->setRefID(nOut);

		std::map<unsigned int, std::vector<MatchingTuple *> >::const_iterator itTpc;
		itTpc = tupmapTpc.find(*itTpcSet);
		// if there are no successful matches with this TPC track, write out ths TPC
		// track and cont.
		if (itTpc == tupmapTpc.end())
			continue;

		int cdcTrackID = -1;
		// now find best CDC match
		double matchQual =
				0.; // criterion for now: smallest sum of constrained (abs) values

		for (unsigned int ict = 0; ict < (itTpc->second).size();
				ict++) { // loop over tuple collection
			double thisMatch = 0.;
			MatchingTuple *theMat = (itTpc->second)[ict];
			if (!theMat->hasBranch(fCdcTrackBranchName))
				continue;
			std::vector<TString> names = theMat->getListOfCritNames();
			for (unsigned int in = 0; in < names.size(); in++) {
				if (!(theMat->getCriterion(names[in]).isConstrained()))
					continue;
				thisMatch += std::fabs((theMat->getCriterion(names[in]).getValue()));
			}
			if (ict == 0 || thisMatch < matchQual) {
				matchQual = thisMatch;
				cdcTrackID = theMat->getIdFromBranch(fCdcTrackBranchName);
			}

		} // end loop over tuple collection with this TPC track ID

		// there was no matched track for this TPC track, written out anyway
		if (cdcTrackID < 0) // should never get here
			continue;

		// from here on: merge tracks
		if (fRealData) {
			std::vector<GFAbsRecoHit *> tpcHits = outTrack->getHits();
			GFAbsRecoHit *lastHit = (GFAbsRecoHit *)tpcHits.back();
			TVectorD lastPos(3);
			lastPos = lastHit->getRawHitCoord();
			TVector3 lastPosV(lastPos[0], lastPos[1], 0.);

			// now build CDC GFRecoHits
			CdcTrack *cdcTrack = (CdcTrack *)fCdcTrackArray->At(cdcTrackID);
			double mx = cdcTrack->GetMx();
			double my = cdcTrack->GetMy();
			TVector3 cent(mx, my, 0.);

			// ASSUMING IDENTICAL INDEXING HERE
			// make a copy pf the existing CDC GFTrack
			GFTrack *cdcGFTrack =
					new GFTrack(*(GFTrack *)fCdcGFTrackArray->At(cdcTrackID));

			std::vector<GFAbsRecoHit *> cdcHits = cdcGFTrack->getHits();

			// now assign sorting parameter: Phi with respect to the circle center
			TVector3 relTpc = lastPosV - cent;

			for (unsigned int icdch = 0; icdch < cdcHits.size(); icdch++) {
				TVectorD raw(3);
				TVector3 hitV;
				if (dynamic_cast<const CdcSpacepoint *>(cdcHits[icdch]) ||
						dynamic_cast<const PseudoSpacePoint *>(cdcHits[icdch])) {
					raw = cdcHits[icdch]->getRawHitCoord();
					hitV.SetXYZ(raw[0], raw[1], 0.);
				}

				else { // RPC / Barrel hit
					GFDetPlane pl =
							cdcHits[icdch]->getDetPlane(cdcGFTrack->getCardinalRep());
					hitV = pl.getO();
					hitV.SetZ(0.);
				}
				double sort = hitV.Angle(lastPosV);
				cdcHits[icdch]->setRho(sort);
			}

			try {
				cdcGFTrack->sortHits();
			} catch (GFException &ex) {
				std::cout << ex.what() << std::endl;
				delete cdcGFTrack;
				continue;
			}

			// merge tracks:
			outTrack->mergeHits(cdcGFTrack);
			delete cdcGFTrack;
		}
		// CdcTrack and CdcGFTrack IDs are assumed to be the same, using CdcTrack
		// BranchName here...
		iTup->addIdAndBranch(cdcTrackID, fCdcTrackBranchName);
		iTup->setMatched(true);

	} // end loop over TPC track ID

	// END TPC-CDC MATCHING
	// ---------------------------------------------------------------

	// TODO: create and write out final tuples for QA
}

void FopiTrackMergingTask::AddTupleBranchName(const TString &t) {
	if (fTupleBranchMap.count(t)) {
		TString err("Branch entry '");
		err += t;
		err += "' already existing!";
		Fatal("AddTupleBranchName", err);
	}

	FairRootManager *ioman = FairRootManager::Instance();
	if (ioman == 0)
		Fatal("AddTupleBranchName", "RootManager not instantiated!");

	fTupleBranchMap[t] = NULL;
}

ClassImp(FopiTrackMergingTask)