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


#include <iostream>
#include <set>

#include "TClonesArray.h"

#include "CbmRootManager.h"
#include "CbmRunAna.h"
#include "CbmBaseParSet.h"
#include "CbmRuntimeDb.h"
#include "CbmDetector.h"
#include "CbmTrdPoint.h"
#include "CbmTrdHit.h"
#include "CbmStsTrack.h"
#include "CbmTrdTrack.h"

#include "CbmLitEnvironment.h"

#include "CbmLitTrdTrackFinder.h"


// ----------------------- Default constructor ---------------------------
CbmLitTrdTrackFinder::CbmLitTrdTrackFinder()
{
   fEvents = 0;
   fVerbose = 1;
}
// -----------------------------------------------------------------------


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

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


// ------------------- Determines the TRD geometry ------------------------
void CbmLitTrdTrackFinder::ReadTrdGeometry()
{
   CbmLitEnvironment* env = CbmLitEnvironment::Instance();
   
   env->GetTrdLayout(fNofTrdStations, fNofTrdLayers, fNofTrdLayersPerStation);

   fMaxSigmaX.resize(fNofTrdLayers);
   fMaxSigmaY.resize(fNofTrdLayers);
}
// -----------------------------------------------------------------------


// ------------------ Arrange TRD hits by layers -------------------------
void CbmLitTrdTrackFinder::ArrangeHits()
{

   fHitIdMap.clear();

   Int_t nTrdHits = fArrayTrdHit->GetEntriesFast();


   for(Int_t iHit = 0; iHit < nTrdHits; iHit++) {
   
      if (fUsedHitsSet.find(iHit) != fUsedHitsSet.end()) continue;
      
      CbmTrdHit* hit = (CbmTrdHit*) fArrayTrdHit->At(iHit);
      
      if(NULL == hit) continue;
      
      Int_t layer = hit->GetPlaneID() - 1;
      if(layer < 0 || layer > fNofTrdLayers) {
         cout << "-W- CbmLitTrdTrackFinder::ArrangeHits: "
              << "wrong layer number." << endl;
         continue;
      }
      
      fHits[layer].push_back(hit);
      fHitIdMap[hit] = iHit;
   }
     
   if (fVerbose > 1) {
      cout << "-I- CbmLitTrdTrackFinder::ArrangeHits: " << endl;
      for(int i = 0; i < fNofTrdLayers; i++) {
         cout << "   TRD layer " << i << " has " 
            << fHits[i].size() << " CbmTrdHits" << endl;
      }
   }
   
   if (fVerbose > 1) {
      cout << "-I- CbmLitTrdTrackFinder::ArrangeHits: "
           << endl << "   ";
   }

   // sort hits on each layer by x or y coordinate
   for (int i = 0; i < fNofTrdLayers; i++)
   if (fHits[i][0]->GetDx() > fHits[i][0]->GetDy()) {
      sort(fHits[i].begin(), fHits[i].end(), CmpYUp);
      fMaxSigmaY[i] = fHits[i][0]->GetDy() * 1e-4;
      fMaxSigmaX[i] = 0;
      if (fVerbose > 1) cout << "-y-";
   } else {
      sort(fHits[i].begin(), fHits[i].end(), CmpXUp);
      fMaxSigmaX[i] = fHits[i][0]->GetDx() * 1e-4;
      fMaxSigmaY[i] = 0;
      if (fVerbose > 1) cout << "-x-";
   }
   
   if (fVerbose > 1) cout << endl;
   
   
   if (fVerbose > 1) {
      cout << "-I- CbmLitTrdTrackFinder::ArrangeHits: "
           << " max sigma: " << endl;
      for (Int_t i = 0; i < fNofTrdLayers; i++)
      cout << "   fMaxSigmaX[" << i << "]=" << fMaxSigmaX[i]
           << "   fMaxSigmaY[" << i << "]=" << fMaxSigmaY[i] << endl;
   }
}
// -----------------------------------------------------------------------

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


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

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

   for (Int_t i = 0; i < fNofTrdLayersPerStation[Station]; i++) {
      if (Station != 0 || i != 0 ) {
         fLinePropagator->Propagate(pTracks[i]->GetParamLast(), 
                                    fHits[layer + i][0]->GetZ());
      }
   }

   CbmTrdHit* pHits[fNofTrdLayersPerStation[Station]];

   Int_t IndMin[fNofTrdLayersPerStation[Station]];
   Int_t IndMax[fNofTrdLayersPerStation[Station]];
   Int_t HitCnt[fNofTrdLayersPerStation[Station]];
   Int_t nofMissingHits[fNofTrdLayersPerStation[Station]];
   for (Int_t i = 0; i < fNofTrdLayersPerStation[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;
          
   for (Int_t iHit0 = IndMin[0]; iHit0 < IndMax[0] + extraLoop; iHit0++) { //1
   
      if (ProcessLayer(0, layer, iHit0, IndMax, nofMissingHits, 
                         pTracks, pHits) == false) continue;
      
      for (Int_t iHit1 = IndMin[1]; iHit1 < IndMax[1] + extraLoop; iHit1++) { //2
      
         if (ProcessLayer(1, layer, iHit1, IndMax, nofMissingHits, 
                            pTracks, pHits) == false) continue;
         
         if (fNofTrdLayersPerStation[Station] < 3) {
            AddTrackCandidate(pTrack, pHits, Station);
            continue;
         }
         
         for (Int_t iHit2 = IndMin[2]; iHit2 < IndMax[2] + extraLoop; iHit2++) { //3
            
            if (ProcessLayer(2, layer, iHit2, IndMax, nofMissingHits, 
                               pTracks, pHits) == false) continue;
  
            if (fNofTrdLayersPerStation[Station] < 4) {
               AddTrackCandidate(pTrack, pHits, Station);
               continue;
            }
            
            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 < fNofTrdLayersPerStation[Station]; i++)
      delete pTracks[i];

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

// ----------------Process the layer in the station-----------------------
Bool_t CbmLitTrdTrackFinder::ProcessLayer(
                              Int_t layerInStation,
                              Int_t layer,
                              Int_t iHit,
                              Int_t IndMax[],
                              Int_t nofMissingHits[],  
                              CbmTrdTrack* pTracks[],          
                              CbmTrdHit* 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;
   } 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 CbmLitTrdTrackFinder::IsIn(CbmTrdTrack *pTrack, CbmTrdHit *pHit) 
{

   CbmTrackParam *par = pTrack->GetParamLast();
   Double_t Cov00 = par->GetCovariance(0, 0);
   Double_t Cov11 = 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() * 1e-4;  
   
   Double_t y1 = par->GetY();
   Double_t y2 = pHit->GetY();
   Double_t dy2 = pHit->GetDy() * 1e-4;
   
   Double_t devX = fSigmaCoef * TMath::Sqrt(Cov00 + dx2 * dx2);
   
   // for standalone track-finder (3x3 geometry) 
   Int_t layer = pHit->GetPlaneID() - 1;
   if (fIsStandalone && fNofTrdLayers == 9 && 
          (layer == 3 || layer == 4 || layer == 5)) 
              devX = fSigmaCoef * fSigmaX_2[layer - 3];
   //
   
   Double_t devY = fSigmaCoef * TMath::Sqrt(Cov11 + dy2 * dy2);

   return ( ( (x1 + devX) >= x2 ) &&
            ( (x1 - devX) <= x2 ) &&
            ( (y1 + devY) >= y2 ) &&
            ( (y1 - devY) <= y2 ) );
}
// -----------------------------------------------------------------------

// ---------------Kalman Filter update of the tracks----------------------
void CbmLitTrdTrackFinder::UpdateTracks(Int_t Station) 
{

 //  Int_t layerStart = 0;
 //  for (Int_t i = 0; i < Station; i++) 
 //        layerStart += fNofTrdLayersPerStation[i]; 
            
   
   for ( std::vector<CbmTrdTrack*>::iterator iTrack = fFoundTracks.begin();
            iTrack != fFoundTracks.end(); iTrack++) {
      
      (*iTrack)->SetChi2(0.0);
      
      Int_t nofHits = (*iTrack)->GetNofTrdHits();
      
      
      
      for (Int_t iHit = 0; iHit < nofHits; iHit++) {
         CbmTrdHit * pHit = (CbmTrdHit*) fArrayTrdHit->At((*iTrack)
                                ->GetTrdHitIndex(iHit));
                  
         Double_t Ze = pHit->GetZ();
         
         fLinePropagator->Propagate((*iTrack)->GetParamLast(), Ze);
         
         fFilter->Filter((*iTrack)->GetParamLast(), pHit);
         (*iTrack)->SetChi2((*iTrack)->GetChi2() + CalcChi2((*iTrack)->GetParamLast(), pHit));
                  
      }
      // TODO
      (*iTrack)->SetChi2( (*iTrack)->GetChi2() / (Double_t) nofHits ); 
      /*
      Int_t layer;
      for (Int_t iLayer = 0; iLayer < fNofTrdLayersPerStation[Station];
            iLayer++) {
         
         layer = layerStart + iLayer;

         CbmTrdHit * pHit = (CbmTrdHit*) fArrayTrdHit->At((*iTrack)
                                ->GetTrdHitIndex(layer));
                  
         Double_t Ze = pHit->GetZ();

         if (layer > 0) fLinePropagator->Propagate((*iTrack)->GetParamLast(), Ze);

         //fKf->Filter(*iTrack, pHit);
         fFilter->Filter((*iTrack)->GetParamLast(), pHit);
         (*iTrack)->SetChi2((*iTrack)->GetChi2() + CalcChi2((*iTrack)->GetParamLast(), pHit));
      }
      */
   }
}
// -----------------------------------------------------------------------

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

   std::sort(fFoundTracks.begin(), fFoundTracks.end(), CompareChi2);
   
   //set to store hits id,
   //to check if the hit was already used by another track
   set<Int_t> HitsId;
   Int_t cnt = 0;

   for (std::vector<CbmTrdTrack*>::iterator iTrack = fFoundTracks.begin();
         iTrack != fFoundTracks.end(); iTrack++) {
   
      Int_t NofHits = (*iTrack)->GetNofTrdHits();
      
      for(Int_t iHit = 0; iHit < NofHits; iHit++) {
         Int_t HitId = (*iTrack)->GetTrdHitIndex(iHit);
         
         if(HitsId.find(HitId) != HitsId.end()) {
            (*iTrack)->SetFlag(1);
            break;
         }
      }
   
      if ( (*iTrack)->GetFlag() ) continue;
      
      for(Int_t iHit = 0; iHit < NofHits; iHit++) {
         Int_t HitId = (*iTrack)->GetTrdHitIndex(iHit);
         HitsId.insert(HitId);
      }
      cnt++;
   }
   
   if (fVerbose > 1) {
      cout << "-I- CbmLitTrdTrackFinder::CheckCommonHits(): "
         << "count good = " << cnt << " " << "HitsId.size() = "
         << HitsId.size() << endl;
   }
   
   HitsId.clear();
}
// -----------------------------------------------------------------------

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

   std::vector<CbmTrdHit*>::iterator it_min;
   std::vector<CbmTrdHit*>::iterator it_max;
   CbmTrdHit* hit = new CbmTrdHit();
   
   CbmTrackParam* par = pTrack->GetParamLast();
   
   if (fMaxSigmaY[layer] != 0) {
      Double_t Cov11 = par->GetCovariance(1, 1);
      Double_t devY = fSigmaCoef * TMath::Sqrt( Cov11 + 
                        fMaxSigmaY[layer] * fMaxSigmaY[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 Cov00 = par->GetCovariance(0, 0);
      Double_t devX = fSigmaCoef * TMath::Sqrt( Cov00 + 
                        fMaxSigmaX[layer] * fMaxSigmaX[layer] );
      
      //for standalone (3x3 geometry)
      if (fIsStandalone && fNofTrdLayers == 9 && 
            (layer == 3 || layer == 4 || layer == 5)) 
               devX = fSigmaCoef * fSigmaX_2[layer - 3];
      //
      
      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 CbmLitTrdTrackFinder::AddTrackCandidate(CbmTrdTrack* pTrack,
                                             CbmTrdHit* pHits[], 
                                             Int_t Station)
{
   // Add new track to the fFoundTracks vector
   CbmTrdTrack* pNewTrack = new CbmTrdTrack(*pTrack);

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

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

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

// -----------------------------------------------------------------------
Double_t CbmLitTrdTrackFinder::CalcChi2(CbmTrackParam *pParam, CbmHit* pHit)
{
   Double_t errX2 = pHit->GetDx() * pHit->GetDx() * 1e-8;
   Double_t errY2 = pHit->GetDy() * pHit->GetDy() * 1e-8;
   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;
}
// -----------------------------------------------------------------------

ClassImp(CbmLitTrdTrackFinder);