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


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


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

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


// ------------------- Determines the TRD geometry ------------------------
void CbmLitTrdTrackFinder::ReadTrdGeometry()
{
   // Get the pointer to the CbmRunAna object
   CbmRunAna* ana = CbmRunAna::Instance();
   if(NULL == ana) {
      cout << "-E- CbmLitTrdTrackFinder::ReadTrdGeometryt :"
           <<" no CbmRunAna object!" << endl;
      return;
   }
   
   // Get the pointer to run-time data base
   CbmRuntimeDb* rtdb = ana->GetRuntimeDb();
   if(NULL == rtdb) {
      cout << "-E- CbmLitTrdTrackFinder::ReadTrdGeometry :"
           <<" no runtime database!" << endl;
      return;
   }
   
   // Get the pointer to container of base parameters
   CbmBaseParSet* baseParSet =
      (CbmBaseParSet*) rtdb->getContainer("CbmBaseParSet");
   if(NULL == baseParSet) {
      cout << "-E- CbmLitTrdTrackFinder::ReadTrdGeometry :"
           <<" no container of base parameters!" << endl;
      return;
   }
   
   // Get the pointer to detector list
   TObjArray *detList = baseParSet->GetDetList();
   if(NULL == detList) {
      cout << "-E- CbmLitTrdTrackFinder::ReadTrdGeometry :"
           << " no detector list!" << endl;
      return;
   }
   
   // Find TRD detector
   CbmDetector* trd = (CbmDetector*) detList->FindObject("TRD");
   if(NULL == trd) {
      cout << "-E- CbmLitTrdTrackFinder::ReadTrdGeometry :"
           << " no TRD detector!" << endl;
      return;
   }
   // Determine the geometry version
   TString name = trd->GetGeometryFileName();
   if(name.Contains("9")) {
      cout << "-I- CbmLitTrdTrackFinder::ReadTrdGeometry :"
           << " TRD geometry : 3x3." << endl;
      fNofTrdStations = 3;
      fNofTrdLayersPerStation.resize(3);
      for (Int_t i = 0; i < 3; i++) fNofTrdLayersPerStation[i] = 3;
      fNofTrdLayers = 9;
   } else if(name.Contains("12") || name.Contains("standard")) {
      cout << "-I- CbmLitTrdTrackFinder::ReadTrdGeometry :"
           << " TRD geometry : 3x4." << endl;
      fNofTrdStations = 3;
      fNofTrdLayersPerStation.resize(3);
      for (Int_t i = 0; i < 3; i++) fNofTrdLayersPerStation[i] = 4;
      fNofTrdLayers = 12;
   } else if(name.Contains("6x2")) {
      cout << "-I- CbmLitTrdTrackFinder::ReadTrdGeometry :"
           << " TRD geometry : 6x2." << endl;
      fNofTrdStations = 6;
      fNofTrdLayersPerStation.resize(6);
      for (Int_t i = 0; i < 6; i++) fNofTrdLayersPerStation[i] = 2;
      fNofTrdLayers = 12;
	} else if(name.Contains("trd_special")) {
      cout << "-I- CbmLitTrdTrackFinder::ReadTrdGeometry :"
      << " TRD geometry : 4-3-3." << endl;
      fNofTrdStations = 3;
      fNofTrdLayersPerStation.resize(3);
      fNofTrdLayersPerStation[0] = 4;
      fNofTrdLayersPerStation[1] = 3;
      fNofTrdLayersPerStation[2] = 3;
      fNofTrdLayers = 10;
   }


   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 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() 
{

   // linear extrapolation
   fKf->SetMethod(0);
   
   //loop over stations
   for (int iStation = 0; iStation < fNofTrdStations; iStation++) { 
   
      //int iStation = 0;
      if (iStation > 0 || !fIsStandalone) {
      
         for (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;
      }
      
      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 < fNofTrdStations - 1) {
         for (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) 
{

   Int_t IndMin[fNofTrdLayersPerStation[Station]];
   Int_t HitCnt[fNofTrdLayersPerStation[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 ) {
         fKf->Propagate(pTracks[i], fHits[layer + i][0]->GetZ());
      }
   }

   CbmTrdHit* pHits[fNofTrdLayersPerStation[Station]];

   MinMaxIndex(pTracks[0], IndMin[0], HitCnt[0], layer);
   for (Int_t iHit0 = IndMin[0]; iHit0 < IndMin[0] + HitCnt[0]; iHit0++) { //1
   
      pHits[0] = fHits[layer][iHit0];
      
      if (IsIn(pTracks[0], pHits[0]) == false) continue;
      
      MinMaxIndex(pTracks[1], IndMin[1], HitCnt[1], layer + 1);
      for (Int_t iHit1 = IndMin[1]; iHit1 < IndMin[1] + HitCnt[1]; iHit1++) { //2
      
         pHits[1] = fHits[layer + 1][iHit1];
         
         if (IsIn(pTracks[1], pHits[1]) == false) continue;
         
         if (fNofTrdLayersPerStation[Station] < 3) {
            AddTrackCandidate(pTrack, pHits, Station);
            continue;
         }
         
         MinMaxIndex(pTracks[2], IndMin[2], HitCnt[2], layer + 2);
         for (Int_t iHit2 = IndMin[2]; iHit2 < IndMin[2] + HitCnt[2]; iHit2++) { //3
         
            pHits[2] = fHits[layer + 2][iHit2];
            
            if (IsIn(pTracks[2], pHits[2]) == false) continue;
            
            if (fNofTrdLayersPerStation[Station] < 4) {
               AddTrackCandidate(pTrack, pHits, Station);
               continue;
            }
            
            MinMaxIndex(pTracks[3], IndMin[3], HitCnt[3], layer + 3);
            for (Int_t iHit3 = IndMin[3]; iHit3 < IndMin[3] + HitCnt[3]; iHit3++) { //4
            
               pHits[3] = fHits[layer + 3][iHit3];
                                 
               if (IsIn(pTracks[3], pHits[3]) == false)  continue;
                        
               AddTrackCandidate(pTrack, pHits, Station);
            
            } //4
         } //3
      } //2
   } //1
   
   for (int i = 0; i < fNofTrdLayersPerStation[Station]; i++)
   delete pTracks[i];

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

// ----------------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 ( vector<CbmTrdTrack*>::iterator iTrack = fFoundTracks.begin();
            iTrack != fFoundTracks.end(); iTrack++) {

      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) fKf->Propagate(*iTrack, Ze);

         fKf->Filter(*iTrack, pHit);
      }
   }
}
// -----------------------------------------------------------------------

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

   sort(fFoundTracks.begin(), fFoundTracks.end(), CbmTrdTrack::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 (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)
{

   vector<CbmTrdHit*>::iterator it_min;
   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 i = 0; i < fNofTrdLayersPerStation[Station]; i++) {
      // Int_t hitId = fArrayTrdHit->IndexOf(pHits[i]);
      Int_t hitId = fHitIdMap[pHits[i]];
      pNewTrack->AddHit(hitId, pHits[i]);
   }
   
   pNewTrack->SortHits();
   fFoundTracks.push_back(pNewTrack);
}
// -----------------------------------------------------------------------

// ----------------Removes hits which belong to found tracks--------------
void CbmLitTrdTrackFinder::RemoveHits()
{
   for(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;
   }
}
// -----------------------------------------------------------------------


ClassImp(CbmLitTrdTrackFinder);