/*
 * PndIsochroneTrackFinder.cxx
 *
 *  Created on: 10.10.2014
 *      Author: Stockmanns
 */

#include "PndIsochroneTrackFinder.h"

#include "PndSttHit.h"

#include "TClonesArray.h"
#include "TH2D.h"
#include "TVector2.h"
#include "TMath.h"
#include "FairRootManager.h"

ClassImp(PndIsochroneTrackFinder);

PndIsochroneTrackFinder::PndIsochroneTrackFinder(): fHoughHisto(0), fHitsToFit(0), fClockwise(kFALSE) {
	// TODO Auto-generated constructor stub
	fHoughHisto = new TH2D("histo","histo", 1000, -200, 200, 1000, -200, 200);
	fWeightMap["STTHit"] = 1;
	fWeightMap["MVDHitPixel"] = 2;
	fWeightMap["MVDHitStrip"] = 2;
	fWeightMap["GemHit"] = 2;

}

PndIsochroneTrackFinder::~PndIsochroneTrackFinder() {
	// TODO Auto-generated destructor stub
}

void PndIsochroneTrackFinder::AddHits(TClonesArray* hits, Int_t branchId)
{
	if (branchId == FairRootManager::Instance()->GetBranchId("STTHit")){
		std::cout << "PndIsochroneTrackFinder::AddHits size() " << hits->GetEntriesFast() << std::endl;
		for (int i = 0; i < hits->GetEntriesFast(); i++){
			PndSttHit* mySttHit = (PndSttHit*)hits->At(i);
			mySttHit->SetEntryNr(FairLink(-1, FairRootManager::Instance()->GetEntryNr(), branchId, i));
			std::cout << "PndIsochroneTrackFinder::AddHits EntryNr SttHit: " << mySttHit->GetEntryNr() << std::endl;
			PndGeneralHoughHit myHoughHit(*mySttHit);
			std::cout << "PndIsochroneTrackFinder::AddHits EntryNr HoughHit: " << myHoughHit.GetEntryNr() << std::endl;

			myHoughHit.SetIsochroneRadius(mySttHit->GetIsochrone());
			fHitsToFit.push_back(myHoughHit);
		}
	} else {
		for (int i = 0; i < hits->GetEntriesFast(); i++){
			FairHit* myFairHit = (FairHit*)hits->At(i);
			myFairHit->SetEntryNr(FairLink(-1, FairRootManager::Instance()->GetEntryNr(), branchId, i));
			std::cout << "PndIsochroneTrackFinder::AddHits EntryNr FairHit: " << myFairHit->GetEntryNr() << std::endl;
			PndGeneralHoughHit myHoughHit(*myFairHit);
			std::cout << "PndIsochroneTrackFinder::AddHits EntryNr HoughHit: " << myHoughHit.GetEntryNr() << std::endl;
			myHoughHit.SetIsochroneRadius(0);
			fHitsToFit.push_back(myHoughHit);
		}
	}
}

std::vector<PndTrack> PndIsochroneTrackFinder::FindTracks()
{
	std::vector<PndTrack> result;

	for (int i = 0; i < fHitsToFit.size(); i++){
		Double_t xVal(.0);
		Double_t yVal(.0);
		Double_t isoChrone(.0);
		PndGeneralHoughHit myHit = fHitsToFit.at(i);
		std::cout << "PndIsochroneTrackFinder::FindTracks() myHit: " << myHit.GetEntryNr() << std::endl;
		if (myHit.GetEntryNr().GetType() == FairRootManager::Instance()->GetBranchId("STTHit")){
			if (myHit.GetZ() > 35.5 || myHit.GetZ() < 34.5) //skip skewed Straw!
				continue;
		}

		xVal = myHit.GetX();
		yVal = myHit.GetY();
		isoChrone = myHit.GetIsochroneRadius();

		std::cout << "Processing Hit " << i << std::endl;


		TVector2 oldCC(0,0);
		for (Int_t phi = 0; phi < 360; phi += 1){
			Double_t phiRad = phi * TMath::DegToRad();		//Todo implement check for skewed layers
			TVector2 circleCenter = IsoCalc(xVal, yVal, isoChrone , phiRad);
			if (oldCC.Mod() > 0){
				TVector2 temp = oldCC;
				temp -= circleCenter;
				if (TMath::Abs(oldCC.Mod()) < 300 && TMath::Abs(circleCenter.Mod()) < 300){
					std::cout << "PndIsochroneTrackFinder::FindTracks weight for " << myHit.GetEntryNr() << " weight: " << fWeightMap[FairRootManager::Instance()->GetBranchName(myHit.GetEntryNr().GetType())] << std::endl;
					InterpolateHisto(fHoughHisto, oldCC, circleCenter, i, fWeightMap[FairRootManager::Instance()->GetBranchName(myHit.GetEntryNr().GetType())]);
				}
				oldCC = circleCenter;
			} else {
				Double_t x, y;
				x = circleCenter.X();
				y = circleCenter.Y();
				Int_t bin = fHoughHisto->Fill(x, y, fWeightMap[FairRootManager::Instance()->GetBranchName(myHit.GetEntryNr().GetType())]);
				fBinsWithContent.insert(bin);
				fBinToHit.insert(std::pair<Int_t, Int_t>(bin, i));
				oldCC = circleCenter;
			}
		}
	}
	std::multimap<Int_t, Int_t> peakBins = FindPeaks(fBinsWithContent, fHoughHisto, 6);
	std::cout << "PeakFinder: " << std::endl;
	for (std::multimap<Int_t, Int_t>::reverse_iterator peakIter = peakBins.rbegin(); peakIter != peakBins.rend(); ++peakIter){
		Int_t binX, binY, binZ;
		fHoughHisto->GetBinXYZ(peakIter->second, binX, binY, binZ);
		TVector2 circle(fHoughHisto->GetXaxis()->GetBinCenter(binX), fHoughHisto->GetYaxis()->GetBinCenter(binY));
		std::cout << peakIter->first << " : " << circle.X() << "/" << circle.Y() << std::endl;
		PndTrackCand myTrackCand = FillTrackCand(peakIter->second);

		if (myTrackCand.GetNHits() > 4){
			PndTrack myTrack = CalcPndTrack(myTrackCand, circle);
			fTrackCands.push_back(myTrackCand);
			result.push_back(myTrack);
			std::cout << "PndTrack: " << myTrack << std::endl;

		}
	}
	return result;
}

PndTrack PndIsochroneTrackFinder::CalcPndTrack(PndTrackCand cand, TVector2 circle, Double_t mag) {
	FairLink first = cand.GetSortedHit(0);
	FairLink last = cand.GetSortedHit(cand.GetNHits()-1);
	FairHit firstHit = (FairHit)fHitsToFit.at(first.GetIndex());
	FairHit lastHit =  (FairHit)fHitsToFit.at(last.GetIndex());

	FairTrackParP firstPar = CalcTrackParP(firstHit, circle, mag);
	FairTrackParP lastPar  = CalcTrackParP(lastHit, circle, mag);
	PndTrack result(firstPar, lastPar, cand);
	return result;
}

FairTrackParP PndIsochroneTrackFinder::CalcTrackParP(FairHit hit, TVector2 circleCenter, Double_t mag)
{
	TVector3 hitPos;
	TVector3 hitPosError(0.1, 0.1, 0.1);
	TVector3 momError(2, 2, 2);
	TVector3 dj(1,0,0);
	TVector3 dk(0,1,0);
	TVector3 origin(0, 0, 1);

	hitPos = CalcHitPosInTrack(hit, circleCenter);
	Double_t pt = CalcPt(mag, circleCenter);
	TVector2 hit2D(hit.GetX(), hit.GetY());
	TVector2 ptDir = CalcPtDir(hit2D, circleCenter);
	TVector2 ptVec2D = pt * ptDir;
	TVector3 ptVec(ptVec2D.X(), ptVec2D.Y(), 0);
	Double_t q = 0;
	if (fClockwise)
		q = 1;
	else
		q = -1;
	FairTrackParP result(hitPos, ptVec, hitPosError, momError, q, origin, dj, dk);
	return result;
}

Double_t PndIsochroneTrackFinder::CalcPt(Double_t mag, TVector2 circleCenter) {
	double result = 0.3 * mag * circleCenter.Mod() * 0.01;
	std::cout << "PndIsochroneTrackFinder::CalcPt " << result << std::endl;
	return result;
}

TVector2 PndIsochroneTrackFinder::CalcPtDir(TVector2 hit, TVector2 circleCenter) {
	TVector2 hitDirUnit = (hit - circleCenter).Unit();
	TVector2 result;
	result = hitDirUnit.Rotate(90.0 * TMath::DegToRad());
	if (fClockwise)
		result = result.Rotate(180.0 * TMath::DegToRad());

	return result;
}

TVector3 PndIsochroneTrackFinder::CalcHitPosInTrack(FairHit hit, TVector2 circleCenter)
{
	Double_t s = CalcArcLength(hit, circleCenter);
	TVector2 negOrigin = (-1)*circleCenter;

	Double_t startPhi = negOrigin.Phi();
	TVector2 radVec(circleCenter.Mod(), 0);
	TVector2 rotVec = radVec.Rotate(startPhi + s);
	rotVec += circleCenter;
	TVector3 result(rotVec.X(), rotVec.Y(), hit.GetZ());

	std::cout << "PndIsochroneTrackFinder::CalcHitPosInTrack s: " << s * TMath::RadToDeg() << " startPhi: " << startPhi * TMath::RadToDeg()
		<< " RadVec: " << radVec.X() << "/" << radVec.Y() << " RotVec: " << rotVec.X() << "/" << rotVec.Y() << std::endl;

	std::cout << "PndIsochroneTrackFinder::CalcHitPosInTrack Hit " << hit << " Circle: " << circleCenter.X() << "/" << circleCenter.Y()
			<< " s " << s << " startPhi: " << startPhi << " Result " << result.X() << "/" << result.Y() << "/" << result.Z() << std::endl;
	return result;
}

Double_t PndIsochroneTrackFinder::CalcArcLength(FairHit hit, TVector2 circleCenter){
	TVector2 hitVec(hit.GetX(), hit.GetY());
	TVector2 opCircleCenter = (-1)*circleCenter;
	TVector2 dirVec = opCircleCenter;

	dirVec += hitVec;
	std::cout << "PndIsochroneTrackFinder::CalcArcLength: Hit: " << hitVec.X() << "/" << hitVec.Y()
			<< " circleCenter " << circleCenter.X() << "/" << circleCenter.Y()
			<< " arcLength: " << TVector2::Phi_0_2pi(dirVec.DeltaPhi(opCircleCenter)) * TMath::RadToDeg() << std::endl;
	return TVector2::Phi_0_2pi(dirVec.DeltaPhi(opCircleCenter));

}

PndTrackCand PndIsochroneTrackFinder::FillTrackCand(Int_t histoBin) {
	PndTrackCand result;
	Int_t binX, binY, binZ;
	fHoughHisto->GetBinXYZ(histoBin, binX, binY, binZ);
	TVector2 circleCenter(fHoughHisto->GetXaxis()->GetBinCenter(binX), fHoughHisto->GetYaxis()->GetBinCenter(binY));

	std::pair<std::multimap<Int_t, Int_t>::iterator, std::multimap<Int_t, Int_t>::iterator > hitIter;
	hitIter = fBinToHit.equal_range(histoBin);
	std::map<Double_t, FairLink> trackCandHits;
	for (std::multimap<Int_t, Int_t>::iterator iter = hitIter.first; iter!= hitIter.second; iter++){

		if (fHitsUsed.find((*iter).second) != fHitsUsed.end()) continue;
		else fHitsUsed.insert((*iter).second);

		std::cout << "PndIsochroneTrackFinder::FillTrackCand " << (*iter).second << " Entries: " << fHitsToFit.size() << std::endl;

		FairHit myHit = (FairHit)fHitsToFit.at((*iter).second);
		Double_t rho = CalcArcLength(myHit, circleCenter);
		std::cout << "PndIsochroneTrackFinder::FillTrackCand: SttHit " << myHit << " EntryNr: " << myHit.GetEntryNr() << " Rho: " << rho << std::endl;

		trackCandHits[rho] = myHit.GetEntryNr();
	}

	fClockwise = kFALSE;  //check the orientation of your circle. Standard is counterclockwise.
	if (trackCandHits.rbegin()->first > TMath::Pi()){
		if ((2 * TMath::Pi()) - trackCandHits.rbegin()->first < trackCandHits.begin()->first){
			fClockwise = kTRUE;
		}
	}
	for (std::map<Double_t, FairLink>::iterator iter = trackCandHits.begin(); iter != trackCandHits.end(); iter++){
			Double_t correctedRho;
			if (fClockwise == kTRUE){
				correctedRho = 2 * TMath::Pi() - iter->first;
			} else {
				correctedRho = iter->first;
			}
			result.AddHit(iter->second, correctedRho);
	}
	return result;
}

/**
 * Calculates a point lying on a isochrone around a point, at given angle
 * @param x0 x coordinate of center point
 * @param y0 y coordinate of center point
 * @param isoR isochrone radius around center point
 * @param phi position around center point where isochrone should be evaluated, in rad
 * @return TVector2 coordinate of calculated point
 */
TVector2 PndIsochroneTrackFinder::IsoCalc(Double_t x0, Double_t y0, Double_t isoR, Double_t phi) {  // Calculate a point lying on a isochrone of radius isoR around (x0,y0) at angle phi
	Double_t s = (isoR*isoR - x0 * x0 - y0 * y0) /(2 * (x0 * TMath::Cos(phi) + y0 * TMath::Sin(phi) + isoR));
	TVector2 result(x0 + s * TMath::Cos(phi), y0 + s * TMath::Sin(phi));
	return result;
}

/**
 * Gets the global bin index in a two-dimensional histogram of a given (x,y) value pair
 * @param histo Histogram (filled)
 * @param xVal Some point's x value
 * @param yVal Some point's y value
 * @return Int_t Global index of bin where (x,y) value is found
 */
Int_t PndIsochroneTrackFinder::GetHistoBin(TH2* histo, Double_t xVal, Double_t yVal) {
	Int_t binx = histo->GetXaxis()->FindBin(xVal);
	Int_t biny = histo->GetYaxis()->FindBin(yVal);
	return histo->GetBin(binx, biny);
}

void PndIsochroneTrackFinder::InterpolateHisto(TH2* histo, TVector2& in, TVector2& out, Int_t& hitId, Double_t weight) {//, std::set<Int_t>& binsWithContent){
//	std::cout << "InterpolateBetween: " << in.X() << "/" << in.Y() << " and " << out.X() << "/" << out.Y() << std::endl;
	Double_t binWidthX = histo->GetXaxis()->GetBinWidth(1);
	Double_t binWidthY = histo->GetYaxis()->GetBinWidth(1);
	Double_t binWidth = TMath::Min(binWidthX, binWidthY);

	TVector2 dir = (out - in);
	Double_t dist = dir.Mod();
	dir = dir.Unit();

	Double_t s = 0;
	Double_t stepLength = binWidth/20;
	TVector2 runVec = in;
	TVector2 addVec = stepLength * dir;
	Int_t currentHistoIndex = GetHistoBin(histo, in.X(), in.Y());
//	std::cout << "currentHistoIndex: " << currentHistoIndex << " for " << in.X() << "/" << in.Y() << std::endl;
	s += stepLength;
	while (s < dist){
		runVec += addVec;
//		std::cout << "currentHistoIndex: " << currentHistoIndex << " FutureHistoIndex " << GetHistoBin(histo, runVec.X(), runVec.Y()) << " for " << runVec.X() << "/" << runVec.Y() << std::endl;
		if (currentHistoIndex != GetHistoBin(histo, runVec.X(), runVec.Y())){
			currentHistoIndex = GetHistoBin(histo, runVec.X(), runVec.Y());
			Int_t bin = histo->Fill(runVec.X(), runVec.Y(), 1);//, weight);
			fBinsWithContent.insert(bin);
			fBinToHit.insert(std::pair<Int_t, Int_t>(bin, hitId));
//			std::cout << "Written into Bin: " << bin << std::endl;
		}
		s += stepLength;
	}
}

/**
 * Simple threshold peakfinder
 */
std::multimap<Int_t, Int_t> PndIsochroneTrackFinder::FindPeaks(std::set<Int_t>& binsWithContent, TH2* histo, Int_t threshold) {
	std::multimap<Int_t, Int_t> result;
	for (std::set<Int_t>::iterator binIter = binsWithContent.begin(); binIter != binsWithContent.end(); binIter++){
		if (histo->GetBinContent(*binIter) > threshold){
			result.insert(std::pair<Int_t, Int_t>(histo->GetBinContent(*binIter), *binIter));
		}
	}
	return result;
}