#include "PndSttCellTrackFinderTask.h"

#include <iostream>
#include <stdio.h>

// Root includes
#include "TROOT.h"
#include "TString.h"
#include "TClonesArray.h"
#include "TParticlePDG.h"

// framework includes
#include "FairRootManager.h"
#include "FairRun.h"
#include "FairRuntimeDb.h"
#include "FairRunAna.h"
#include "FairEventHeader.h"

// PndMvd includes
#include "PndTrackCand.h"
#include "PndRiemannTrack.h"
#include "PndSttMapCreator.h"
#include "PndSttCellTrackFinder.h"

using std::cout;
using std::endl;

//Macro for printing the calculation times into a file called calcTimesCPU.txt
//#define PRINT_CALC_TIMES


ClassImp(PndSttCellTrackFinderTask);

void PndSttCellTrackFinderTask::SetParContainers() {
	FairRuntimeDb* rtdb = FairRunAna::Instance()->GetRuntimeDb();
	fSttParameters = (PndGeoSttPar*) rtdb->getContainer("PndGeoSttPar");
}

InitStatus PndSttCellTrackFinderTask::Init() {
	FairRootManager* ioman = FairRootManager::Instance();

	if (!ioman) {
		std::cout << "-E- PndSttCellTrackFinderTask::Init: "
				<< "RootManager not instantiated!" << std::endl;
		return kFATAL;
	}

	fEventHeader = (TClonesArray*) ioman->GetObject("EventHeader.");
	if (!fEventHeader) {
		cout
				<< "-W-  PndSttCellTrackFinderTask::Init: No EventHeader array! Needed for EventNumber"
				<< endl;
		return kERROR;
	}

	if (fHitBranch.size() == 0) {
		std::cout << "-W- PndSttCellTrackFinderTask::Init: "
				<< "No Branch Names given with AddHitBranch(TString branchName)! Standard BranchNames taken!"
				<< std::endl;
		fHitBranch.push_back("STTHit");
		fHitBranch.push_back("STTCombinedSkewedHits");
	}

	for (int i = 0; i < (int) fHitBranch.size(); i++) {
		InitHitArray(fHitBranch[i]);
	}

	PndSttMapCreator *mapper = new PndSttMapCreator(fSttParameters);
	fTubeArray = mapper->FillTubeArray();
	fTrackFinder.SetSttTubeArray(fTubeArray);
	fTrackFinder.SetCalcFirstTrackletInf(fAnalyseSteps);
	fTrackFinder.SetVerbose(fVerbose);

	fFirstTrackCandArray = ioman->Register("FirstTrackCand", "PndTrackCand",
			"STT", fPersistence);
	fFirstRiemannTrackArray = ioman->Register("FirstRiemannTrack",
			"PndRiemannTrack", "STT", fPersistence);

	fCombiTrackCandArray = ioman->Register("CombiTrackCand", "PndTrackCand",
			"STT", fPersistence);
	fCombiTrackArray = ioman->Register("CombiTrack", "PndTrack", "STT",
			fPersistence);
	fCombiRiemannTrackArray = ioman->Register("CombiRiemannTrack",
			"PndRiemannTrack", "STT", fPersistence);

	fCorrectedIsochronesArray = ioman->Register("CorrectedIsochrones",
			"FairHit", "STT", fPersistence);

	fCorrectedCombiTrackCandArray = ioman->Register("CombiTrackCandCorrected",
			"PndTrackCand", "STT", fPersistence);
	fCorrectedCombiTrackArray = ioman->Register("CombiTrackCorrected",
			"PndTrack", "STT", fPersistence);
	fCorrectedCombiRiemannTrackArray = ioman->Register(
			"CombiRiemannTrackCorrected", "PndRiemannTrack", "STT",
			fPersistence);

	std::cout << "-I- PndSttCellTrackFinderTask: Initialisation successfull"
			<< std::endl;
	//fInitDone = kTRUE;
	return kSUCCESS;
}

// -----   Public method Exec   --------------------------------------------
void PndSttCellTrackFinderTask::Exec(Option_t* opt) {

	FairEventHeader* myEventHeader = (FairEventHeader*) fEventHeader;
	int eventNumber = myEventHeader->GetMCEntryNumber();

	if (fVerbose > 0) {
	cout
			<< "====================Begin PndSttCellTrackFinderTask::Exec======================="
			<< endl;

	cout << "Event #" << eventNumber << endl;
	}

// Reset output array
	if (!fFirstTrackCandArray)
		Fatal("Exec", "No trackCandArray");

	fTrackFinder.Reset();

	FairRootManager *ioman = FairRootManager::Instance();

	for (int i = 0; i < (int) fHitBranch.size(); i++) {
		fTrackFinder.AddHits(fHitArray[i], ioman->GetBranchId(fHitBranch[i]));
	}

	fTrackFinder.FindTracks();
	std::map<int, FairHit*> correctedIsochrones =
			fTrackFinder.GetCorrectedIsochrones();

	int indexCounter = 0;
	for (std::map<int, FairHit*>::iterator iter = correctedIsochrones.begin();
			iter != correctedIsochrones.end(); iter++) {
		//std::cout << "Corrected Isochrone for TubeId: " << iter->first << " : "<< *iter->second << std::endl;
		FairHit* myCorrectedHit =
				new (
						(*fCorrectedIsochronesArray)[fCorrectedIsochronesArray->GetEntries()]) FairHit(
						*iter->second);
	}

	if (fAnalyseSteps) {
		for (int i = 0; i < fTrackFinder.NumFirstTrackCands(); i++) {
			PndTrackCand* myCand =
					new ((*fFirstTrackCandArray)[i]) PndTrackCand(
							fTrackFinder.GetFirstTrackCand(i));
		}

//		for (int i = 0; i < fTrackFinder.NumFirstRiemannTracks(); ++i) {
//			PndRiemannTrack* myTrack =
//					new ((*fFirstRiemannTrackArray)[i]) PndRiemannTrack(
//							fTrackFinder.GetFirstRiemannTrack(i));
//		}
	}

	for (int i = 0; i < fTrackFinder.NumCombinedTracks(); ++i) {
		PndTrackCand* myCand = new ((*fCombiTrackCandArray)[i]) PndTrackCand(
				fTrackFinder.GetCombiTrackCand(i));
		PndTrack* myTrack = new ((*fCombiTrackArray)[i]) PndTrack(
				fTrackFinder.GetCombiTrack(i));

		PndTrackCand* myCandCorrected = new (
				(*fCorrectedCombiTrackCandArray)[i]) PndTrackCand(
				fTrackFinder.GetCorrectedCombiTrackCand(i));
		PndTrack* myTrackCorrected =
				new ((*fCorrectedCombiTrackArray)[i]) PndTrack(
						fTrackFinder.GetCorrectedCombiTrack(i));

		myTrack->SetTrackCandRef(myCand);
		myTrack->SetTrackCand(*myCand);

		myTrackCorrected->SetTrackCandRef(myCandCorrected);
		myTrackCorrected->SetTrackCand(*myCandCorrected);
	}

//
////	for (int i = 0; i < fTrackFinder.NumCombinedRiemannTracks(); ++i) {
////
////		PndRiemannTrack* myRiemannTrack =
////				new ((*fCombiRiemannTrackArray)[i]) PndRiemannTrack(
////						fTrackFinder.GetCombiRiemannTrack(i));
////	}
//
//	if (fVerbose > 0)
//		cout << "#FirstTracklets: " << fTrackFinder.NumFirstTrackCands()
//				<< ", #CombinedTracklets: " << fTrackFinder.NumCombinedTracks()
//				<< endl;
//
	fCombiTrackCandArray->Sort();
	fCombiTrackArray->Sort();
	fCorrectedCombiTrackCandArray->Sort();
	fCorrectedCombiTrackArray->Sort();

}

void PndSttCellTrackFinderTask::FinishEvent() {
	fFirstTrackCandArray->Delete();
	fFirstRiemannTrackArray->Delete();
	fCombiTrackCandArray->Delete();
	fCombiTrackArray->Delete();
	fCombiRiemannTrackArray->Delete();
	fCorrectedIsochronesArray->Delete();
	fCorrectedCombiTrackCandArray->Delete();
	fCorrectedCombiTrackArray->Delete();
	fCorrectedCombiRiemannTrackArray->Delete();

	fNumUnambiguousHitsPerEvent.push_back(
			fTrackFinder.NumUnambiguousNeighbors());

}

void PndSttCellTrackFinderTask::FinishTask() {

#ifdef PRINT_CALC_TIMES
	// write calculation times
	vector<vector<Double_t> > calcTimes =
			fTrackFinder.GetTimeStampsTrackletGen();
	FILE* fp = fopen("calcTimesCPU.txt", "w");
	fprintf(fp,
			"%3s \t %3s \t %10s \t %10s \t %10s \t %10s \t %10s \t %10s \t %10s \t %10s \t %10s \t %10s \n",
			"#E", "#UH", "EvaluateState()", "InitStartTracklets()",
			"EvaluateMultiState()", "GenerateTracklets()",
			"CombineTrackletsMultiStages()", "AssignAmbiguousHits()",
			"SplitData()", "AddMissingHits()", "CreatePndTrackCands()",
			"FindTracks()");

	for (int i = 0; i < calcTimes.size(); ++i) {
		fprintf(fp, "%3i \t %3i", i,fNumUnambiguousHitsPerEvent.at(i));
		for (int j = 0; j < calcTimes[i].size(); j += 2) {
			fprintf(fp, "\t %.15f", (calcTimes[i].at(j + 1) - calcTimes[i].at(j)) * 1000);
		}
		fprintf(fp, "\n");
	}

	fclose(fp);
#endif

}

void PndSttCellTrackFinderTask::InitHitArray(TString branchName) {
	TClonesArray* tempArray =
			(TClonesArray*) FairRootManager::Instance()->GetObject(branchName);
	if (tempArray == 0) {
		std::cout << "-W- PndSttCellTrackFinderTask::Init: "
				<< "No hitArray for BranchName " << branchName.Data()
				<< std::endl;
	}
	fHitArray.push_back(tempArray);
}