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

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



// ------------------------- 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;
   }
   
   //Create Kalman Filter
   fKf = new CbmLitKalman();
   if(NULL == fKf) {
      cout << "-E- CbmLitTrdTrackFinderSts::Init : "
           << "kalman filter is not initialized." << endl;
   }
   fKf->Init();

   // Determine the TRD geometry
   ReadTrdGeometry();

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


// -------------------- Algorithm implementation -------------------------
Int_t CbmLitTrdTrackFinderSts::DoFind(TClonesArray* hitArray,
                                      TClonesArray* trackArray)
{

   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(1);
      if (iIter == 1) SetIterPar(3);
      if (iIter == 2) SetIterPar(5);
      if (iIter == 3) SetIterPar(7);
         
            
      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();

   for(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) 
{
   fSigmaCoef = SigmaCoef;

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

// --------------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->SetStsTrackIndex(iStsTrack);
       fTrdTracksS.push_back(TrdTrack);
    }

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

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

      fKf->SetMethod(3);
      fKf->Propagate((*iTrack)->GetParamLast(), Ze);

      (*iTrack)->SetParamFirst(*(*iTrack)->GetParamLast());
      (*iTrack)->SetChi2( 0.0 );
   }

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

// -----------------------------------------------------------------------
void CbmLitTrdTrackFinderSts::InitTracks() {
   for (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);
   }
}
// -----------------------------------------------------------------------

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

ClassImp(CbmLitTrdTrackFinderSts);