#include "PndFtsHoughTrackFinder.h"

#include <iostream>

#include "TMath.h"
#include <math.h>
#include <algorithm>

#include <set>
#include <vector>
#include <map>

// FTS
#include "PndGeoFtsPar.h"
#include "PndFtsMapCreator.h"
#include "PndFtsHit.h"
#include "FairHit.h"

// magnetic field
#include "FairField.h"
#include "TVector3.h"

// (Hough) tracking
#include "PndFtsHoughTrackerTask.h"
#include "PndTrackCand.h"
#include "PndTrack.h"
#include "FairTrackParP.h"
#include "PndFtsHoughTrackCand.h"

// histogramming / plotting
#include "TH1.h"
#include "TH2.h"
#include "TGraph.h"

// peak finder
#include "TSpectrum2.h"

// root IO
#include "TClonesArray.h"
#include "FairRunAna.h"
#include "FairRootManager.h"
#include "FairRuntimeDb.h"
#include "FairTask.h"

#include "PndFtsHit.h"
#include "TString.h"
#include "FairHit.h"


#include "PndFtsTube.h"
#include "FairEventHeader.h"




using namespace std;

ClassImp(PndFtsHoughTrackFinder)
;





PndFtsHoughTrackFinder::PndFtsHoughTrackFinder(PndFtsHoughTrackerTask *trackerTask, Int_t branchId, TClonesArray* hits, FairField* field) :
										fTrackerTask(trackerTask),
										fFtsHitArray(hits),
										fFtsBranchId(branchId),
										fField(field),

										// Hough spaces
										fHoughSpaceZxLineBeforeDipole(0),
										fHoughspaceZxParabola(0),
										fHoughSpaceZxLineBehindDipole(0),
										fHoughspaceZyLine(0),

										// min peak heights
										fMinPeakHeightZxLineParabola(4),
										fMinPeakHeightZxParabola(6),
										fMinPeakHeightZxParabolaLine(4),
										fMinPeakHeightZyLine(4),

										// general
										fSaveDebugInfo(kFALSE),
										fVerbose(0)
{
	if (0==fTrackerTask){
		std::cout << "PndFtsHoughTrackFinder FATAL ERROR Tracker task not set.\n";
	}
}


PndFtsHoughTrackFinder::~PndFtsHoughTrackFinder()
{
	if(3<fVerbose) Info(MESSAGE_ORIGIN,"Destructor of PndFtsHoughTrackFinder");
}







//void PndFtsHoughTrackFinder::CreatePndTrackCands() {
//
//}


void PndFtsHoughTrackFinder::FindTracks() {

	if (0<fVerbose) {
		cout << "PndFtsHoughTrackFinder::FindTracks()" << endl;
	}

	// reset
	fHoughTrackCands.clear();
	fHoughTrackCandsNew.clear();





	// Do we have enough hits in the FTS?
	if(fMinPeakHeightZxParabola > fFtsHitArray->GetEntriesFast()) {
		if(0<fVerbose) Info("Exec","Skip the event, since we have too few hits in FTS");
		return;
	}








	// zx plane: Straight line Hough transform before dipole
	static const Int_t stepsPerThetaDegLineBeforeDipole = 4; // greater number means finer scanning in theta
	static const Int_t thetaDegLowLineBeforeDipole = -18; // in degree
	static const Int_t thetaDegHighLineBeforeDipole = 18; // in degree
	delete fHoughSpaceZxLineBeforeDipole;
	fHoughSpaceZxLineBeforeDipole = new PndFtsHoughSpace(
			"lineBeforeDipole",
			stepsPerThetaDegLineBeforeDipole*(thetaDegHighLineBeforeDipole-thetaDegLowLineBeforeDipole),
			thetaDegLowLineBeforeDipole / 180. * meinpi, // in rad
			thetaDegHighLineBeforeDipole / 180. * meinpi, // in rad

			stepsPerThetaDegLineBeforeDipole*16, // TODO: Check values
			-80., // in cm // TODO: Check values
			80., // in cm

			fZLineParabola,
			0.,

			fFtsBranchId,
			fFtsHitArray,

			fField
	);
	fHoughSpaceZxLineBeforeDipole->GetXaxis()->SetTitle("#theta [rad]");
	fHoughSpaceZxLineBeforeDipole->GetYaxis()->SetTitle("x_{LP} [cm]");



	// Do straight line hough transform on non-skewed hits from stations 1+2
	if ( kTRUE == fHoughSpaceZxLineBeforeDipole->MakeHoughSpace() )
	{
		if (0<fVerbose) {
			std::cout << "Hough Space for zx line before dipole was created successfully!" << std::endl
					<< "We have " << fHoughSpaceZxLineBeforeDipole->GetNHits() << " hits in the line Hough space.\n";
		}
		if (fTrackerTask->GetSaveDebugInfo()){
			fTrackerTask->WriteHistogram(fHoughSpaceZxLineBeforeDipole);
		}
	}
	else
	{
		std::cout << "Hough Space for zx line before dipole could not be created! " << std::endl;
	}



	// find peaks for line hough space and store in vector
	const Int_t minHeightLineBeforeDipole = 6;
	std::vector<PndFtsHoughTracklet> trackletsLineBeforeDipole;


	// Call peak finder and plot the solution if it was found!
	if (kTRUE == fHoughSpaceZxLineBeforeDipole->FindAllPeaks(minHeightLineBeforeDipole, trackletsLineBeforeDipole))
	{
		if (0<fVerbose) {
			PrintFoundTracklets(trackletsLineBeforeDipole, fHoughSpaceZxLineBeforeDipole->GetName());
			fTrackerTask->WriteHistogram(fHoughSpaceZxLineBeforeDipole);
		}
	}
	else
	{
		std::cout << "Error: Peak finder had a problem with hough space for zx line before dipole!!!" << std::endl;
	}







	// loop over all line tracklets which were found by line HT before dipole field

	for (UInt_t iTrackletLine=0; iTrackletLine < trackletsLineBeforeDipole.size(); ++iTrackletLine)
	{
		const Double_t peakThetaRadLineBeforeDipole = trackletsLineBeforeDipole[iTrackletLine].getThetaRadVal();
		const Double_t peakInterceptLineBeforeDipole = trackletsLineBeforeDipole[iTrackletLine].getSecondVal();
		const Double_t peakThetaRadHwLineBeforeDipole = trackletsLineBeforeDipole[iTrackletLine].getThetaRadHw();

		// determine where to look for parabola
		const Int_t stepsPerThetaDegParabola = 10; // greater number means finer scanning in theta
		//!TODO: Rework this
		const Double_t thetaRadLowParabola = peakThetaRadLineBeforeDipole - peakThetaRadHwLineBeforeDipole; // in rad
		const Double_t thetaRadHighParabola = peakThetaRadLineBeforeDipole + peakThetaRadHwLineBeforeDipole; // in rad
		const Double_t thetaDegLowParabola = thetaRadLowParabola / meinpi * 180.; // in deg
		const Double_t thetaDegHighParabola = thetaRadHighParabola / meinpi * 180.; // in deg

		if (thetaRadHighParabola==thetaRadLowParabola) std::cout << "ERROR: low and high are the same for parabola!\n";
		UInt_t thetaBins = ceil(stepsPerThetaDegParabola*(thetaDegHighParabola-thetaDegLowParabola));

		std::cout << "event: " << fTrackerTask->GetEventNr() << "\nthetaDegLowParabola=" << thetaDegLowParabola << " thetaDegHighParabola=" << thetaDegHighParabola << " thetaBins=" << thetaBins << "  peakThetaDegLineBeforeDipole=" << peakThetaRadLineBeforeDipole/meinpi*180. << " peakThetaDegHwLineBeforeDipole=" << peakThetaRadHwLineBeforeDipole/meinpi*180. << '\n';

		delete fHoughspaceZxParabola;
		fHoughspaceZxParabola= new PndFtsHoughSpace(
				"parabola",

				thetaBins,
				thetaRadLowParabola, // in rad
				thetaRadHighParabola, // in rad

				stepsPerThetaDegParabola*300, // 300 is good as factor
				-0.015, // a.u.
				0.015, // a.u.

				fZLineParabola,
				peakInterceptLineBeforeDipole,

				fFtsBranchId,
				fFtsHitArray,

				fField
		);
		fHoughspaceZxParabola->GetXaxis()->SetTitle("#theta [rad]");
		fHoughspaceZxParabola->GetYaxis()->SetTitle("#frac{Q}{p_zx} [a.u.]");



		// Do parabola hough transform (shifts FTS hits by hitshiftinx) for non-skewed hits in stations 3+4+5
		if ( kTRUE == fHoughspaceZxParabola->MakeHoughSpace() )
		{
			if (0<fVerbose) {
				std::cout << "Line " << iTrackletLine << ": Hough Space for zx parabola was created successfully!" << std::endl
						<< "We have " << fHoughspaceZxParabola->GetNHits() << " hits in the line Hough space.\n";
			}
			if (fTrackerTask->GetSaveDebugInfo()){
				fTrackerTask->WriteHistogram(fHoughspaceZxParabola, iTrackletLine);
			}

		}
		else
		{
			std::cout << "Hough Space for zx parabola could not be created! " << std::endl;
		}


		// TODO: Rewrite the following conversion
		// Get momenta from arbitrary units to meaningful ones
		//		pzinvpeak = 1./pzinvpeak*0.00299792458;
		//	pzinvpeakWithBField = pzinvpeak/BMeanForParabola;


		// find peaks for zx parabola hough space
		const Int_t minHeightZxParabola = 8;
		std::vector<PndFtsHoughTracklet> zxParabolaTracklets;


		if (kTRUE == fHoughspaceZxParabola->FindAllPeaks(minHeightZxParabola, zxParabolaTracklets))
		{
			if (0<fVerbose)
			{
				PrintFoundTracklets(zxParabolaTracklets, fHoughspaceZxParabola->GetName());
			}

		}
		else
		{
			std::cout << "Error: Peak finder had a problem for parabola!!!" << std::endl;
		}

		if(1<fVerbose) {
			std::cout << "Create track candidates from line in zx before dipole "<< iTrackletLine << " and all parabolas within dipole." << std::endl;
		}
		for (UInt_t iTrackletParabola=0; iTrackletParabola < zxParabolaTracklets.size(); ++iTrackletParabola)
		{
			PndFtsHoughTrackCand newHoughTrackCand(fFtsBranchId, fFtsHitArray);
			newHoughTrackCand.SetZxFirstLine(trackletsLineBeforeDipole[iTrackletLine]);
			newHoughTrackCand.SetZxParabola(zxParabolaTracklets[iTrackletParabola]);
			fHoughTrackCands.push_back(newHoughTrackCand);
		}

	} // loop over all line tracklets which were found by line HT before dipole field




	//	std::cout << "At this point I have track candidates from line in zx before dipole and parabola within. Now add line after dipole field in zx" << std::endl;
	// TODO






	// zx plane: Straight line Hough transform behind dipole
	static const Int_t stepsPerThetaDegLineBehindDipole = 2; // greater number means finer scanning in theta
	static const Int_t thetaDegLowLineBehindDipole = -89; // in degree
	static const Int_t thetaDegHighLineBehindDipole = 89; // in degree
	delete fHoughSpaceZxLineBehindDipole;
	fHoughSpaceZxLineBehindDipole = new PndFtsHoughSpace(
			"lineBehindDipole",
			stepsPerThetaDegLineBehindDipole*(thetaDegHighLineBehindDipole-thetaDegLowLineBehindDipole),
			thetaDegLowLineBehindDipole / 180. * meinpi, // in rad
			thetaDegHighLineBehindDipole / 180. * meinpi, // in rad

			stepsPerThetaDegLineBehindDipole*20, // TODO: Check values
			-200., // in cm // TODO: Check values
			200., // in cm

			fZParabolaLine,
			0.,

			fFtsBranchId,
			fFtsHitArray,

			fField
	);
	fHoughSpaceZxLineBehindDipole->GetXaxis()->SetTitle("#theta [rad]");
	fHoughSpaceZxLineBehindDipole->GetYaxis()->SetTitle("x_{LP} [cm]");



	// Do straight line hough transform on non-skewed hits from stations 1+2
	if ( kTRUE == fHoughSpaceZxLineBehindDipole->MakeHoughSpace() )
	{
		if (0<fVerbose) {
			std::cout << "Hough Space for zx line behind dipole was created successfully!" << std::endl
					<< "We have " << fHoughSpaceZxLineBehindDipole->GetNHits() << " hits in the line Hough space.\n";
		}

	}
	else
	{
		std::cout << "Hough Space for zx line behind dipole could not be created! " << std::endl;
	}



	// find peaks for line hough space and store in vector
	const Int_t minHeightLineBehindDipole = 6;
	std::vector<PndFtsHoughTracklet> trackletsLineBehindDipole;


	// Call peak finder and plot the solution if it was found!
	if (kTRUE == fHoughSpaceZxLineBehindDipole->FindAllPeaks(minHeightLineBehindDipole, trackletsLineBehindDipole))
	{
		if (0<fVerbose) {
			PrintFoundTracklets(trackletsLineBehindDipole, fHoughSpaceZxLineBehindDipole->GetName());
		}
		if (fTrackerTask->GetSaveDebugInfo()){
			fTrackerTask->WriteHistogram(fHoughSpaceZxLineBehindDipole);
		}
	}
	else
	{
		std::cout << "Error: Peak finder had a problem with hough space for zx line behind dipole!!!" << std::endl;
	}














	//	if (0<fSaveDebugInfo){
	//		WriteHistograms();
	//	}


	if(0<fVerbose) {

		// TODO Add missing output




		//		if (kTRUE == correctpz)
		//		{
		//			std::cout << "Watch out! p_z value in plot is already corrected!!!"  << std::endl;
		//		}
	}


	delete fHoughSpaceZxLineBeforeDipole;
	delete fHoughspaceZxParabola;
	delete fHoughSpaceZxLineBehindDipole;
	delete fHoughspaceZyLine;

}









Bool_t PndFtsHoughTrackFinder::FilterTrackletsBasedOnSharedHits(
		UInt_t maxAcceptableSharedHits,
		std::vector<PndFtsHoughTracklet> &tracklets
)
{
	// to store indices that I want to "delete"
	std::set<UInt_t> indicesToDelete;
	std::set<UInt_t>::iterator findIndex;


	// compare all tracklets with each other, keep track which tracklets should be "deleted"
	// delete all tracklets that share more than maxAcceptableSharedHits hits
	// keep only the heighest peak
	for (UInt_t iTrackletLeft = 0; iTrackletLeft < tracklets.size(); ++iTrackletLeft)
	{
		for (UInt_t iTrackletRight = iTrackletLeft+1; iTrackletRight < tracklets.size(); ++iTrackletRight)
		{
			// check if two tracks share more than maxAcceptableSharedHits hits
			PndFtsHoughTracklet trackletLeft = tracklets[iTrackletLeft];
			PndFtsHoughTracklet trackletRight = tracklets[iTrackletRight];
			const UInt_t nSharedHits = trackletLeft.getNSharedHits(trackletRight);
			if (nSharedHits>maxAcceptableSharedHits){
				// mark the peak with the lower "height" for deletion
				const Double_t heightLeft = trackletLeft.getPeakHeightFromPeakFinder();
				const Double_t heightRight = trackletRight.getPeakHeightFromPeakFinder();
				if (heightLeft>heightRight){
					indicesToDelete.insert(iTrackletRight);
				} else if (heightLeft<heightRight){
					indicesToDelete.insert(iTrackletLeft);
				} else if (heightLeft==heightRight){
					std::cout << "WARNING: Found two peaks of the same height that share " <<  nSharedHits << " hits.\n";
				}
			}

		}
	} // end for loop: compare all tracklets with each other

	// if we have no tracklets to "delete", we don't need to do anything
	if (0==indicesToDelete.size()){
		if (0<fVerbose) {  std::cout << "FilterTrackletsBasedOnSharedHits: No tracklets are marked for deletion." << std::endl; }
		return kTRUE;
	}

	// create new tracklets vector and copy all entries from input vector to it which are not marked for deletion
	std::vector<PndFtsHoughTracklet> filteredTracklets;
	for (UInt_t iTracklet = 0; iTracklet < tracklets.size(); ++iTracklet)
	{
		// check if the index is marked for "deletion"
		findIndex = indicesToDelete.find(iTracklet);
		if (findIndex == indicesToDelete.end()){
			// index was not found -> entry is not marked for deletion -> I should copy it
			filteredTracklets.push_back(tracklets[iTracklet]);
		}
	}
	// replace input tracklets with filtered ones
	tracklets = filteredTracklets;
	return kTRUE;
}