#include "PndMvdTpcRiemannCorrelatorTask.h"

#include <iostream>

// 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 "PndMCTrack.h"

#include "PndDetectorList.h"

// PndMvd includes
#include "PndSdsRecoHit.h"
// #include "PndMvdGFTrackCand.h"
#include "GFTrack.h"
#include "GFTrackCand.h"
#include "PndSdsHit.h"
#include "PndSdsMCPoint.h"
#include "PndSdsCluster.h"
#include "PndSdsDigi.h"
#include "PndMvdTpcRiemannCorrelatorTask.h"
//#include "PndTpcHit.h"
#include "PndTpcCluster.h"
#include "PndRiemannHit.h"
//#include "PndRiemannTrack.h"

PndMvdTpcRiemannCorrelatorTask::PndMvdTpcRiemannCorrelatorTask() : FairTask("MVD TPC Riemann Track Correlator"),
fMaxDist(1), fMaxSZ(1), fMaxSZChi2(1)
{
	fHitBranchMVDPixel = "MVDHitsPixel";
	fHitBranchMVDStrip = "MVDHitsStrip";
	fHitBranchTPC = "PndTpcClusterMerged";
	fTrackBranch = "MVDRiemannGFTrackCand";

	fEventNr = 0;
	//fTrackBranch = "MCTrack";
}

PndMvdTpcRiemannCorrelatorTask::~PndMvdTpcRiemannCorrelatorTask()
{
}

void PndMvdTpcRiemannCorrelatorTask::SetParContainers()
{
  // Get Base Container
/*
  FairRun* ana = FairRun::Instance();
  FairRuntimeDb* rtdb=ana->GetRuntimeDb();
  fGeoPar = (PndMvdGeoPar*)(rtdb->getContainer("PndMvdGeoPar"));
*/
}

InitStatus PndMvdTpcRiemannCorrelatorTask::ReInit()
{

  InitStatus stat=kERROR;
  return stat;

  /*
  FairRun* ana = FairRun::Instance();
  FairRuntimeDb* rtdb=ana->GetRuntimeDb();
  fGeoPar=(PndMvdGeoPar*)(rtdb->getContainer("PndMvdGeoPar"));

  return kSUCCESS;
  */
}

// -----   Public method Init   --------------------------------------------
InitStatus PndMvdTpcRiemannCorrelatorTask::Init()
{

  FairRootManager* ioman = FairRootManager::Instance();

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

  // Get input array
  fHitArrayMVDPixel = (TClonesArray*) ioman->GetObject(fHitBranchMVDPixel);
  if ( !fHitArrayMVDPixel){
    std::cout << "-W- PndMvdIdealTrackFinderTask::Init: " << "No hitArray Pixel!" << std::endl;
    return kERROR;
  }

  fHitArrayMVDStrip = (TClonesArray*) ioman->GetObject(fHitBranchMVDStrip);
   if ( !fHitArrayMVDStrip){
     std::cout << "-W- PndMvdIdealTrackFinderTask::Init: " << "No hitArray Strip!" << std::endl;
     return kERROR;
   }

  fHitArrayTPC = (TClonesArray*) ioman->GetObject(fHitBranchTPC);
    if ( !fHitArrayTPC){
      std::cout << "-W- PndMvdIdealTrackFinderTask::Init: " << "No hitArray TPC!" << std::endl;
      return kERROR;
    }

   fTrackCandArray = (TClonesArray*) ioman->GetObject(fTrackBranch);
     if ( !fTrackCandArray){
       std::cout << "-W- PndMvdIdealTrackFinderTask::Init: " << "No hitArray2!" << std::endl;
       return kERROR;
     }

  fCombinedArray = new TClonesArray("GFTrackCand");
  ioman->Register("MVD_TPC_CombinedRiemannTrack", "MVD", fCombinedArray, kTRUE);

  std::cout << "-I- PndMvdTpcRiemannCorrelator: Initialisation successfull" << std::endl;
  return kSUCCESS;
}

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

  // Reset output array
  if ( ! fCombinedArray )
    Fatal("Exec", "No CombinedGFTrackCandArray");
  fCombinedArray->Delete();

  std::vector<TH1*> distHistos;
  std::vector<TH1*> szDistHistos;
  std::vector<TH1*> szChi2Histos;

  TString distHName("distHisto_");
  TString szDistHName("szdistHisto_");
  TString szChi2HName("szChi2Histo_");

  if (fVerbose > 0)
 	std::cout << "Analysing Event Nr: " << fEventNr << std::endl;

  int foundTracks = 0;
  if (fTrackCandArray->GetEntriesFast() > 100)
	  return;
  for (int trackInd = 0; trackInd < fTrackCandArray->GetEntriesFast(); trackInd++){				//run through all tracks;
	  GFTrackCand* myMvdTrackCand = (GFTrackCand*)fTrackCandArray->At(trackInd);
	 // PndRiemannTrack myTrack = GetRiemannTrack(myMvdGFTrackCand);
	 // PndRiemannTrack newTrack = GetRiemannTrack(myMvdGFTrackCand);
	  TString name = distHName;
/*	  name+= fEventNr;
	  name+= "_";
	  name+= trackInd;
	  TH1D* myDistHisto = new TH1D(name.Data(), "Riemann Distance to Track", 200, -10,10);
	  name = szDistHName;
	  name+= fEventNr;
	  name+= "_";
	  name+= trackInd;
	  TH1D* mySZDistHisto = new TH1D(name.Data(), "Riemann SZDistance", 200, -10,10);
	  name = szChi2HName;
	  name+= fEventNr;
	  name+= "_";
	  name+= trackInd;
	  TH1D* mySZChi2Histo = new TH1D(name.Data(), "Riemann SZ Chi2", 200, 0,10);
	*/
	  if (fVerbose > 0)
		  std::cout << "TrackNr: " << trackInd << std::endl;
	  if (fVerbose > 0)
		  std::cout << "First iteration - Candidates: ";
	  GFTrackCand newCand = AddTPCHits(myMvdTrackCand); //first iteration
	  if (fVerbose > 0){
		  std::cout << newCand.getNHits() << std::endl;
		  std::cout << "Second iteration - Candidates: ";
	  }

//	  GFTrackCand finalCand = newCand;
	  GFTrackCand finalCand = AddTPCHits(&newCand);
	  if (fVerbose > 0)
		  std::cout << finalCand.getNHits() << std::endl;
	/*  for (int tpcCluster = 0; tpcCluster < fHitArrayTPC->GetEntriesFast(); tpcCluster++){
		  PndTpcCluster* myCluster = (PndTpcCluster*) fHitArrayTPC->At(tpcCluster);
		  PndTpcHit myTpcHit(*myCluster);
		  PndRiemannHit myRHit((FairHit*)&myTpcHit);
		  double trackDist = myTrack.dist(&myRHit);
		  double szDist = myTrack.szDist(&myRHit);
		  double szChi2 = myTrack.calcSZChi2(&myRHit);

		  myDistHisto->Fill(trackDist);
		  mySZDistHisto->Fill(szDist);
		  mySZChi2Histo->Fill(szChi2);

		  if (fabs(szDist) < fMaxSZ && fabs(trackDist) < fMaxDist && fabs(szChi2) < fMaxSZChi2){
			  if (fVerbose > 1){
				  std::cout << "TPCCluster " << tpcCluster << " added to track " << trackInd << std::endl;
				  std::cout << "trackDist: " << trackDist << " szDist: " << szDist << " szChi2: " << szChi2 << std::endl;
			  }
			  myMvdGFTrackCand->addHit(3, tpcCluster);
			  newTrack.addHit(&myRHit);
		  }

	  }*/
//	  distHistos.push_back(myDistHisto);
//	  szDistHistos.push_back(mySZDistHisto);/
//	  szChi2Histos.push_back(mySZChi2Histo);
//	  newTrack.refit();
//	  newTrack.szFit();
  	PndRiemannTrack finalTrack = GetRiemannTrack(&finalCand);

	  finalCand.setCurv(1/finalTrack.r());
	  finalCand.setDip(finalTrack.dip());

	  if (finalCand.getNHits() > 10){
		  new ((*fCombinedArray)[foundTracks])GFTrackCand(finalCand);
		  foundTracks++;
	  }
  }
//  fHRiemannDistVector.push_back(distHistos);
//  fHSZDistVector.push_back(szDistHistos);
//  fHSZChi2Vector.push_back(szChi2Histos);
  fEventNr++;


//  for (int i = 0; i < trackFinder.NTracks(); i++){
//	  new ((*fTrackCandArray)[i])GFTrackCand(trackFinder.GetTrackCand(i));
//  }
}

PndRiemannTrack PndMvdTpcRiemannCorrelatorTask::GetRiemannTrack(GFTrackCand* cand)
{
	PndRiemannTrack result;
	for (unsigned int i = 0; i < cand->getNHits(); i++){
		unsigned int detId, hitId;
		cand->getHit(i, detId, hitId);
		FairHit* myHit = 0;
		if (detId == 1){
			myHit = (FairHit*)fHitArrayMVDPixel->At(hitId);
		}
		else if (detId == 0){
			myHit = (FairHit*)fHitArrayMVDStrip->At(hitId);
		}
		if (myHit != 0){
			PndRiemannHit hit(myHit);
			result.addHit(hit);
		}
	}
	result.refit(true);
	result.szFit(true);
	return result;
}

GFTrackCand PndMvdTpcRiemannCorrelatorTask::AddTPCHits(GFTrackCand* myCand)
{
	PndRiemannTrack myTrack = GetRiemannTrack(myCand);
	GFTrackCand newCand = *myCand;
	TVector3 pos, sig;
	for (int tpcCluster = 0; tpcCluster < fHitArrayTPC->GetEntriesFast(); tpcCluster++){
		//??? if (myCand->HitInTrack(3, tpcCluster)) continue;
		PndTpcCluster* myCluster = (PndTpcCluster*) fHitArrayTPC->At(tpcCluster);
//		PndTpcHit myTpcHit(*myCluster);
		//Since the TPC code has no FairHit data type We create a simple FairHit just with
		//the 3d position from the TpcCluster
		pos = myCluster->pos();
		sig = myCluster->sig();
		FairHit *myTpcHit = new FairHit(kTpcCluster, pos, sig,0);
		//caution: the index here is not pointing to an actual CbmPoint
		//TODO: change the detector number (3) to a kTpcSomething
		PndRiemannHit myRHit(myTpcHit);
		double trackDist = myTrack.dist(&myRHit);
		double szDist = myTrack.szDist(&myRHit);
		double szChi2 = myTrack.calcSZChi2(&myRHit);

//		  myDistHisto->Fill(trackDist);
//		  mySZDistHisto->Fill(szDist);
//		  mySZChi2Histo->Fill(szChi2);

		if (fabs(szDist) < fMaxSZ && fabs(trackDist) < fMaxDist && fabs(szChi2) < fMaxSZChi2){
			if (fVerbose > 1){
				std::cout << "TPCCluster " << tpcCluster << " added" <<  std::endl;
				std::cout << "trackDist: " << trackDist << " szDist: " << szDist << " szChi2: " << szChi2 << std::endl;
			}
			newCand.addHit(3, tpcCluster);
			//newTrack.addHit(&myRHit);
		}
	}
	return newCand;
}

void PndMvdTpcRiemannCorrelatorTask::DrawDistHistos(TCanvas* can, int event)
{
	DrawHistos(can, fHRiemannDistVector[event]);
}

void PndMvdTpcRiemannCorrelatorTask::DrawSZHistos(TCanvas* can, int event)
{
	DrawHistos(can, fHSZDistVector[event]);
}

void PndMvdTpcRiemannCorrelatorTask::DrawSZChi2Histos(TCanvas* can, int event)
{
	DrawHistos(can, fHSZChi2Vector[event]);
}

void PndMvdTpcRiemannCorrelatorTask::DrawHistos(TCanvas* can, std::vector<TH1*> histos)
{
	can->Divide(3,3);

	for (unsigned int i = 0; i < histos.size() && i < 9; i++){
		can->cd(i+1);
		can->Update();
		histos[i]->DrawClone();
	}
}


ClassImp(PndMvdTpcRiemannCorrelatorTask);