/** CbmLitFindMvdTracks.cxx
 * @author Andrey Lebedev <andrey.lebedev@gsi.de>
 * @since 2011
 * @version 1.0
 **/
#include "finder/CbmLitFindMvdTracks.h"
#include "base/CbmLitEnvironment.h"
#include "base/CbmLitToolFactory.h"
#include "data/CbmLitHit.h"
#include "data/CbmLitPixelHit.h"
#include "data/CbmLitStripHit.h"
#include "data/CbmLitTrack.h"
#include "data/CbmLitTrackParam.h"
#include "finder/CbmLitTrackFinderNN.h"
#include "utils/CbmLitConverter.h"
#include "utils/CbmLitMemoryManagment.h"

#include "CbmStsTrack.h"
#include "FairHit.h"
#include "FairRunAna.h"
#include "FairRuntimeDb.h"

#include <iostream>
#include <algorithm>

CbmLitFindMvdTracks::CbmLitFindMvdTracks()
{

}

CbmLitFindMvdTracks::~CbmLitFindMvdTracks()
{

}

InitStatus CbmLitFindMvdTracks::Init()
{
   DetermineSetup();
   ReadAndCreateDataBranches();

   CbmLitEnvironment* env = CbmLitEnvironment::Instance();
   env->GetMvdLayout();

   InitTrackReconstruction();

   return kSUCCESS;
}

void CbmLitFindMvdTracks::Exec(
   Option_t* opt)
{
   ConvertInputData();

   RunTrackReconstruction();

   ConvertOutputData();

   ClearArrays();

   std::cout << "-I- Event: " << fEventNo++ << std::endl;
}

void CbmLitFindMvdTracks::SetParContainers()
{
   FairRunAna* ana = FairRunAna::Instance();
   FairRuntimeDb* rtdb = ana->GetRuntimeDb();
   rtdb->getContainer("FairBaseParSet");
}

void CbmLitFindMvdTracks::Finish()
{

}

void CbmLitFindMvdTracks::DetermineSetup()
{
   CbmLitEnvironment* env = CbmLitEnvironment::Instance();
   fIsMvd = env->IsMvd();
   fIsSts = env->IsSts();

   std::cout << "-I- Found detectors:";
   if (fIsMvd) { std::cout << "MVD "; }
   if (fIsSts) { std::cout << "STS "; }
   std::cout << std::endl;
}

void CbmLitFindMvdTracks::ReadAndCreateDataBranches()
{
   FairRootManager* ioman = FairRootManager::Instance();
   if (NULL == ioman) { Fatal("Init","CbmRootManager is not instantiated"); }

   // MVD data
   if (fIsMvd) {
      fMvdHits = (TClonesArray*) ioman->GetObject("MvdHit");
      if (NULL == fMvdHits) { Fatal("Init", "No MvdHit array!"); }
   }

   //STS data
   if (fIsSts) {
      fStsTracks = (TClonesArray*) ioman->GetObject("StsTrack");
      if (NULL == fStsTracks) { Fatal("Init","No StsTrack array!"); }
      fStsHits = (TClonesArray*) ioman->GetObject("StsHit");
      if (NULL == fStsHits) { Fatal("Init","No StsHit array!"); }
   }
}

void CbmLitFindMvdTracks::ConvertInputData()
{
   CbmLitConverter::StsTrackArrayToTrackVector(fStsTracks, fLitStsTracks);
   // Change last and first parameters of the track seeds
   for(int iTrack = 0; iTrack < fLitStsTracks.size(); iTrack++) {
      CbmLitTrack* track = fLitStsTracks[iTrack];
      const CbmLitTrackParam* parLast = track->GetParamLast();
      const CbmLitTrackParam* parFirst = track->GetParamFirst();
      track->SetParamLast(parFirst);
      track->SetParamFirst(parLast);
   }
   std::cout << "-I- Number of STS tracks: " << fLitStsTracks.size() << std::endl;

   CbmLitConverter::MvdHitArrayToHitVector(fMvdHits, fLitMvdHits);
   // Make reverse order of the hits
   CbmLitEnvironment* env = CbmLitEnvironment::Instance();
   int nofPlanes = env->GetMvdLayout().GetNofPlanes();
   for(int iHit = 0; iHit < fLitMvdHits.size(); iHit++) {
      CbmLitHit* hit = fLitMvdHits[iHit];
      hit->SetPlaneId(nofPlanes - hit->GetPlaneId() - 1);
   }
   std::cout << "-I- Number of MVD hits: " << fLitMvdHits.size() << std::endl;
}

void CbmLitFindMvdTracks::ConvertOutputData()
{
   for(int iTrack = 0; iTrack < fLitOutputTracks.size(); iTrack++) {
      CbmLitTrack* litTrack = fLitOutputTracks[iTrack];
      int trackId = litTrack->GetPreviousTrackId();
      CbmStsTrack* track = static_cast<CbmStsTrack*>(fStsTracks->At(trackId));
      for (int iHit = 0; iHit < litTrack->GetNofHits(); iHit++) {
         const CbmLitHit* litHit = litTrack->GetHit(iHit);
         int refId = litHit->GetRefId();
         FairHit* hit = static_cast<FairHit*>(fMvdHits->At(refId));
         track->AddMvdHit(refId, hit);
      }
      track->SortMvdHits();

      FairTrackParam parFirst;
      CbmLitConverter::LitTrackParamToTrackParam(litTrack->GetParamLast(), &parFirst);
      track->SetParamFirst(parFirst);
   }
}

void CbmLitFindMvdTracks::InitTrackReconstruction()
{
   CbmLitToolFactory* factory = CbmLitToolFactory::Instance();
   fFinder = factory->CreateTrackFinder("mvd_nn");//"mvd_" + fTrackingType);
}

void CbmLitFindMvdTracks::ClearArrays()
{
   // Free memory
   for_each(fLitStsTracks.begin(), fLitStsTracks.end(), DeleteObject());
   for_each(fLitMvdHits.begin(), fLitMvdHits.end(), DeleteObject());
   for_each(fLitOutputTracks.begin(), fLitOutputTracks.end(), DeleteObject());
   fLitStsTracks.clear();
   fLitMvdHits.clear();
   fLitOutputTracks.clear();
}

void CbmLitFindMvdTracks::RunTrackReconstruction()
{
   fFinder->DoFind(fLitMvdHits, fLitStsTracks, fLitOutputTracks);
}