// -------------------------------------------------------------------------
// -----                  CbmLitKalman source file               -----
// -----                  Created 14/08/06  by A. Lebedev               -----
// -------------------------------------------------------------------------
 


#include "CbmGeoTrdPar.h"
#include "CbmTrackParam.h"


#include "CbmLitKalman.h"


//constructors
CbmLitKalman::CbmLitKalman() {

}

//Destructor
CbmLitKalman::~CbmLitKalman() {
   ////
}


void CbmLitKalman::Init()
{

   ReadMaterials();
}

void CbmLitKalman::Propagate(CbmTrackParam *pParam, double z_out) 
{

   CbmLitState State;
 
   CopyFromTrackParam(pParam, State);
   
   State.SetZ(pParam->GetZ());
     
   KalmanPrediction(State, z_out);

   pParam->SetZ(z_out);
   CopyToTrackParam(pParam, State);
}

void CbmLitKalman::Propagate(CbmTrdTrack *pTrack, double z_out, Bool_t Downstream) 
{

   CbmLitState State;

   if (Downstream) {
      CopyFromTrackParam(pTrack->GetParamLast(), State);
   
      State.SetZ(pTrack->GetParamLast()->GetZ());
   
      KalmanPrediction(State, z_out);
   
      pTrack->GetParamLast()->SetZ(z_out);
      CopyToTrackParam(pTrack->GetParamLast(), State);
   } else {
      CopyFromTrackParam(pTrack->GetParamFirst(), State);
   
      State.SetZ(pTrack->GetParamFirst()->GetZ());
   
      KalmanPrediction(State, z_out);
   
      pTrack->GetParamFirst()->SetZ(z_out);
      CopyToTrackParam(pTrack->GetParamFirst(), State);   
   }

}


void CbmLitKalman::Filter(CbmTrdTrack *pTrack, CbmTrdHit *pHit, Bool_t Downstream){
 
   CbmLitNode Node;
   
   if (Downstream)
       CopyFromTrackParam(pTrack->GetParamLast(), Node.GetPredictedState());
   else 
       CopyFromTrackParam(pTrack->GetParamFirst(), Node.GetPredictedState());

   Node.GetMeas().SetX(pHit->GetX());
   Node.GetMeas().SetY(pHit->GetY());

   Double_t cov[4]; 
   cov[0] = pHit->GetDx() * pHit->GetDx() * 1e-8;
   cov[3] = pHit->GetDy() * pHit->GetDy() * 1e-8;
   
   Node.GetMeas().SetCovMatrix(cov);

   KalmanFilter(Node);
   
   if (Downstream)
      CopyToTrackParam(pTrack->GetParamLast(), Node.GetFilteredState());
   else 
      CopyToTrackParam(pTrack->GetParamFirst(), Node.GetFilteredState());

   pTrack->SetChi2(pTrack->GetChi2() + Node.GetFilteredChi2());

}


void CbmLitKalman::Filter(CbmStsTrack *pTrack, CbmHit *pHit)
{
   
 
   CbmLitNode Node;
   
   CopyFromTrackParam(pTrack->GetParamLast(), Node.GetPredictedState());

   Node.GetMeas().SetX(pHit->GetX());
   Node.GetMeas().SetY(pHit->GetY());

   Double_t cov[4]; 
   cov[0] = pHit->GetDx() * pHit->GetDx();
   cov[3] = pHit->GetDy() * pHit->GetDy();
   
   Node.GetMeas().SetCovMatrix(cov);

   KalmanFilter(Node);
   
   CopyToTrackParam(pTrack->GetParamLast(), Node.GetFilteredState());

   pTrack->SetChi2(pTrack->GetChi2() + Node.GetFilteredChi2());
}


void CbmLitKalman::ReFit(CbmTrdTrack *pTrack, 
                         TClonesArray* TrdHits,
                         Bool_t Downstream)
{
   Int_t NofHits = pTrack->GetNofTrdHits();
   CbmTrdHit *pHit;
   
   if (Downstream) {
      pTrack->SetParamLast(*pTrack->GetParamFirst());
      for (Int_t iHit = 0; iHit < NofHits; iHit++) {
         Int_t HitIndex = pTrack->GetTrdHitIndex(iHit);
         pHit = (CbmTrdHit*) TrdHits->At(HitIndex);
         Double_t z_out = pHit->GetZ();         
         Propagate(pTrack, z_out);
         Filter(pTrack, pHit);   
      }
   } else {
      pTrack->SetParamFirst(*pTrack->GetParamLast());
      for (Int_t iHit = NofHits - 1; iHit > -1; iHit--) {
         Int_t HitIndex = pTrack->GetTrdHitIndex(iHit);
         pHit = (CbmTrdHit*) TrdHits->At(HitIndex);
         Double_t z_out = pHit->GetZ();
         Propagate(pTrack, z_out, kFALSE);
         Filter(pTrack, pHit, kFALSE);   
      }  
   }
}

void CbmLitKalman::CopyToTrackParam(CbmTrackParam *par, CbmLitState &State) 
{
   
   Double_t *X = State.GetStateVector();
   
   par->SetX(X[0]);
   par->SetY(X[1]);
   par->SetTx(X[2]);
   par->SetTy(X[3]);
   par->SetQp(X[4]);

  // par->SetCovMatrix(State.GetCovMatrix());
   Double_t *C = State.GetCovMatrix();
   for (int i = 0; i < 5; i++)
      for (int j = 0; j < 5; j++) par->SetCovariance(i,j,C[5*i + j]);
}

void CbmLitKalman::CopyFromTrackParam(CbmTrackParam *par, CbmLitState &State) 
{

   Double_t *X = State.GetStateVector();
  
   X[0] = par->GetX();
   X[1] = par->GetY();
   X[2] = par->GetTx();
   X[3] = par->GetTy();
   X[4] = par->GetQp();
   
 //  TMatrixFSym cov(5);
 //  par->CovMatrix(cov);
 //  State.SetCovMatrix(cov);
   Double_t *C = State.GetCovMatrix();
   TMatrixFSym cov(5);
   par->CovMatrix(cov);
   for (int i = 0; i < 5; i++)
      for (int j = 0; j < 5; j++) C[5*i + j] = cov[i][j];

}

/*
void CbmLitKalman::Smoother(CbmTrdTrack *pTrack, TClonesArray* TrdHits, bool Downstream)
{
   vector<CbmLitKalmanMeasurement*> M;
   int NHits = pTrack->GetNofTrdHits();
   M.resize(NHits);
   CbmTrdHit* pHit;

   for (int iHit = 0; iHit < NHits; iHit++) {
      int HitIndex = pTrack->GetTrdHitIndex(iHit);
      pHit = (CbmTrdHit*) TrdHits->At(HitIndex);
      M[iHit] = new CbmLitKalmanMeasurement();
      M[iHit]->Create(pHit);
   }

   //if (Downstream) sort(M.begin(),M.end(), CbmLitKalmanMeasurement::CmpDown);
   //else sort(M.begin(),M.end(), CbmLitKalmanMeasurement::CmpDown);

   sort(M.begin(),M.end(), CbmLitKalmanMeasurement::CmpUp);
   //if (Downstream)CopyFromTrackParam(pTrack->GetParamLast(), M[0]->Xp, M[0]->Cp);
   //else
   CopyFromTrackParam(pTrack->GetParamFirst(), M[0]->fXp, M[0]->fCp);

   FitTrackFilter(M);
   //CopyToTrackParam(pTrack->GetParamLast(), M.back()->Xf, M.back()->Cf);

   KalmanSmoother(M, true);
   CopyToTrackParam(pTrack->GetParamFirst(), M.back()->fXs, M.back()->fCs);
/*
   for (int iHit = 0; iHit < NHits; iHit++) {
      for (int i = 0; i < 5; i++) {
         M[iHit]->fXf[i] = M[iHit]->fXs[i];
         for (int j = 0; j < 5; j++) M[iHit]->fCf[i][j] = M[iHit]->fCs[i][j];
      }
   }

   KalmanSmoother(M, false);

   CopyToTrackParam(pTrack->GetParamLast(), M.back()->fXs, M.back()->fCs);
*/


   //if (Downstream)CopyFromTrackParam(pTrack->GetParamFirst(), M.back()->Xs, M.back()->Cs);
   //else CopyFromTrackParam(pTrack->GetParamLast(), M.back()->Xs, M.back()->Cs);

/*   for (int iHit = 0; iHit < NHits; iHit++) delete M[iHit];
   M.clear();
}*/
/*
double CbmLitKalman::GetChi2HitPredicted(CbmTrdTrack* pTrack, CbmTrdHit *pHit)
{
   //calculate chi2 predicted
   //chi2 =(rt)*(R^-1)*(r)
//   t1 = m0 - X_in0;
//   t8 = C_in10 * C_in10;
//   t10 = 0.1e1 / (V00 * V11 + V00 * C_in11 + C_in00 * V11 + C_in00 * C_in11 -
//   t8);
//   t12 = m1 - X_in1;
//   t24 = t1 * (t1 * (V11 + C_in11) * t10 - t12 * C_in10 * t10) + t12 * (-t1 *
//   C_in10 * t10 + t12 * (V00 + C_in00) * t10);

/*   TMatrixFSym C(5);
   CbmTrackParam * par = pTrack->GetParamLast();
   par->CovMatrix(C);

   double Dx = pHit->GetDx();
   double Dy = pHit->GetDy();
   double dx = pHit->GetX() - par->GetX();
   double dy = pHit->GetY() - par->GetY();
   double t10 = 0.1e1 / (Dx * Dy + Dx * C[1][1] + C[0][0] * Dy + C[0][0] * C[1][1] - C[1][0] * C[1][0]);

   double chi2 = t10 * ( dx * dx * (Dy + C[1][1]) + dy * dy * (Dx + C[0][0])
            - 2 * dy * dx * C[1][0] );

   return chi2;
*/
/*   if(pHit->Dx() < pHit->Dy()) {
      double dx = pTrack->GetParamLast()->GetX() - pHit->X();
      double c = pTrack->GetParamLast()->GetCovariance(0, 0)
             + TMath::Power(pHit->Dx() * 1e-4, 2);
      return (dx * dx / c);
   } else {
      double dy = pTrack->GetParamLast()->GetY() - pHit->Y();
      double c = pTrack->GetParamLast()->GetCovariance(1, 1)
             + TMath::Power(pHit->Dy() * 1e-4, 2);
      return (dy*dy / c);
   }
   */
/*   return 0;
}
*/


ClassImp(CbmLitKalman);