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


// ----------------------- Default constructor ---------------------------
CbmLitTrdTrackFinderSts::CbmLitTrdTrackFinderSts()
{
   fIsStandalone = false;
}
// -----------------------------------------------------------------------


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

   delete fFieldPropagator;
   delete fFieldExtrapolator;
   delete fLinePropagator;
   delete fLineExtrapolator;
   delete fFilter;
}
// -----------------------------------------------------------------------



// ------------------------- Initialisation ------------------------------
void CbmLitTrdTrackFinderSts::Init()
{
   // Activate data branches
   // Get pointer to the ROOT manager
   CbmRootManager* rootMgr = CbmRootManager::Instance();
   if(NULL == rootMgr) {
      cout << "-E- CbmLitTrdTrackFinderSts::Init : "
           << "ROOT manager is not instantiated" << endl;
      return;
   }
   
   // Get pointer to the STS track array
   fArrayStsTrack = (TClonesArray*) rootMgr->GetObject("STSTrack");
   if(NULL == fArrayStsTrack) {
      cout << "-E- CbmLitTrdTrackFinderSts::Init : "
           << "no STS track array" << endl;
      return;
   }
   
   fFieldExtrapolator = new CbmLitRK4TrackExtrapolator();
   fFieldPropagator = new CbmLitTrackPropagator(fFieldExtrapolator);
   if(NULL == fFieldPropagator) {
      cout << "-E- CbmLitTrdTrackFinderSts::Init : "
           << "field propagator is not initialized." << endl;
   }
   fFieldPropagator->Init();
   
   fLineExtrapolator = new CbmLitLineTrackExtrapolator();
   fLinePropagator = new CbmLitTrackPropagator(fLineExtrapolator);
   if(NULL == fLinePropagator) {
      cout << "-E- CbmLitTrdTrackFinderSts::Init : "
           << "line propagator is not initialized." << endl;
   }
   fLinePropagator->Init();
   
   fFilter = new CbmLitKalmanFilter();
   if(NULL == fFilter) {
      cout << "-E- CbmLitTrdTrackFinderSts::Init : "
           << "filter is not initialized." << endl;
   }
   fFilter->Init();
   
   // Determine the TRD geometry
   ReadTrdGeometry();

   fHits.clear();
   fHits.resize(fNofTrdLayers);
}
// -----------------------------------------------------------------------


// -------------------- Algorithm implementation -------------------------
Int_t CbmLitTrdTrackFinderSts::DoFind(TClonesArray* hitArray,
                                      TClonesArray* trackArray)
{
  // return 0;
   
   cout << "==============================================================================" << endl;
   cout << "-I- CbmLitTrdTrackFinderSts::DoFind: ACTIVE" << endl;
   
   
   fArrayTrdHit = hitArray;

   fTrdTracks.clear();
   fFoundTracks.clear();
   fStsIdSet.clear();


   fUsedHitsSet.clear();

   fHits.clear();
   fHits.resize(fNofTrdLayers);
   
   ArrangeHits();

   CreateTrdFromSts();

   ExtrapolateToTrd();
 
   
   fNofIter = 4;
   
   for (Int_t iIter = 0; iIter < fNofIter; iIter++) {
       
      if (fVerbose > 1) {
         cout << "-I- CbmLitTrdTrackFinderSts::DoFind: "
              << "iteration " << iIter << endl;
      }
   
      //fHits.resize(fNofTrdLayers);
      //fHitIdMap.clear();
      
      if (iIter == 0) SetIterPar(3, 0, fNofTrdStations - 1, 0);
      if (iIter == 1) SetIterPar(5, 0, fNofTrdStations - 2, 0);
      if (iIter == 2) SetIterPar(5, 1, fNofTrdStations - 1, 0);
      if (iIter == 3) SetIterPar(5, 0, fNofTrdStations - 1, 1);
    
            
      if (iIter != 0) ArrangeHits();
      
      InitTracks();

      TrackFollowing();
     
      RemoveHits();

      CopyToOutput(trackArray);
   
    //  for(UInt_t i = 0; i < fTrdTracks.size(); i++)
    //       delete fTrdTracks[i];
    //  fTrdTracks.clear();

      for (Int_t i = 0; i < fNofTrdLayers; i++) 
      fHits[i].clear();
      //fHits.clear();
   }

   fUsedHitsSet.clear();
   fHits.clear();
   
   for(std::vector<CbmTrdTrack*>::iterator iTrack = fTrdTracksS.begin();
      iTrack != fTrdTracksS.end(); iTrack++) {
      delete (*iTrack);
   }
   fTrdTracksS.clear();

   cout << "-I- CbmLitTrdTrackFinderSts::DoFind : "
        << trackArray->GetEntriesFast() << " TRD tracks found." << endl;

   fEvents++;

   cout << "-I- CbmLitTrdTrackFinderSts : "
        << fEvents << " events processed" << endl;
        
   cout << "--------------------------------------------------------" << endl;

   return trackArray->GetEntriesFast();
}
// -----------------------------------------------------------------------

// ------------Set iteration parameters-----------------------------------
void CbmLitTrdTrackFinderSts::SetIterPar(
                               Double_t SigmaCoef, 
                               Int_t beginStation, 
                               Int_t endStation,
                               Int_t maxNofMissingHitsInStation) 
{
   fSigmaCoef = SigmaCoef;
   fBeginStation = beginStation;
   fEndStation= endStation;
   fMaxNofMissingHitsInStation = maxNofMissingHitsInStation;
   
}
// -----------------------------------------------------------------------

// --------------Creates Trd tracks from Sts tracks-----------------------
void CbmLitTrdTrackFinderSts::CreateTrdFromSts()
{

    Int_t nStsTracks = fArrayStsTrack->GetEntriesFast();

    for(Int_t iStsTrack = 0; iStsTrack < nStsTracks; iStsTrack++) {
       
       CbmStsTrack* StsTrack = (CbmStsTrack*) fArrayStsTrack->At(iStsTrack);
   
       if(StsTrack == NULL) continue;
            
       CbmTrdTrack* TrdTrack = new CbmTrdTrack();
       TrdTrack->SetParamLast(*StsTrack->GetParamLast());
       TrdTrack->SetChi2(StsTrack->GetChi2());
       TrdTrack->SetNDF(1);
       TrdTrack->SetStsTrackIndex(iStsTrack);
       fTrdTracksS.push_back(TrdTrack);
       
       TrdTrack = NULL; ///?!?!?
    }

    if (fVerbose > 1) {
      cout << "-I- CbmLitTrdTrackFinderSts::CreateTrdFromSts : " << endl
            << nStsTracks << " STS tracks were loaded, "
            << fTrdTracksS.size() << " TRD tracks were created" << endl;
    }
}
// -----------------------------------------------------------------------

// --------------------extrapolate STS tracks to first TRD layer----------
void CbmLitTrdTrackFinderSts::ExtrapolateToTrd()
{
   Double_t Ze = fHits[0][0]->GetZ();
   
   for (std::vector<CbmTrdTrack*>::iterator iTrack = fTrdTracksS.begin();
            iTrack != fTrdTracksS.end(); iTrack++) {

      fFieldPropagator->Propagate((*iTrack)->GetParamLast(), Ze);
    //  fKf->SetMethod(1);
    //  fKf->Propagate((*iTrack)->GetParamLast(), Ze);
    //  (*iTrack)->GetParamLast()->Print();
      (*iTrack)->SetParamFirst(*(*iTrack)->GetParamLast());
      (*iTrack)->SetChi2( 0.0 );
   }

   if (fVerbose > 1) {
      cout << "-I- CbmLitTrdTrackFinderSts::ExtrapolateToTrd() : "
           << "Extrapolation to TRD finished " << endl;
   }
}
// -----------------------------------------------------------------------

// -----------------------------------------------------------------------
void CbmLitTrdTrackFinderSts::InitTracks() {

   for (std::vector<CbmTrdTrack*>::iterator iTrack = fTrdTracksS.begin();
           iTrack != fTrdTracksS.end(); iTrack++) {
      
      if (fStsIdSet.find((*iTrack)->GetStsTrackIndex()) 
                    != fStsIdSet.end()) continue;
            
      CbmTrdTrack* track = new CbmTrdTrack(*(*iTrack));
      
      fTrdTracks.push_back(track);
      
      track = NULL;
   }
   
   // extrapolate to the begin layer for the tracking
   Int_t beginLayer = 0;
   if (fBeginStation > 0) {    
   
      for(Int_t i = 0; i < fBeginStation; i++) 
         beginLayer += fNofTrdLayersPerStation[i];   
   
      Double_t Ze = fHits[beginLayer][0]->GetZ();
      
      for (std::vector<CbmTrdTrack*>::iterator iTrack = fTrdTracks.begin();
           iTrack != fTrdTracks.end(); iTrack++) {
    
         fLinePropagator->Propagate((*iTrack)->GetParamLast(), Ze);
      
      } 
   }
}
// -----------------------------------------------------------------------

// --------------------Copy to output TClonesArray------------------------
void CbmLitTrdTrackFinderSts::CopyToOutput(TClonesArray* trackArray)
{
   Int_t NofTracks = trackArray->GetEntriesFast();
   for(std::vector<CbmTrdTrack*>::iterator iTrack = fFoundTracks.begin();
         iTrack != fFoundTracks.end(); iTrack++) {
      if( (*iTrack)->GetFlag() == 0) {
         if (fStsIdSet.find((*iTrack)->GetStsTrackIndex()) == fStsIdSet.end()) {
         
            new ((*trackArray)[NofTracks]) CbmTrdTrack(*(*iTrack));
                                    
            fStsIdSet.insert((*iTrack)->GetStsTrackIndex());
            
            NofTracks++;
         }
      }
      delete (*iTrack);
   }
   
   //clear array
   fFoundTracks.clear();
}
// -----------------------------------------------------------------------

ClassImp(CbmLitTrdTrackFinderSts);