//-----------------------------------------------------------
// File and Version Information:
// $Id$
//
// Description:
//      Implementation of class PndTpcRiemannTrackFinder
//      see PndTpcRiemannTrackFinder.h for details
//
// Environment:
//      Software developed for the PANDA Detector at FAIR.
//
// Author List:
//      Sebastian Neubert    TUM            (original author)
//      Johannes Rauch       TUM
//
//
//-----------------------------------------------------------

// Panda Headers ----------------------

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

// C/C++ Headers ----------------------
#include <algorithm>
#include <iostream>

// Collaborating Class Headers --------
#include "DebugLogger.h"
//#include "PndTpcClusterRadius.h"
#include "PndTpcClusterZ.h"
#include "GFTrackCand.h"
#include "PndTpcCluster.h"
#include "PndTpcRiemannTrack.h"
#include "PndTpcRiemannHit.h"
#include "PndTpcProximityHTCorrelator.h"
#include "PndTpcRiemannHTCorrelator.h"
#include "PndTpcRiemannTTCorrelator.h"

// initialisation of static member variables
double PndTpcRiemannTrackFinder::_interactionZ = 0.;
int PndTpcRiemannTrackFinder::_sorting = 3;

// Class Member definitions -----------
PndTpcRiemannTrackFinder::PndTpcRiemannTrackFinder()
  : _minHitsForFit(5), _sortingMode(false)
{
  // correlators in decreasing priority!
  
  //addCorrelator(new PndTpcProximityHTCorrelator(1.));
  //addCorrelator(new PndTpcRiemannHTCorrelator(5.E-4));
}


PndTpcRiemannTrackFinder::~PndTpcRiemannTrackFinder()
{
  for(int i=0;i<_correlators.size();++i){
    if(_correlators[i]!=NULL){
      delete _correlators[i];
      _correlators[i]=NULL;
    }
  }
  _correlators.clear();

  for(int i=0;i<_TTcorrelators.size();++i){
    if(_TTcorrelators[i]!=NULL){
      delete _TTcorrelators[i];
      _TTcorrelators[i]=NULL;
    }
  }
  _TTcorrelators.clear();
}

void
PndTpcRiemannTrackFinder::addCorrelator(PndTpcAbsHitTrackCorrelator* c)
{
  _correlators.push_back(c);
  _found.push_back(false);
  _bestMatchQuality.push_back(9999);
  _bestMatchIndex.push_back(0);
}

void
PndTpcRiemannTrackFinder::addTTCorrelator(PndTpcAbsTrackTrackCorrelator* c)
{
  _TTcorrelators.push_back(c);
}


unsigned int
PndTpcRiemannTrackFinder::buildTracks(std::vector<PndTpcCluster*>& cll,
				   std::vector<PndTpcRiemannTrack*>& candlist)
{
  sortClusters(cll);
  unsigned int ncl=cll.size();

  for(unsigned int icl=0;icl<ncl;++icl){ // loop over clusters
    /*if(icl<30){
      std::cout<<"Perp: "<<cll[icl]->pos().Perp()
	       <<"   Mag: "<<cll[icl]->pos().Mag()	
	       <<"   Z: "<<cll[icl]->pos().Z()<<std::endl;
    }*/

    PndTpcRiemannHit* rhit=new PndTpcRiemannHit(cll[icl]);
    unsigned int ntrks=candlist.size();
    unsigned int maxlevel=0; // index of deepest correlator reached
    bool foundAtAll=false;
    //std::cout<<"hit"<<icl<<" testing "<<ntrks<<" tracks"<<std::endl;
    for(unsigned int itrk=0;itrk<ntrks;++itrk){ // loop over tracks
      PndTpcRiemannTrack* trk=candlist[itrk];
      if(trk==NULL)continue;
      // WE STEP THROUGH THE INDIVIDUAL CORRELATORS
      // IF A TRACK SURVIVES EACH CORRELATOR
      // THE HIT IS ASSIGNED TO THE BEST (smallest!) MATCH 
      bool trksurvive=false;
      for(int icor=0;icor<_correlators.size();++icor){ // loop through correlators
        // CORRELATE HIT WITH TRACK
        double matchQuality=99999;
        bool survive=false;
        bool applicable=_correlators[icor]->corr(trk,rhit,survive,matchQuality);
        if(!applicable)continue; // try the next correlator
        if(!survive){
          trksurvive=false;
          break; // track has failed this level --> can be excluded
        }
        if(icor<maxlevel)continue; // there are cands that reached deeper level
        maxlevel=icor;
        if(_bestMatchQuality[icor]>matchQuality){
          _bestMatchQuality[icor]=matchQuality;
          _bestMatchIndex[icor]=itrk;
          trksurvive=true;
        }
      } // end loop over correlator
      foundAtAll|=trksurvive;
    } // end loop over tracks

    if(!foundAtAll)// new track
      {
        PndTpcRiemannTrack* trk=new PndTpcRiemannTrack();
        trk->setSort(_sortingMode);
        candlist.push_back(trk);
        //std::cout<<"Creating new track"<<std::endl;
        trk->addHit(rhit);
      }
    else {
      // add hit to best match
      // use the bestMatch from deepest level
      // std::cout<<"bestMatch[0]="<<_bestMatchIndex[0]
// 	       <<"   bestMatch[1]="<<_bestMatchIndex[1]<<std::endl;
//       std::cout<<"choosing "<<_bestMatchIndex[maxlevel]<<std::endl;
      PndTpcRiemannTrack* theTrk=candlist[_bestMatchIndex[maxlevel]];
      theTrk->addHit(rhit);
      if(theTrk->getNumHits()>=_minHitsForFit){
        theTrk->refit();
        theTrk->szFit();
      }
    }
    resetFlags();
  } // end loop over cluster
  

  std::cout<<candlist.size()<<" Riemann Tracks found."<<std::endl;
  for(int i=0;i<candlist.size();++i){
    std::cout<<"Track"<<i<<": "
  	     <<candlist[i]->getNumHits()<<" hits  R=";
    if(candlist[i]->getNumHits()>=_minHitsForFit)std::cout<<candlist[i]->r();
    std::cout<<std::endl;
    //candlist[i]->Plot(1);
  }
  
 return candlist.size();
}

void
PndTpcRiemannTrackFinder::mergeTracks(std::vector<PndTpcRiemannTrack*>& candlist){
  std::cout<<"PndTpcRiemannTrackFinder::mergeTracks"<<std::endl;

  unsigned int ntr=candlist.size();
  if (ntr<2) return; // need at least 2 trackcands to merge

  for(unsigned int itrk1=0; itrk1<ntr-1; ++itrk1){ // loop over tracks
    if(candlist[itrk1]==NULL)continue;
    PndTpcRiemannTrack* trk1=candlist[itrk1];

    for(unsigned int itrk2=itrk1+1; itrk2<ntr; ++itrk2){ // loop over the other tracks to be tested
      if(candlist[itrk2]==NULL)continue;

      PndTpcRiemannTrack* trk2=candlist[itrk2];

      // WE STEP THROUGH THE INDIVIDUAL CORRELATORS
      // IF THE TRACK trk2 SURVIVES EACH CORRELATOR
      // IT IS MERGED WITH THE TRACK trk1
      bool survive=true;
      for(unsigned int icor=0; icor<_TTcorrelators.size(); ++icor){ // loop through correlators
        // CORRELATE trk1 WITH trk2
        // all correlators must be applicable and must be survived!
        double matchQuality;
        bool applicable;

        // make sure that the first track in the correlator is the bigger one
        if (trk1->getNumHits()>=trk2->getNumHits())
          applicable=_TTcorrelators[icor]->corr(trk1,trk2,survive,matchQuality);
        else
          applicable=_TTcorrelators[icor]->corr(trk2,trk1,survive,matchQuality);

        if(!applicable) {
          survive = false;
          break;
        }
        if(!survive) break;
      } // end loop through correlators
      if (!survive) continue; // test next trk2

      // merge tracks if survived
      std::cout<<"merge track "<<itrk1<<" with track "<<itrk2<<std::endl;
      unsigned int nhits1 = trk1->getNumHits();
      unsigned int nhits2 = trk2->getNumHits();

      if(!_sortingMode){ // we have to collect clusters from both tracks, sort them and build a new track
        // collect clusters from both tracks and sort
        std::vector<PndTpcCluster*> clusters;
        for(unsigned int i=0; i<nhits1; ++i){
          clusters.push_back(trk1->getHit(i)->cluster());
        }
        for(unsigned int i=0; i<nhits2; ++i){
          clusters.push_back(trk2->getHit(i)->cluster());
        }
        sortClusters(clusters);

        // fill clusters into new RiemannTrack and refit
        PndTpcRiemannTrack* mergedTrack = new PndTpcRiemannTrack();
        mergedTrack->setSort(false);
        for(unsigned int i=0; i<clusters.size(); ++i){
          PndTpcRiemannHit* rhit = new PndTpcRiemannHit(clusters[i]);
          mergedTrack->addHit(rhit);
        }
        if(mergedTrack->getNumHits()>=_minHitsForFit){
          mergedTrack->refit();
          mergedTrack->szFit();
        }

        // delete old trackcands and store new trackcand
        delete candlist[itrk1];
        trk1=mergedTrack;
        candlist[itrk1]=mergedTrack;
      }
      else{ // we can just add the hits from trk2 to trk1 and the sorting is done internally
        for(unsigned int i=0; i<nhits2; ++i){
          trk1->addHit(trk2->getHit(i));
        }
        // refit if we have enough hits
        if(trk1->getNumHits()>=_minHitsForFit){
          trk1->refit();
          trk1->szFit();
        }
      }

      // delete trk2
      delete candlist[itrk2];
      candlist[itrk2]=NULL;

    } // end loop over the other tracks to be tested
  } // end loop over tracks

  // clean up candlist
  for(int i=0;i<candlist.size();++i){
    if (candlist[i]==NULL) {
      candlist.erase(candlist.begin()+i);
      --i; // go one step back because "erase" shifts back the rest
    }
  }

  std::cout<<candlist.size()<<" Merged Riemann Tracks: "<<std::endl;
  for(int i=0;i<candlist.size();++i){
    std::cout<<"Track "<<i<<": "
         <<candlist[i]->getNumHits()<<" hits  R=";
    if(candlist[i]->getNumHits()>=_minHitsForFit)std::cout<<candlist[i]->r();
    std::cout<<std::endl;
  }

}


void
PndTpcRiemannTrackFinder::sortClusters(std::vector<PndTpcCluster*>& cll){
  if(_sorting==-1) return;
  std::sort(cll.begin(),cll.end(),sortCluster);
}


bool
sortCluster(PndTpcCluster* s1, PndTpcCluster* s2)
{
  // -1: no sorting, 0: sort Clusters by X, 1: Y, 2: Z, 3: R, 4: distance to origin
  int sorting = PndTpcRiemannTrackFinder::getSorting();
  double a1;
  double a2;
  TVector3 d1;
  TVector3 d2;
  switch (sorting){
    case -1: //no sorting
      return false;
      break;
    case 0:
      a1=s1->pos().X();
      a2=s2->pos().X();
      return a1>a2;
      break;
    case 1:
      a1=s1->pos().Y();
      a2=s2->pos().Y();
      return a1>a2;
      break;
    case 2:
      a1=s1->pos().Z();
      a2=s2->pos().Z();
      return a1>a2;
      break;
    case 4:
      d1 = s1->pos();
      d1(2) -= PndTpcRiemannTrackFinder::getInteractionZ();
      d2 = s2->pos();
      d2(2) -= PndTpcRiemannTrackFinder::getInteractionZ();
      a1=d1.Mag();
      a2=d2.Mag();
      return a1>a2;
      break;
    default: // also case 3 -> R
      a1=s1->pos().Perp();
      a2=s2->pos().Perp();
      return a1>a2;
  }
}


bool
sortTracksN(PndTpcRiemannTrack* t1, PndTpcRiemannTrack* t2)
{
  int a1=t1->getNumHits();
  int a2=t2->getNumHits();
  return a1>a2;
}

bool
sortTracksFirstClusterPos(PndTpcRiemannTrack* t1, PndTpcRiemannTrack* t2)
{
  return sortCluster(t1->getFirstHit()->cluster(), t2->getFirstHit()->cluster());
}



void
PndTpcRiemannTrackFinder::resetFlags(){
// reset all flags
  for(int k=0;k<_correlators.size();++k){
    _found[k]=false;
    _bestMatchQuality[k]=99999;
    _bestMatchIndex[k]=0;
  }
}