#include "CbmStsKFTrackFitter.h"

#include "CbmRootManager.h"
#include "CbmStsTrack.h"
#include "CbmVertex.h"

#include "CbmKFMath.h"
#include "CbmKFTrack.h"
#include "CbmKFVertex.h"
#include "CbmKFStsHit.h"

ClassImp(CbmStsKFTrackFitter);

CbmStsKFTrackFitter::CbmStsKFTrackFitter(){ 
  fHitsArray = 0;
  fIsInitialised = 0;
}

void CbmStsKFTrackFitter::Init(){ 
  CbmRootManager *fManger = CbmRootManager::Instance();
  fHitsArray   = (TClonesArray *)  fManger->GetObject("STSHit");
  fIsInitialised = 1;
};

void CbmStsKFTrackFitter::SetKFHits(CbmKFTrack &T, CbmStsTrack* track){

  T.fHits.clear();

  if( !fIsInitialised ) Init();
  
  Int_t NHits   = ( fHitsArray   ) ?track->GetNHits()   :0;
  if ( NHits <= 0 ) return ;
  
  fHits.resize(NHits);
  for (Int_t i=0; i<NHits;i++){
    Int_t j = track->GetHitIndex(i);
    fHits[i].Create((CbmStsHit*) fHitsArray->At(j));
    T.fHits.push_back(&(fHits[i]));
  }
}

Int_t CbmStsKFTrackFitter::DoFit( CbmStsTrack* track, Int_t pidHypo ){

  track->SetPidHypo(pidHypo);

  CbmKFTrack T;
  T.SetPID(pidHypo);
  SetKFHits( T, track);
  for( Int_t i=0; i<6; i++ ) T.GetTrack()[i] = 0.; // no guess taken
  T.Fit( 1 ); // fit downstream 
  T.Fit( 0 ); // fit upstream  
  T.Fit( 1 ); // fit downstream 
  T.GetTrackParam( *track->GetParamLast() ); // store fitted track & cov.matrix
  T.Fit( 0 ); // fit upstream
  T.GetStsTrack( *track, 1 );          // store fitted track & cov.matrix & chi2 & NDF

  return 0;
}


void CbmStsKFTrackFitter::Extrapolate( CbmTrackParam* track, Double_t z, CbmTrackParam* e_track )
{
  if( !track ) return;
  CbmKFTrack T;  
  T.SetTrackParam( *track );
  T.Extrapolate( z );
  if( e_track ) T.GetTrackParam( *e_track );
}


void CbmStsKFTrackFitter::Extrapolate( CbmStsTrack* track, Double_t z, CbmTrackParam* e_track )
{
  if( !track ) return;
  CbmKFTrack T;
  T.SetPID( track->GetPidHypo() );
  CbmTrackParam *fpar = track->GetParamFirst(), *lpar = track->GetParamLast();
  
  if( z<=fpar->GetZ() ){ // extrapolate first parameters
    T.SetTrackParam( *fpar );
    T.Extrapolate( z );
  }else   if( z<fpar->GetZ()+0.1 ){ // extrapolate first parameters
    T.SetTrackParam( *fpar );
    T.Propagate( z );
  }else if( lpar->GetZ()<=z ){ // extrapolate last parameters
    T.SetTrackParam( *lpar );
    T.Extrapolate( z );
  }else if( lpar->GetZ()-0.1<z ){ // extrapolate last parameters
    T.SetTrackParam( *lpar );
    T.Propagate( z );
  }else { // refit with smoother
    SetKFHits( T, track);
    T.SetTrackParam( *fpar );
    T.Smooth( z );
  }
  if( e_track ) T.GetTrackParam( *e_track );
}

 
Double_t  CbmStsKFTrackFitter::GetChiToVertex( CbmStsTrack* track,  CbmVertex *vtx )
{
  if( !vtx ){  
    CbmRootManager *fManger = CbmRootManager::Instance();
    vtx = (CbmVertex *)  fManger->GetObject("PrimaryVertex");
    if( !vtx ){
      cout<< "-W- CbmStsKFTrackFitter::GetChiToVertex: No Primary Vertex found!"<<endl;
      return 100.;
    }
  }
  CbmKFTrack T;
  T.SetStsTrack( *track, 1 );    
  T.Extrapolate( vtx->GetZ() );

  TMatrixFSym tmp(3);
  vtx->CovMatrix( tmp );
  Double_t Cv[3] = { tmp(0,0), tmp(0,1), tmp(1,1) };

  return CbmKFMath::getDeviation( T.GetTrack()[0], T.GetTrack()[1], T.GetCovMatrix(), 
				  vtx->GetX(), vtx->GetY(), Cv ); 
}


Double_t CbmStsKFTrackFitter::FitToVertex( CbmStsTrack* track, CbmVertex *vtx, CbmTrackParam *v_track )
{
  Double_t ret = 100.;  
  if( !track || !vtx || !v_track ) return ret;
  CbmKFTrack T( *track );
  CbmKFVertex V( *vtx );
  T.Fit2Vertex( V );
  T.GetTrackParam( *v_track );
  if( T.GetRefNDF()>0 && T.GetRefChi2()>=0 ){
    ret = T.GetRefChi2()/T.GetRefNDF();
    if( finite(ret) ) ret =  sqrt(ret);
  }
  return ret;
}