// $Id: AliHLTTPCCAFitter.cxx,v 1.9 2010/08/26 15:05:50 ikulakov Exp $
// **************************************************************************
// This file is property of and copyright by the ALICE HLT Project          *
// ALICE Experiment at CERN, All rights reserved.                           *
//                                                                          *
// Primary Authors: Sergey Gorbunov <sergey.gorbunov@kip.uni-heidelberg.de> *
//                  Ivan Kisel <kisel@kip.uni-heidelberg.de>                *
//                  for The ALICE HLT Project.                              *
//                                                                          *
// Developed by:   Igor Kulakov <I.Kulakov@gsi.de>                          *
//                 Maksym Zyzak <M.Zyzak@gsi.de>                            *
//                                                                          *
// Permission to use, copy, modify and distribute this software and its     *
// documentation strictly for non-commercial purposes is hereby granted     *
// without fee, provided that the above copyright notice appears in all     *
// copies and that both the copyright notice and this permission notice     *
// appear in the supporting documentation. The authors make no claims       *
// about the suitability of this software for any purpose. It is            *
// provided "as is" without express or implied warranty.                    *
//                                                                          *
//***************************************************************************

#include "arbb.hpp"
#include <cstdio>
#include <cstring>
#include <cstdlib>

#include "AliHLTTPCCAFitter.h"
#include "AliHLTTPCCATrackParamArBB.h"
#include "TStopwatch.h"
#include "AliHLTTPCCATrackLinearisationArBB.h"

using namespace std;

using arbb::u8;
using arbb::u16;
using arbb::u32;
using arbb::i8;
using arbb::i16;
using arbb::i32;
using arbb::f32;
using arbb::boolean;

using arbb::dense;
using arbb::range;
using arbb::const_range;
using arbb::fill;

AliHLTTPCCAFitter::AliHLTTPCCAFitter()
{
  //* constructor
  fTrInfo.fNSectors = static_cast<int>(fgkNSectors);
}

AliHLTTPCCAFitter::~AliHLTTPCCAFitter()
{
}

void AliHLTTPCCAFitter::ReadData(fstream &InputTraks, const int &NTracks, const int &NClusters)
{
  if(!(InputTraks.is_open()))
  {
    std::cout << "Error: input file is not opened."<<std::endl;
    return;
  }

  std::string tmpstring;

  fTrInfo.Resize(NTracks);
  fClusters.Resize(NClusters);

  range<f32> wClustersXRange = fClusters.fX.write_only_range();
  range<f32> wClustersYRange = fClusters.fY.write_only_range();
  range<f32> wClustersZRange = fClusters.fZ.write_only_range();
  range<u8>  wClustersISectorRange = fClusters.fISector.write_only_range();
  range<u8>  wClustersIRowRange = fClusters.fIRow.write_only_range();
  range<u32> wClustersICluRange = fClusters.fIClu.write_only_range();

  range<f32>     wInnerParamXRange = fTrInfo.fInnerParam.fX.write_only_range();      // x position
  range<f32>     wInnerParamSignRange = fTrInfo.fInnerParam.fSignCosPhi.write_only_range(); // sign of cosPhi   // phi = arctg (Dy/Dx)
  range<f32>     wInnerParamPRange[5];
  range<f32>     wInnerParamCRange[15];
  for(int i=0; i<5; i++)
    wInnerParamPRange[i] = fTrInfo.fInnerParam.fP[i].write_only_range();
  for(int i=0; i<15; i++)
    wInnerParamCRange[i] = fTrInfo.fInnerParam.fC[i].write_only_range();
  range<f32>     wInnerParamChiRange = fTrInfo.fInnerParam.fChi2.write_only_range();
  range<i16>     wInnerParamNDFRange = fTrInfo.fInnerParam.fNDF.write_only_range();

  range<f32> wInnerAlphaRange = fTrInfo.fInnerAlpha.write_only_range();

  range<u16> wNClustersRange  = fTrInfo.fNClusters.write_only_range();
  range<u8>  wSectorRange = fTrInfo.fSector.write_only_range();

  int TrackCur = 0, NClustCur=0, NClustTrack=0, TmpSector=0, tmpInt = 0;

  // Read the tracks data from file
  for(int iTr = 0; iTr<NTracks; iTr++)
  {
      InputTraks >> tmpstring >> TrackCur;
      if(iTr != TrackCur)
      {
        std::cout << "Input Error: track number is corrupted. iTr = "<<iTr<<", TrackCur = "<< TrackCur<< std::endl;
        return;
      }
      InputTraks >> tmpstring >>wInnerParamXRange[iTr]>>tmpstring;
      for(int iP=0; iP<5; iP++)
        InputTraks >> wInnerParamPRange[iP][iTr];
      InputTraks  >>tmpstring>> wInnerParamSignRange[iTr] >> tmpstring;
      for(int iC=0; iC<15; iC++)
        InputTraks >> wInnerParamCRange[iC][iTr];
      InputTraks >>tmpstring>> wInnerParamChiRange[iTr];
      InputTraks >> tmpstring >> tmpInt; wInnerParamNDFRange[iTr] = tmpInt;

      InputTraks >> tmpstring>> wInnerAlphaRange[iTr];
      InputTraks >> tmpstring>> NClustTrack;
      wNClustersRange[iTr] = NClustTrack;

      InputTraks >> tmpstring>> TmpSector;
      wSectorRange[iTr] = TmpSector;

      for ( int iTrClu = 0; iTrClu < NClustTrack; iTrClu++ ) {
        float tmpFloat = 0.f;
        InputTraks >> tmpstring >> tmpFloat; wClustersXRange[NClustCur+iTrClu] = tmpFloat;
        InputTraks >> tmpFloat; wClustersYRange[NClustCur+iTrClu] = tmpFloat;
        InputTraks >> tmpFloat; wClustersZRange[NClustCur+iTrClu] = tmpFloat;
        InputTraks >> tmpInt;   wClustersISectorRange[NClustCur+iTrClu] = tmpInt;
        InputTraks >> tmpInt;   wClustersIRowRange[NClustCur+iTrClu] = tmpInt;
        InputTraks >> tmpInt;   wClustersICluRange[NClustCur+iTrClu] = tmpInt;
      }

      NClustCur += NClustTrack;
  }

  // if NTracks is smaller, than the number of tracks in the file, we should repack data with clusters: 
  // the estimated number of clusters is higher this case
  fClusters.fValid = select(arbb::indices(i32(0), fClusters.fValid.length(), i32(1))<NClustCur, 1, 0);
  fClusters.Pack();

  // Find the first cluster of the track in the array of clusters
  fTrInfo.fFirstClusterRef[0] = 0;
  _for(i32 iTr=1, iTr<fTrInfo.fInnerParam.ArraySize , ++iTr)
  {
    fTrInfo.fFirstClusterRef[iTr] = fTrInfo.fFirstClusterRef[iTr-1] + static_cast<i32>(fTrInfo.fNClusters[iTr-1]);
  }_end_for

  // Difne the identification number for each track
  fTrInfo.orig_track_id = indices(i32(0), i32(NTracks), i32(1));
}

void AliHLTTPCCAFitter::FitTrackArBB(boolean Dir, ClusterInfoArBB &Clusters, AliHLTTPCCATrackParamArBB &Param, dense<f32> &Alpha,
                                     TrackInfoArBB &TrInfo, dense<f32> SectorAlpha,
                                     AliHLTTPCCAParamArBB &SectorParam, dense<boolean> &ok, dense<boolean> &ClusterMask, boolean &Refit)
{
  AliHLTTPCCATrackParamArBB::AliHLTTPCCATrackFitParamArBB fitPar;
  AliHLTTPCCATrackParamArBB t = Param;
  AliHLTTPCCATrackLinearisationArBB l( Param ); // Calculate track linearization parameters: Sin(Phi), Cos(Phi), dz/ds, q/Pt

  dense<f32> Alpha0 = Alpha;

  boolean first(1);

  t.CalculateFitParameters( fitPar ); // Calculate parameters for the material correction calculations

  dense<u16> nClustersNew = fill(0,Param.ArraySize);

  dense<boolean> mask;

  ClusterMask = fill(0,Clusters.fX.length());

  _for ( u16 ihit = 0, ihit < max_reduce(TrInfo.fNClusters), ihit++ ) {

  // Determine, which tracks still have unfitted hits
    mask = ihit < TrInfo.fNClusters;
  // Calculate current track hit indices
    dense<i32> jhit;
    _if(Dir) {
      jhit = TrInfo.fFirstClusterRef + static_cast< dense<i32> >(TrInfo.fNClusters) - fill<i32>(1 + ihit, Param.ArraySize);
    }
    _else {
      jhit = TrInfo.fFirstClusterRef + fill<i32>(ihit, Param.ArraySize);
    }_end_if

  // Determine the rotation angle of the coordinate system, in which hit is.
    const dense<f32> sliceAlpha = gather(SectorAlpha, arbb::gather(Clusters.fISector,jhit));
    dense<boolean> RotMask = mask && arbb::abs( sliceAlpha - Alpha0 ) > 1.e-4f;
    dense<boolean> UnRotMask = !RotMask;
  // If the track parameters are into the coordinate system of one sector and hit is into the coordinate system of another sector - 
  // rotate the track parameters to the coordinate system of the hit.
    t.Rotate( RotMask, sliceAlpha - Alpha0, l, .999f );
    Alpha0 = select( RotMask, sliceAlpha, Alpha0);
    mask = mask && (UnRotMask || RotMask);
  // Transport track parameters to the hit position.
    dense<f32> x = arbb::gather(Clusters.fX,jhit);
    t.TransportToXWithMaterial( mask, x, l, fitPar, SectorParam.cBz() );
  // If this is the first hit of the current fit procedure - reinitialize the covariance matrix, chi2 and NDF.
    _if ( first ) {
      t.fC[0] = fill(10.f, Param.ArraySize);
      t.fC[1] = fill(0.f, Param.ArraySize);
      t.fC[2] = fill(10.f, Param.ArraySize);
      t.fC[3] = fill(0.f, Param.ArraySize);
      t.fC[4] = fill(0.f, Param.ArraySize);
      t.fC[5] = fill(1.f, Param.ArraySize);
      t.fC[6] = fill(0.f, Param.ArraySize);
      t.fC[7] = fill(0.f, Param.ArraySize);
      t.fC[8] = fill(0.f, Param.ArraySize);
      t.fC[9] = fill(1.f, Param.ArraySize);
      t.fC[10] = fill(0.f, Param.ArraySize);
      t.fC[11] = fill(0.f, Param.ArraySize);
      t.fC[12] = fill(0.f, Param.ArraySize);
      t.fC[13] = fill(0.f, Param.ArraySize);
      t.fC[14] = fill(10.f, Param.ArraySize);
      t.fChi2 = fill(0.f, Param.ArraySize);
      t.fNDF  = fill(-5, Param.ArraySize);
      t.CalculateFitParameters( fitPar );
    }_end_if

    dense<f32> err2Y, err2Z;
  // Calculate the errors of the hit.
    SectorParam.GetClusterErrors2( gather(Clusters.fIRow,jhit), t, err2Y, err2Z );
  // Add hit to the track.
    t.Filter( mask, arbb::gather(Clusters.fY,jhit), arbb::gather(Clusters.fZ,jhit), err2Y, err2Z );

    first = 0;
  // Mark successfully added hits
    ClusterMask = ClusterMask || scatter (mask, jhit, ClusterMask, mask);
    nClustersNew = select(mask, nClustersNew+1, nClustersNew);
  }_end_for

  t.fP[4] = select(arbb::abs(t.fP[4]) < 1.e-8f, 1.e-8f, t.fP[4]);

  const i32 NTracks = TrInfo.fNClusters.length();
  // Determine, whether track is fitted correctly
  ok = fill(1,TrInfo.fNClusters.length());

  for(int iC=0; iC<15; iC++) ok = ok && !(t.fC[iC] == std::numeric_limits<float>::infinity() || t.fC[iC] == std::numeric_limits<float>::quiet_NaN());
  for(int iP=0; iP<5 ; iP++) ok = ok && !(t.fP[iP] == std::numeric_limits<float>::infinity() || t.fP[iP] == std::numeric_limits<float>::quiet_NaN());

  ok = ok && (t.fC[0] > 0) && (t.fC[2] > 0) && (t.fC[5] > 0) && (t.fC[9] > 0) && (t.fC[14] > 0);
  ok = select( arbb::abs(t.SinPhi()) >= 0.999f, 0, ok);
  _if(!Refit) { 
    ok = ok && ( t.Chi2()/(static_cast< dense<f32> >(t.NDF())) < 9.f); 
    ok = ok && ( t.NDF() > 0.f );
    ok = ok && ( static_cast< dense<f32> >(nClustersNew) > (AliHLTTPCCAParameters::MinTrackPurity * static_cast< dense<f32> >(TrInfo.fNClusters)));
  }_end_if
  // Set the sign of the Cos(Phi) according to the track linearisation.
  t.SetSignCosPhi( select( ok && (l.CosPhi() >= 0.f) , 1, -1) );
  // Mark hits of the correctly fitted tracks with 1, incorrectly - with 0.
  ClusterMask = ClusterMask && repeat(ok,TrInfo.fNClusters);
  Clusters.fValid = ClusterMask;
  t.Pack(ok);
  Param = t;
  Alpha = pack(Alpha0, ok);
  Clusters.Pack();
  TrInfo.fNClusters = pack(nClustersNew,ok);
  TrInfo.fFirstClusterRef = pack(TrInfo.fFirstClusterRef, ok);
  TrInfo.fFirstClusterRef[0] = 0;
  _for(i32 iTr=1, iTr<Param.ArraySize , ++iTr)
  {
    TrInfo.fFirstClusterRef[iTr] = TrInfo.fFirstClusterRef[iTr-1] + static_cast<i32>(TrInfo.fNClusters[iTr-1]);
  }_end_for

  TrInfo.fSector   = pack(TrInfo.fSector, ok);
  TrInfo.fInnerRow = pack(TrInfo.fInnerRow, ok);
  TrInfo.fOuterRow = pack(TrInfo.fOuterRow, ok);

#ifdef DO_TPCCATRACKER_EFF_PERFORMANCE
  TrInfo.orig_track_id = pack(TrInfo.orig_track_id, ok);
#endif // DO_TPCCATRACKER_EFF_PERFORMANCE
}

void AliHLTTPCCAFitter::FitArBB(boolean Dir, ClusterInfoArBB &Clusters, TrackInfoArBB &TrInfo, dense<f32> SectorAlpha, 
                                AliHLTTPCCAParamArBB &SectorParam)
{
  dense<boolean> ok; // This mask shows, whether track is fitted or not.
  // We will start from the most inner hit  - so, copy inner parameters to the outer, then outer parameters will be modified.
  // Finally, we will obtained outer parameters in the most outer position of the track.
  TrInfo.fOuterParam = TrInfo.fInnerParam;
  TrInfo.fOuterAlpha = TrInfo.fInnerAlpha;
  dense<boolean> ClusterMask;
//  boolean Refit(0);
  boolean Refit(1);
  FitTrackArBB( Dir, Clusters, TrInfo.fOuterParam, TrInfo.fOuterAlpha, TrInfo, SectorAlpha, SectorParam, ok, ClusterMask, Refit);
  // We will start second fit from the most outer hit  - so, copy outer parameters to the inner, then inner parameters will be modified.
  // Finally, we will obtained inner parameters in the most inner position of the track.
  TrInfo.fInnerParam = TrInfo.fOuterParam;
  TrInfo.fInnerAlpha = TrInfo.fOuterAlpha;
  FitTrackArBB( !Dir, Clusters, TrInfo.fInnerParam, TrInfo.fInnerAlpha, TrInfo, SectorAlpha, SectorParam, ok, ClusterMask, Refit);
  // Repack outer parameters according to the "ok" mask (clean the array from the unfitted tracks)
  TrInfo.fOuterParam.Pack(ok);
  TrInfo.fOuterAlpha = pack(TrInfo.fOuterAlpha, ok);

  TrInfo.fChiPrev = arbb::fill(10000000.f, TrInfo.fInnerParam.ArraySize);
  TrInfo.fChiNext = arbb::fill(10000000.f, TrInfo.fInnerParam.ArraySize);
  TrInfo.fPrevNeighbour = arbb::fill(-1, TrInfo.fInnerParam.ArraySize);
  TrInfo.fNextNeighbour = arbb::fill(-1, TrInfo.fInnerParam.ArraySize);
  TrInfo.fSectorPrevNeighbour = arbb::fill(-1, TrInfo.fInnerParam.ArraySize);
  TrInfo.fSectorNextNeighbour = arbb::fill(-1, TrInfo.fInnerParam.ArraySize);

  TrInfo.fSectorTrackInfoStart = fill(0,TrInfo.fNSectors);
  TrInfo.fSectorNTrackInfos = fill(0,TrInfo.fNSectors);
  // Find the total number of clusters in each sector.
  _for(i32 iTr=0, iTr<TrInfo.fInnerParam.ArraySize, ++iTr)
  {
    TrInfo.fSectorNTrackInfos[TrInfo.fSector[iTr]] = TrInfo.fSectorNTrackInfos[TrInfo.fSector[iTr]] + 1;
  }_end_for
  // Find the first cluster of the sector in the array of clusters.
  _for(i32 iSec=1, iSec< TrInfo.fNSectors, ++iSec)
  {
    TrInfo.fSectorTrackInfoStart[iSec] = TrInfo.fSectorTrackInfoStart[iSec-1] + TrInfo.fSectorNTrackInfos[iSec-1];
  }_end_for
}

void AliHLTTPCCAFitter::FitTrack( bool dir )
{
// dir determines the fit direction:
// dir = 0 -> forward  direction (from the most inner (first) hit to the most outer (last) hit)
// dir = 1 -> backward direction (from the most outer (last) hit to the most inner (first) hit)
  boolean Dir(dir);
  call(FitArBB)(Dir,fClusters,fTrInfo,fSectorAlpha, fSectorParamArBB);
}