// -------------------------------------------------------------------------
// -----                  CbmLitTrackFinder source file               -----
// -----                  Created 06/07/06  by A. Lebedev               -----
// -------------------------------------------------------------------------
 

#include "CbmLitTrackFinder.h"

#include "base/CbmTrackPropagator.h"
#include "base/CbmTrackUpdate.h"
#include "base/CbmTrackExtrapolator.h"
#include "CbmLitKalmanFilter.h"
#include "CbmLitTrackPropagator.h"
#include "CbmLitLineTrackExtrapolator.h"
#include "CbmLitRK4TrackExtrapolator.h"
#include "CbmLitParabolicTrackExtrapolator.h"
#include "CbmLitEnvironment.h"

#include "CbmRootManager.h"
#include "CbmRunAna.h"
#include "CbmRuntimeDb.h"
#include "CbmTrkHit.h"
#include "CbmTrdHit.h"

#include "TClonesArray.h"

#include <set>
#include <map>


// ----------------------- Default constructor ---------------------------
CbmLitTrackFinder::CbmLitTrackFinder():
   fIsStandalone(false)
{

   //TODO init of the parameters
   fEvents = 0;
   fVerbose = 1;
}
// -----------------------------------------------------------------------


// --------------------------- Destructor --------------------------------
CbmLitTrackFinder::~CbmLitTrackFinder()
{

}
// -----------------------------------------------------------------------

// ------------------------- Initialisation ------------------------------
//void CbmLitTrackFinder::Init()
//{
   
//}

// -----------------------------------------------------------------------



// -------------------- Algorithm implementation -------------------------
//Int_t CbmLitTrackFinder::DoFind(TClonesArray* hitArray,
 //                                     TClonesArray* trackArray)
//{
//   return 0;
//}
// -----------------------------------------------------------------------

// ------------------- Determines the Much geometry ------------------------
void CbmLitTrackFinder::ReadDetectorGeometry()
{

   
}
// -----------------------------------------------------------------------

// ------------------ Arrange Much hits by layers -------------------------
void CbmLitTrackFinder::ArrangeHits()
{

   fHitIdMap.clear();

   Int_t nofHits = fHitArray->GetEntriesFast();

   for (Int_t i = 0; i < fNofLayers; i++){
      fMaxErrY[i] = 0;
      fMaxErrX[i] = 0;
   }

   for(Int_t iHit = 0; iHit < nofHits; iHit++) {
   
      if (fUsedHitsSet.find(iHit) != fUsedHitsSet.end()) continue;
      
      //TODO use static_cast
      CbmTrkHit* hit = (CbmTrkHit*) fHitArray->At(iHit);
      
      if(NULL == hit) continue;
      
     Int_t layer = hit->GetStationNr() - 1;
                 
      if(layer < 0 || layer > fNofLayers) {
         std::cout << "-W- CbmLitTrackFinder::ArrangeHits: "
                   << "wrong layer number." << std::endl;
         continue;
      }
      
      if (fMaxErrX[layer] < hit->GetDx()) fMaxErrX[layer] = hit->GetDx();
      if (fMaxErrY[layer] < hit->GetDy()) fMaxErrY[layer] = hit->GetDy();
            
      //std::cout << hit->GetZ() << " ";
      
      fHits[layer].push_back(hit);
      fHitIdMap[hit] = iHit;
   }
     
   if (fVerbose > 1) {
      std::cout << "-I- CbmLitTrackFinder::ArrangeHits: " << std::endl;
      for (Int_t i = 0; i < fNofLayers; i++) {
          std::cout << "  Station " << i << " has " 
                    << fHits[i].size() << " hits" << std::endl;
      }
   }
   
   if (fVerbose > 1) {
      std::cout << "-I- CbmLitTrackFinder::ArrangeHits: "
                << std::endl << "   ";
   }

   // sort hits on each layer by x or y coordinate
   for (Int_t i = 0; i < fNofLayers; i++) {
      if (fMaxErrX[i] > fMaxErrY[i]) {
         std::sort(fHits[i].begin(), fHits[i].end(), CmpYUp);
       // fMaxErrY[i] *= 1e-4;
        fMaxErrX[i] = 0;
        if (fVerbose > 1) std::cout << "-y-";
      } else {
        sort(fHits[i].begin(), fHits[i].end(), CmpXUp);
       // fMaxErrX[i] *= 1e-4;
        fMaxErrY[i] = 0;
        if (fVerbose > 1) std::cout << "-x-";
      }
   }
   
   if (fVerbose > 1) std::cout << std::endl;
   
  
   if (fVerbose > 1) {
      std::cout << "-I- CbmLitTrackFinder::ArrangeHits: "
                << " max hit error: " << std::endl;
      for (Int_t i = 0; i < fNofLayers; i++)
         std::cout << "   fMaxErrX[" << i << "]=" << fMaxErrX[i]
                   << "   fMaxErrY[" << i << "]=" << fMaxErrY[i] << std::endl;
   }
   
}
// -----------------------------------------------------------------------

// ------------------ Track Following ------------------------------------
void CbmLitTrackFinder::TrackFollowing() 
{
   
   //loop over stations
   for (int iStation = 0; iStation < fNofStations; iStation++) { 
      
      if (iStation < fBeginStation) continue;
      if (iStation > fEndStation) continue;
      
    //int iStation = 0;
    //if ((iStation > 0)) {
    if ((iStation > 0 || !fIsStandalone)) {
    
    //TODO protect for standalone not to follow in the station
      
         
            for (std::vector<CbmLitTrack*>::iterator iTrack = fTracks.begin();
                  iTrack != fTracks.end(); iTrack++) {
            
               TrackFollowingStation(*iTrack, iStation);
            // delete (*iTrack);
            
            }     
      } 
      
      
      if (fVerbose > 1) {
         std::cout << "-I- CbmLitTrackFinder::TrackFollowing:" 
                   << fFoundTracks.size() << " followed in the "
                   << iStation << " station" << std::endl;
      }
      
      for(UInt_t i = 0; i < fTracks.size(); i++)
           delete fTracks[i];
      fTracks.clear();
      
      UpdateTracks(iStation);
      
      if (fVerbose > 1) {
         std::cout << "-I- CbmLitTrackFinder::TrackFollowing: "
                   << "Tracks Update Finished" << std::endl;
      }
      
      CheckCommonHits();
      
      if (fVerbose > 1) {
         std::cout << "-I- CbmLitTrackFinder::TrackFollowing: "
              << "Check Of the Common Hits Finished" << std::endl;
      }
      
      if (iStation < fEndStation) { 
         for (std::vector<CbmLitTrack*>::iterator iTrack = fFoundTracks.begin();
                  iTrack != fFoundTracks.end(); iTrack++) {
            if ( (*iTrack)->GetFlag() == 1) {
               delete (*iTrack);
            } else {
               fTracks.push_back(*iTrack);
            }
         }
         fFoundTracks.clear();
      }
   } //loop over stations
}
// -----------------------------------------------------------------------


// --------------------Track Following in the station---------------------
void CbmLitTrackFinder::TrackFollowingStation(CbmLitTrack *pTrack, 
                                              Int_t Station) 
{

   CbmLitTrack* pTracks[fNofLayersPerStation[Station]];
   for (Int_t i = 0; i < fNofLayersPerStation[Station]; i++) {
      pTracks[i] = new CbmLitTrack();
      pTracks[i]->SetParamLast(pTrack->GetParamLast());
   }
   
   Int_t layer = 0;
   for (Int_t i = 0; i < Station; i++) 
       layer += fNofLayersPerStation[i];

//   for (Int_t i = 0; i < fNofLayersPerStation[Station]; i++) {
//      if (Station != 0 || i != 0 ) {
//         fPropagator->Propagate(pTracks[i]->GetParamLast(), 
//                                    fLayerZPos[layer + i]);
//      }
//   }

  // std::cout << std::endl;
      
   CbmTrkHit* pHits[fNofLayersPerStation[Station]];

   Int_t IndMin[fNofLayersPerStation[Station]];
   Int_t IndMax[fNofLayersPerStation[Station]];
   Int_t HitCnt[fNofLayersPerStation[Station]];
   Int_t nofMissingHits[fNofLayersPerStation[Station]];
   for (Int_t i = 0; i < fNofLayersPerStation[Station]; i++) {
    //  MinMaxIndex(pTracks[i], IndMin[i], HitCnt[i], layer + i);
    //  IndMax[i] = IndMin[i] + HitCnt[i];
      nofMissingHits[i] = 0;
   }
   
   Int_t extraLoop = 0;
   if (fMaxNofMissingHitsInStation > 0) extraLoop = 1;        
   
   fPropagator->Propagate(pTracks[0]->GetParamLast(), 
                          fLayerZPos[layer + 0]);
   MinMaxIndex(pTracks[0], IndMin[0], HitCnt[0], layer + 0);
   IndMax[0] = IndMin[0] + HitCnt[0];
   
   for (Int_t iHit0 = IndMin[0]; iHit0 < IndMax[0] + extraLoop; iHit0++) { //1
   
      if (ProcessLayer(0, layer, iHit0, IndMax, nofMissingHits, 
                       pTracks, pHits) == false) continue;
      
      fPropagator->Propagate(pTracks[0]->GetParamLast(), pTracks[1]->GetParamLast(), 
                             fLayerZPos[layer + 1]);
      MinMaxIndex(pTracks[1], IndMin[1], HitCnt[1], layer + 1);
      IndMax[1] = IndMin[1] + HitCnt[1];
                         
      for (Int_t iHit1 = IndMin[1]; iHit1 < IndMax[1] + extraLoop; iHit1++) { //2
      
         if (ProcessLayer(1, layer, iHit1, IndMax, nofMissingHits, 
                          pTracks, pHits) == false) continue;
         
         if (fNofLayersPerStation[Station] < 3) {
            AddTrackCandidate(pTrack, pHits, Station);
            continue;
         }
         
         fPropagator->Propagate(pTracks[1]->GetParamLast(), pTracks[2]->GetParamLast(), 
                                      fLayerZPos[layer + 2]);
         MinMaxIndex(pTracks[2], IndMin[2], HitCnt[2], layer + 2);
         IndMax[2] = IndMin[2] + HitCnt[2];
         
         for (Int_t iHit2 = IndMin[2]; iHit2 < IndMax[2] + extraLoop; iHit2++) { //3
            
            if (ProcessLayer(2, layer, iHit2, IndMax, nofMissingHits, 
                               pTracks, pHits) == false) continue;
  
            if (fNofLayersPerStation[Station] < 4) {
               AddTrackCandidate(pTrack, pHits, Station);
               continue;
            }
            
            fPropagator->Propagate(pTracks[2]->GetParamLast(), pTracks[3]->GetParamLast(), 
                                         fLayerZPos[layer + 3]);
            MinMaxIndex(pTracks[3], IndMin[3], HitCnt[3], layer + 3);
            IndMax[3] = IndMin[3] + HitCnt[3];
            
            
            for (Int_t iHit3 = IndMin[3]; iHit3 < IndMax[3] + extraLoop; iHit3++) { //4
            
               if (ProcessLayer(3, layer, iHit3, IndMax, nofMissingHits, 
                                  pTracks, pHits) == false) continue;
            
               AddTrackCandidate(pTrack, pHits, Station);
            
            } //4
         } //3
      } //2
   } //1
   
   for (Int_t i = 0; i < fNofLayersPerStation[Station]; i++)
      delete pTracks[i];

}
// -----------------------------------------------------------------------

// ----------------Process the layer in the station-----------------------
Bool_t CbmLitTrackFinder::ProcessLayer(
                              Int_t layerInStation,
                              Int_t layer,
                              Int_t iHit,
                              Int_t IndMax[],
                              Int_t nofMissingHits[],  
                              CbmLitTrack* pTracks[],          
                              CbmTrkHit* pHits[])
{
   Bool_t result = true;
   
   if (layerInStation == 0) nofMissingHits[0] = 0; 
   else nofMissingHits[layerInStation] = nofMissingHits[layerInStation - 1];
            
   if (iHit < IndMax[layerInStation]) {
      pHits[layerInStation] = fHits[layer + layerInStation][iHit];
      if (IsIn(pTracks[layerInStation], pHits[layerInStation]) == false) result = false;
      if (fApplyUpdateInLayer && result) fFilter->Update(pTracks[layerInStation]->GetParamLast(), pHits[layerInStation]);
   } else {
      pHits[layerInStation] = NULL;
      nofMissingHits[layerInStation]++;
      if (nofMissingHits[layerInStation] > fMaxNofMissingHitsInStation) result = false;
   }
   
   return result;
}
// -----------------------------------------------------------------------

// ----------------Checks if the hit in the area near the track-----------
Bool_t CbmLitTrackFinder::IsIn(CbmLitTrack *pTrack, CbmTrkHit *pHit) 
{
   CbmTrackParam *par = pTrack->GetParamLast();
   
   Double_t x1 = par->GetX();
   Double_t x2 = pHit->GetX();
         
   Double_t y1 = par->GetY();
   Double_t y2 = pHit->GetY();
   
   Double_t devX;
   Double_t devY;
   
   if (!fPrecalcSearchRegions) {
      //CbmTrackParam *par = pTrack->GetParamLast();
      Double_t Cov00 = TMath::Abs(par->GetCovariance(0, 0));
      Double_t Cov11 = TMath::Abs(par->GetCovariance(1, 1));
      
      if( Cov00 > 100. ) return kFALSE;
      if( Cov11 > 100. ) return kFALSE;  
   
      //Double_t x1 = par->GetX();
      //Double_t x2 = pHit->GetX();
      Double_t dx2 = pHit->GetDx();
      
      //Double_t y1 = par->GetY();
      //Double_t y2 = pHit->GetY();
      Double_t dy2 = pHit->GetDy();      
      
      devX = fSigmaCoef * TMath::Sqrt(Cov00 + dx2 * dx2);      
      devY = fSigmaCoef * TMath::Sqrt(Cov11 + dy2 * dy2);
    // devX = fSigmaCoef * (sqrt(Cov00) + dx2);      
    // devY = fSigmaCoef * (sqrt(Cov11) + dy2);
      
   } else {
      Int_t layer = pHit->GetStationNr() - 1;
      devX = fSigmaCoef * fSigmaX[layer];
      devY = fSigmaCoef * fSigmaY[layer];
   }
  
   return ( ( (x1 + devX) >= x2 ) &&
            ( (x1 - devX) <= x2 ) &&
            ( (y1 + devY) >= y2 ) &&
            ( (y1 - devY) <= y2 ) );
}
// -----------------------------------------------------------------------

// ---------------tracks update of the tracks----------------------
void CbmLitTrackFinder::UpdateTracks(Int_t Station) 
{
  
   for ( std::vector<CbmLitTrack*>::iterator iTrack = fFoundTracks.begin();
            iTrack != fFoundTracks.end(); iTrack++) {
      
      (*iTrack)->SetChi2(0.0);
      
      Int_t nofHits = (*iTrack)->GetNofHits();
      
      for (Int_t iHit = 0; iHit < nofHits; iHit++) {
         CbmTrkHit * pHit = (CbmTrkHit*) fHitArray->At((*iTrack)
                               ->GetHitIndex(iHit));
                  
         Double_t Ze = pHit->GetZ();
         
         fPropagator->Propagate((*iTrack)->GetParamLast(), Ze);
         
         fFilter->Update((*iTrack)->GetParamLast(), pHit);
         (*iTrack)->SetChi2((*iTrack)->GetChi2() + CalcChi2((*iTrack)->GetParamLast(), pHit));
                  
      }
      // TODO check NDF
      if (nofHits > 2) (*iTrack)->SetNDF( 2 * nofHits - 5);
      else (*iTrack)->SetNDF(1);
      //(*iTrack)->SetChi2( (*iTrack)->GetChi2() / (Double_t) nofHits ); 
   }
   
   
}
// -----------------------------------------------------------------------

// ------------Checks found tracks for common hits------------------------
void CbmLitTrackFinder::CheckCommonHits() 
{

   std::sort(fFoundTracks.begin(), fFoundTracks.end(), CompareChi2);
   
   //set to store hits id,
   //to check if the hit was already used by another track
   std::set<Int_t> HitsId;
   Int_t cnt = 0;
   
   for (std::vector<CbmLitTrack*>::iterator iTrack = fFoundTracks.begin();
         iTrack != fFoundTracks.end(); iTrack++) {
      
      (*iTrack)->SetFlag(0); 
    
      Int_t nofSharedHits = 0;
      
      Int_t NofHits = (*iTrack)->GetNofHits();
      
      for(Int_t iHit = 0; iHit < NofHits; iHit++) {
         Int_t HitId = (*iTrack)->GetHitIndex(iHit);
         
         if(HitsId.find(HitId) != HitsId.end()) {
            nofSharedHits++;
            if (nofSharedHits > fNofSharedHits) {
               (*iTrack)->SetFlag(1);
               break;
            }
         }
      }
          
   
      if ( (*iTrack)->GetFlag() == 1 ) continue;
      
      for(Int_t iHit = 0; iHit < NofHits; iHit++) {
         Int_t HitId = (*iTrack)->GetHitIndex(iHit);
         HitsId.insert(HitId);
      }
      cnt++;      
   }
   
   if (fVerbose > 1) {
      std::cout << "-I- CbmLitTrackFinder::CheckCommonHits(): "
         << "count good = " << cnt << " " << "HitsId.size() = "
         << HitsId.size() << std::endl;
   }
   
   HitsId.clear();
}
// -----------------------------------------------------------------------

// ---------Defines min and max index of hits-----------------------------
void CbmLitTrackFinder::MinMaxIndex(CbmLitTrack* pTrack, 
                                    Int_t &IndMin,
                                    Int_t &HitCnt, 
                                    Int_t layer)
{

   std::vector<CbmTrkHit*>::iterator it_min;
   std::vector<CbmTrkHit*>::iterator it_max;
   CbmTrkHit* hit = new CbmTrdHit();
   
   CbmTrackParam* par = pTrack->GetParamLast();
   
   if (fMaxErrY[layer] != 0) {
      Double_t devY;
      if (!fPrecalcSearchRegions) {
         Double_t Cov11 = TMath::Abs(par->GetCovariance(1, 1));
         if (TMath::Abs(Cov11) > 100.) devY = 0;
         else devY = fSigmaCoef * TMath::Sqrt( Cov11 + 
                   fMaxErrY[layer] * fMaxErrY[layer] );
        //else devY = fSigmaCoef * (sqrt(Cov11) + fMaxErrY[layer]);
      } else {
         devY = fSigmaCoef * fSigmaY[layer];
      }
                        
      hit->SetY(par->GetY() - devY);
      it_min = lower_bound(fHits[layer].begin(), fHits[layer].end(),
                               hit, CmpYUp);
     
      hit->SetY(par->GetY() + devY);
      it_max = lower_bound(fHits[layer].begin(), fHits[layer].end(),
                               hit, CmpYUp) - 1;
   } else {
      Double_t devX;
      if (!fPrecalcSearchRegions) {
         Double_t Cov00 = TMath::Abs(par->GetCovariance(0, 0));
         if (TMath::Abs(Cov00) > 100.) devX = 0;
         else devX = fSigmaCoef * TMath::Sqrt( Cov00 + 
                    fMaxErrX[layer] * fMaxErrX[layer] );
        // else devX = fSigmaCoef * (sqrt(Cov00) + fMaxErrX[layer]);
      } else {
         devX = fSigmaCoef * fSigmaY[layer];
      } 
      
      hit->SetX(par->GetX() - devX);
      it_min = lower_bound(fHits[layer].begin(), fHits[layer].end(),
                              hit, CmpXUp);
      
      hit->SetX(par->GetX() + devX);
      it_max = lower_bound(fHits[layer].begin(), fHits[layer].end(),
                              hit, CmpXUp) - 1 ;
   }

   IndMin = it_min - fHits[layer].begin();
   //IndMin = distance(fHits[layer].begin(), it_min);
   HitCnt = it_max - it_min + 1;
   //HitCnt = distance(it_min, it_max) + 1;
   
   delete hit;
}
// -----------------------------------------------------------------------

// -------------Adds track candidate--------------------------------------
void CbmLitTrackFinder::AddTrackCandidate(CbmLitTrack* pTrack,
                                          CbmTrkHit* pHits[], 
                                          Int_t Station)
{
   // Add new track to the fFoundTracks vector
   CbmLitTrack* pNewTrack = new CbmLitTrack(*pTrack);

   for (Int_t i = 0; i < fNofLayersPerStation[Station]; i++) {
      // Int_t hitId = fArrayMuchHit->IndexOf(pHits[i]);
      if (pHits[i] == NULL) continue;
      Int_t hitId = fHitIdMap[pHits[i]];
      pNewTrack->AddHit(hitId, pHits[i]);
   }
     
   pNewTrack->SortHits(); 
     
   if (pNewTrack->GetNofHits() == 0) {
      delete pNewTrack;
      return;
   }   

   fFoundTracks.push_back(pNewTrack);
   
   pNewTrack = NULL; /// ?!?!
}

// -----------------------------------------------------------------------
Double_t CbmLitTrackFinder::CalcChi2(CbmTrackParam *pParam, CbmTrkHit* pHit)
{
   Double_t errX2 = pHit->GetDx() * pHit->GetDx();
   Double_t errY2 = pHit->GetDy() * pHit->GetDy();
   Double_t C0 = pParam->GetCovariance(0, 0);
   Double_t C5 = pParam->GetCovariance(1, 1);
   Double_t C1 = pParam->GetCovariance(1, 0);
   
   Double_t dx = pHit->GetX() - pParam->GetX(); 
   Double_t dy = pHit->GetY() - pParam->GetY();
  
   Double_t denominator = errX2 * errY2 - errX2 * C5 - C0 * errY2 + 
                          C0 * C5 - C1 * C1;
   
   Double_t chi2 = ((dx * (errY2 - C5) + dy * C1) * dx + 
                    (dy * (errX2 - C0) + dx * C1) * dy) / denominator;
   
   return chi2;
}
// -----------------------------------------------------------------------

// -----------------------------------------------------------------------

// ----------------Removes hits which belong to found tracks--------------
void CbmLitTrackFinder::RemoveHits()
{
   for(std::vector<CbmLitTrack*>::iterator iTrack = fFoundTracks.begin();
         iTrack != fFoundTracks.end(); iTrack++) {
      if( (*iTrack)->GetFlag() == 0 ) {
      
         Int_t nofHits = (*iTrack)->GetNofHits();
         
         for (Int_t iHit = 0; iHit < nofHits; iHit++) {
            Int_t HitId = (*iTrack)->GetHitIndex(iHit);
               
            fUsedHitsSet.insert(HitId);               
         }
      }
   }
   
   if (fVerbose > 1) {
      std::cout << "-I- CbmLitTrackFinder::RemoveHits(): "
                << "fUsedHitsSet.size() = " 
                << fUsedHitsSet.size() << std::endl;
   }
}
// -----------------------------------------------------------------------

ClassImp(CbmLitTrackFinder);