/*
 *=====================================================
 *
 *  CBM Level 1 Reconstruction
 *
 *  Authors: M.Zyzak
 *
 *  e-mail :
 *
 *=====================================================
 *
 *  SIMD Fitter for CbmL1Track class
 *
 */

#include "CbmKFTrackFitter.h"

//L1Algo tools
#include "CbmL1Track.h"

#include "L1AlgoInputData.h"
#include "L1Algo.h"
#include "L1Extrapolation.h"
#include "L1AddMaterial.h"

#include <iostream>

using std::vector;
 
CbmKFTrackFitter::CbmKFTrackFitter():fHits(0),fAlgo(0)
{
}

CbmKFTrackFitter::~CbmKFTrackFitter()
{
}

void CbmKFTrackFitter::GetChiToVertex(vector<CbmL1Track> &Tracks, vector<L1FieldRegion> &field, vector<float> &chiToVtx, float chiPrim)
{
  fvec ONE(1.f);
  chiToVtx.reserve(Tracks.size());

  int nTracks_SIMD = fvecLen;
  L1TrackPar T; // fitting parametr coresponding to current track
  
  CbmL1Track *t[fvecLen];

  int nStations = fAlgo->NStations;
  int NMvdStations = fAlgo->NMvdStations;
  const L1Station *sta = fAlgo->vStations;
  fvec x;
  fvec u;
  fvec v;
  fvec y;
  fvec x_temp,y_temp;
  fvec* zSta = new fvec[nStations];
  for(int iSta=0; iSta<nStations; iSta++)
    zSta[iSta] = sta[iSta].z;

  field.reserve(Tracks.size());

  L1FieldRegion fld _fvecalignment;
  L1FieldValue fB[3], fB_temp _fvecalignment;
  fvec zField[3];

  unsigned short N_vTracks = Tracks.size();
  int ista;
  for(unsigned short itrack = 0; itrack < N_vTracks; itrack+=fvecLen)
  {
    if(N_vTracks - itrack < static_cast<unsigned short>(fvecLen))
      nTracks_SIMD = N_vTracks - itrack;

    fvec mass2;
    for(int iVec=0; iVec<nTracks_SIMD; iVec++)
    {
      t[iVec] = & Tracks[itrack+iVec]; // current track
      T.x[iVec]  = t[iVec]->GetTrack()[0];
      T.y[iVec]  = t[iVec]->GetTrack()[1];
      T.tx[iVec] = t[iVec]->GetTrack()[2];
      T.ty[iVec] = t[iVec]->GetTrack()[3];
      T.qp[iVec] = t[iVec]->GetTrack()[4];
      T.z[iVec]  = t[iVec]->GetTrack()[5];
      T.C00[iVec] = t[iVec]->GetCovMatrix()[0];
      T.C10[iVec] = t[iVec]->GetCovMatrix()[1];
      T.C11[iVec] = t[iVec]->GetCovMatrix()[2];
      T.C20[iVec] = t[iVec]->GetCovMatrix()[3];
      T.C21[iVec] = t[iVec]->GetCovMatrix()[4];
      T.C22[iVec] = t[iVec]->GetCovMatrix()[5];
      T.C30[iVec] = t[iVec]->GetCovMatrix()[6];
      T.C31[iVec] = t[iVec]->GetCovMatrix()[7];
      T.C32[iVec] = t[iVec]->GetCovMatrix()[8];
      T.C33[iVec] = t[iVec]->GetCovMatrix()[9];
      T.C40[iVec] = t[iVec]->GetCovMatrix()[10];
      T.C41[iVec] = t[iVec]->GetCovMatrix()[11];
      T.C42[iVec] = t[iVec]->GetCovMatrix()[12];
      T.C43[iVec] = t[iVec]->GetCovMatrix()[13];
      T.C44[iVec] = t[iVec]->GetCovMatrix()[14];
//      float mass = TDatabasePDG::Instance()->GetParticle(t[iVec]->GetPidHypo())->Mass();
      const float mass = t[iVec]->GetMass();//KFParticleDatabase::Instance()->GetMass(t[iVec]->GetPidHypo());
      mass2[iVec] = mass*mass;

      for(int iH = 0; iH < 2; iH++ ) //it is assumed, that track has at least 3 hits
      {
        const L1StsHit &hit = fHits->GetStsHits()[t[iVec]->StsHits[iH]];
        ista = fHits->GetSFlag()[hit.f]/4;

        u[iVec]  = fHits->GetStsStrips()[hit.f].f;
        v[iVec]  = fHits->GetStsStripsB()[hit.b].f;
        fAlgo->StripsToCoor(u[iVec], v[iVec], x_temp, y_temp, sta[ista]);
        x[iVec]  = x_temp[iVec];
        y[iVec]  = y_temp[iVec];
                
        sta[ista].fieldSlice.GetFieldValue( x, y, fB_temp );
        fB[iH+1].x[iVec] = fB_temp.x[iVec];
        fB[iH+1].y[iVec] = fB_temp.y[iVec];
        fB[iH+1].z[iVec] = fB_temp.z[iVec];
        zField[iH+1][iVec] = fHits->GetStsZPos()[hit.iz];
      }
    }
    
    fB[0] = fAlgo->GetVtxFieldValue();
    zField[0] = 0;
    fld.Set( fB[2], zField[2], fB[1], zField[1], fB[0], zField[0] );
    field.push_back(fld);
        
    for(int iSt= nStations-4; iSt>=0; iSt--)
    {
      fvec w(1.f);
      fvec initialized = fvec(T.z > (zSta[iSt]+2.5));
      w = fvec(w & initialized);

      L1Extrapolate( T, zSta[iSt], T.qp, fld, &w );
      if(iSt == NMvdStations - 1)
      {
        L1AddPipeMaterial( T, T.qp, w, mass2);
        EnergyLossCorrection( T, mass2, PipeRadThick, T.qp, fvec(1.f), w);
      }
      L1AddMaterial( T, fAlgo->fRadThick[iSt].GetRadThick(T.x, T.y), T.qp, w, mass2);
      EnergyLossCorrection( T, mass2, fAlgo->fRadThick[iSt].GetRadThick(T.x, T.y), T.qp, fvec(1.f), w);
    }
    if( NMvdStations <= 0 )
    {
      L1Extrapolate( T, 22.5, T.qp, fld);
      L1AddPipeMaterial( T, T.qp, ONE, mass2);
      EnergyLossCorrection( T, mass2, PipeRadThick, T.qp, fvec(1.f), ONE);
    }
//     L1Extrapolate( T, primVtx.GetRefZ(), T.qp, fld);
// 
//     Double_t Cv[3] = { primVtx.GetCovMatrix()[0], primVtx.GetCovMatrix()[1], primVtx.GetCovMatrix()[2] };
// 
//     fvec dx = T.x - primVtx.GetRefX();
//     fvec dy = T.y - primVtx.GetRefY();

    L1Extrapolate( T, 0, T.qp, fld);
    L1AddTargetMaterial( T, T.qp, ONE, mass2);
    EnergyLossCorrection( T, mass2, TargetRadThick, T.qp, fvec(1.f), ONE);

    fvec Cv[3] = { 1, 1, 0.0125*0.0125 };

    fvec dx = T.x - 0;
    fvec dy = T.y - 0;

    fvec c[3] = { T.C00, T.C10, T.C11 };
    c[0]+= Cv[0];  c[1]+= Cv[1];  c[2]+= Cv[2];
    fvec d = c[0]*c[2] - c[1]*c[1] ;
    fvec chi = sqrt( fabs( 0.5*(dx*dx*c[0]-2*dx*dy*c[1]+dy*dy*c[2])/d ) );
    fvec isNull = fvec(fabs(d)<1.e-20);
    chi = fvec(fvec(!isNull) & chi) + fvec(isNull & fvec(0));

    for(int iVec=0; iVec<nTracks_SIMD; iVec++)
      chiToVtx.push_back(chi[iVec]);
    
    if(chiPrim>0)
    {
      for(int iVec=0; iVec<nTracks_SIMD; iVec++)
      {
        if( chi[iVec] < chiPrim )
        {
          t[iVec]->GetTrack()[0] = T.x[iVec];
          t[iVec]->GetTrack()[1] = T.y[iVec];
          t[iVec]->GetTrack()[2] = T.tx[iVec];
          t[iVec]->GetTrack()[3] = T.ty[iVec];
          t[iVec]->GetTrack()[4] = T.qp[iVec];
          t[iVec]->GetTrack()[5] = T.z[iVec];
          t[iVec]->GetCovMatrix()[0] = T.C00[iVec];
          t[iVec]->GetCovMatrix()[1] = T.C10[iVec];
          t[iVec]->GetCovMatrix()[2] = T.C11[iVec];
          t[iVec]->GetCovMatrix()[3] = T.C20[iVec];
          t[iVec]->GetCovMatrix()[4] = T.C21[iVec];
          t[iVec]->GetCovMatrix()[5] = T.C22[iVec];
          t[iVec]->GetCovMatrix()[6] = T.C30[iVec];
          t[iVec]->GetCovMatrix()[7] = T.C31[iVec];
          t[iVec]->GetCovMatrix()[8] = T.C32[iVec];
          t[iVec]->GetCovMatrix()[9] = T.C33[iVec];
          t[iVec]->GetCovMatrix()[10] = T.C40[iVec];
          t[iVec]->GetCovMatrix()[11] = T.C41[iVec];
          t[iVec]->GetCovMatrix()[12] = T.C42[iVec];
          t[iVec]->GetCovMatrix()[13] = T.C43[iVec];
          t[iVec]->GetCovMatrix()[14] = T.C44[iVec];
        }
      }
    }
  }
  delete[] zSta;
}

void CbmKFTrackFitter::CalculateFieldRegion(const vector<CbmL1Track> &Tracks, vector<L1FieldRegion> &field)
{
  L1FieldRegion fld _fvecalignment;

  int nTracks_SIMD = fvecLen;
  L1TrackPar T; // fitting parametr coresponding to current track

  const CbmL1Track *t[fvecLen];

  int ista;
  const L1Station *sta = fAlgo->vStations;
  fvec x;
  fvec u;
  fvec v;
  fvec y;
  fvec x_temp,y_temp;
  L1FieldValue fB[3], fB_temp _fvecalignment;
  fvec zField[3];

  unsigned short N_vTracks = Tracks.size();

  for(unsigned short itrack = 0; itrack < N_vTracks; itrack+=fvecLen)
  {
    if(N_vTracks - itrack < static_cast<unsigned short>(fvecLen))
      nTracks_SIMD = N_vTracks - itrack;

    for(int i=0; i<nTracks_SIMD; i++)
      t[i] = & Tracks[itrack+i]; // current track

    for(int iVec=0; iVec<nTracks_SIMD; iVec++)
    {
      for(int iH = 0; iH < 2; iH++ ) //it is assumed, that track has at least 3 hits
      {
        const L1StsHit &hit = fHits->GetStsHits()[t[iVec]->StsHits[iH]];
        ista = fHits->GetSFlag()[hit.f]/4;

        u[iVec]  = fHits->GetStsStrips()[hit.f].f;
        v[iVec]  = fHits->GetStsStripsB()[hit.b].f;
        fAlgo->StripsToCoor(u[iVec], v[iVec], x_temp, y_temp, sta[ista]);
        x[iVec]  = x_temp[iVec];
        y[iVec]  = y_temp[iVec];
        sta[ista].fieldSlice.GetFieldValue( x, y, fB_temp );
        fB[iH+1].x[iVec] = fB_temp.x[iVec];
        fB[iH+1].y[iVec] = fB_temp.y[iVec];
        fB[iH+1].z[iVec] = fB_temp.z[iVec];
        zField[iH+1][iVec] = fHits->GetStsZPos()[hit.iz];
      }
    }

    fB[0] = fAlgo->GetVtxFieldValue();
    zField[0] = 0;
    fld.Set( fB[2], zField[2], fB[1], zField[1], fB[0], zField[0] );
    // fld.Set( fB[1], zField[1], fB[0], zField[0] );
    field.push_back(fld);
  }
}