/** CbmLitTrackFitterWeight.cxx
 *@author A.Lebedev <alebedev@jinr.ru>
 *@since 2008
 **/

#include "fitter/CbmLitTrackFitterWeight.h"

#include "base/CbmLitEnums.h"
#include "data/CbmLitHit.h"
#include "data/CbmLitStripHit.h"
#include "data/CbmLitPixelHit.h"
#include "data/CbmLitTrack.h"
#include "data/CbmLitTrackParam.h"
#include "interface/CbmLitWeightCalculator.h"
#include "utils/CbmLitMemoryManagment.h"
#include "utils/CbmLitComparators.h"
#include "weight/CbmLitWeightCalculatorSimple.h"
#include "weight/CbmLitWeightCalculatorGauss.h"
#include "weight/CbmLitWeightCalculatorTukey.h"
#include "weight/CbmLitWeightedHitCalculatorImp.h"

#include <vector>
#include <cmath>
#include <iostream>
#include <algorithm>

CbmLitTrackFitterWeight::CbmLitTrackFitterWeight(
   TrackFitterPtr fitter,
   TrackFitterPtr smoother)
{
   fFitter = fitter;
   fSmoother = smoother;
   fWeightedHitCalculator = WeightedHitCalculatorPtr(new CbmLitWeightedHitCalculatorImp());
   fSimpleWeightCalculator = WeightCalculatorPtr(new CbmLitWeightCalculatorSimple());
// fGaussWeightCalculator = WeightCalculatorPtr(new CbmLitWeightCalculatorGauss());
   fTukeyWeightCalculator = WeightCalculatorPtr(new CbmLitWeightCalculatorTukey());

   //track fit parameters
   fNofIterations = 3;
   fAnnealing.push_back(0.);
   fAnnealing.push_back(30.);
   fAnnealing.push_back(15.);
   fOutlierCut = 1e-20;
}

CbmLitTrackFitterWeight::~CbmLitTrackFitterWeight()
{
}

LitStatus CbmLitTrackFitterWeight::Initialize()
{
   return kLITSUCCESS;
}

LitStatus CbmLitTrackFitterWeight::Finalize()
{
   return kLITSUCCESS;
}

LitStatus CbmLitTrackFitterWeight::Fit(
   CbmLitTrack* track,
   bool downstream)
{
   if (track->GetNofHits() < 1) { return kLITERROR; }
   track->SortHits();

   CbmLitTrack etrack;
   etrack.SetPDG(track->GetPDG());
// etrack.SetParamFirst(track->GetParamFirst());
   for(int iter = 0; iter < fNofIterations; iter++) {
      etrack.SetParamFirst(track->GetParamFirst());
      if (CreateEffectiveTrack(track, iter, &etrack) == kLITERROR) {
         return kLITERROR;
      }
      if (CheckEffectiveTrack(&etrack) == kLITERROR) {
         return kLITERROR;
      }
      if (fFitter->Fit(&etrack) == kLITERROR) {
         return kLITERROR;
      }
      if (fSmoother->Fit(&etrack) == kLITERROR) {
         return kLITERROR;
      }
   }
   CreateOutputTrack(track);

   return kLITSUCCESS;
}

LitStatus CbmLitTrackFitterWeight::CreateEffectiveTrack(
   CbmLitTrack* track,
   int iter,
   CbmLitTrack* etrack) const
{
   std::vector<HitPtrIteratorPair> bounds;
   track->GetHitBounds(bounds);
   etrack->ClearHits();
   for (unsigned int i = 0; i < bounds.size(); i++) {
      int nofHits = bounds[i].second - bounds[i].first;
      if (nofHits > 1) {
         //Create effective hit, from set of hits
         //Add effective hit to track
         CbmLitHit* ehit;
         if ((*bounds[i].first)->GetType() == kLITPIXELHIT) { ehit = new CbmLitPixelHit; }
         else if ((*bounds[i].first)->GetType() == kLITSTRIPHIT) { ehit = new CbmLitStripHit; }
         const CbmLitTrackParam* par = (iter > 0)? etrack->GetFitNode(i)->GetSmoothedParam() : NULL;
         LitStatus result = CreateEffectiveHit(bounds[i].first, bounds[i].second, par, iter, ehit);
         if (result == kLITSUCCESS) {
            etrack->AddHit(ehit);
         } else {
            // hit not added
         }
         delete ehit;
      } else {
         //if only one hit, add it directly to the track
         etrack->AddHit(*bounds[i].first);
      }
   }
   return kLITSUCCESS;
}

LitStatus CbmLitTrackFitterWeight::CreateEffectiveHit(
   HitPtrIterator itBegin,
   HitPtrIterator itEnd,
   const CbmLitTrackParam* par,
   int iter,
   CbmLitHit* hit) const
{
   //Calculate hit weights
   //const CbmLitTrackParam* par;// = (iter > 0)? etrack->GetFitNode(i)->GetSmoothedParam() : NULL;
// do {
   if (CalculateWeights(par, itBegin, itEnd, iter) == kLITERROR) {
      return kLITERROR;
   }
   MarkOutliers(itBegin, itEnd);
   if (AreAllOutliers(itBegin, itEnd)) { return kLITERROR; }
   // repeat weights recalculation until none of the weights will be marked
// } while (MarkOutliers(itBegin, itEnd));

   fWeightedHitCalculator->DoCalculate(itBegin, itEnd, hit);

   return kLITSUCCESS;
}

bool CbmLitTrackFitterWeight::AreAllOutliers(
   HitPtrIterator itBegin,
   HitPtrIterator itEnd) const
{
   for(HitPtrIterator it = itBegin; it != itEnd; it++) {
      if (!(*it)->IsOutlier()) {
         return false;
      }
   }
   return true;
}

bool CbmLitTrackFitterWeight::MarkOutliers(
   HitPtrIterator itBegin,
   HitPtrIterator itEnd) const
{
   bool result = false;
   for(HitPtrIterator it = itBegin; it != itEnd; it++) {
      if ((*it)->GetW() < fOutlierCut) {
         (*it)->IsOutlier(true);
         result = true;
      }
   }
   return result;
}

LitStatus CbmLitTrackFitterWeight::CalculateWeights(
   const CbmLitTrackParam* par,
   HitPtrIterator itBegin,
   HitPtrIterator itEnd,
   int iter) const
{
   if (iter == 0) {
      fSimpleWeightCalculator->DoCalculate(par, itBegin, itEnd, 0);
   } else {
      myf T = fAnnealing[iter];
      fTukeyWeightCalculator->DoCalculate(par, itBegin, itEnd, T);
   }
   return kLITSUCCESS;
}


LitStatus CbmLitTrackFitterWeight::CheckEffectiveTrack(
   const CbmLitTrack* track) const
{
   if (track->GetNofHits() < 8) { return kLITERROR; }
   return kLITSUCCESS;
}

LitStatus CbmLitTrackFitterWeight::CreateOutputTrack(
   CbmLitTrack* track)
{
   //find hit with the best weight for each station
   //other hits will be removed
   std::vector<HitPtrIteratorPair> bounds;
   track->GetHitBounds(bounds);
// track->ClearHits();

   HitPtrVector newHits;
   for (unsigned int i = 0; i < bounds.size(); i++) {
      std::sort(bounds[i].first, bounds[i].second, CompareHitPtrWLess());
      CbmLitHit* hit = (*(bounds[i].second-1));
      if (!hit->IsOutlier()) {
         if (hit->GetType() == kLITSTRIPHIT) {
            CbmLitStripHit* newHit = new CbmLitStripHit(*static_cast<const CbmLitStripHit*>(hit));
            newHits.push_back(newHit);
         } else if (hit->GetType() == kLITPIXELHIT) {
            CbmLitPixelHit* newHit = new CbmLitPixelHit(*static_cast<const CbmLitPixelHit*>(hit));
            newHits.push_back(newHit);
         }
      }
   }
   track->ClearHits();
   for (unsigned int i = 0; i < newHits.size(); i++) {
      track->AddHit(newHits[i]);
   }
   for_each(newHits.begin(), newHits.end(), DeleteObject());

   if (track->GetNofHits() == 0) { return kLITERROR; }

   if (fFitter->Fit(track) == kLITERROR) { return kLITERROR; }
   if (fSmoother->Fit(track) == kLITERROR) { return kLITERROR; }

   track->SetLastPlaneId(track->GetHit(track->GetNofHits()-1)->GetPlaneId());

   return kLITSUCCESS;
}