// *****************************************************************************
// This file is property of and copyright by the CBM and the ALICE HLT Project *
// All rights reserved.                                                        *
//                                                                             *
// Primary Authors: Ivan Kisel <kisel@kip.uni-heidelberg.de>                   *
//                  Igor Kulakov <I.Kulakov@gsi.de>                            *
//                  Maksym Zyzak <M.Zyzak@gsi.de>                              *
//                  GSI, Uni Frankfurt and HICforFAIR.                         *
//                                                                             *
// 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.                       *
//                                                                             *
//******************************************************************************

#ifdef DO_TPCCATRACKER_EFF_PERFORMANCE


#include "FitPerformance.h"

#include "TMath.h"
#include "TROOT.h"
#include "Riostream.h"
#include "TFile.h"
#include "TH1.h"
#include "TH2.h"
#include "TProfile.h"
#include "TStyle.h"
#include "AliHLTTPCCAFitter.h"

using arbb::dense;
using arbb::pack;
using arbb::scatter;
using arbb::gather;
using arbb::fill;
using arbb::range;
using arbb::value;

using arbb::boolean;
using arbb::f32;


FitPerformance::FitPerformance()
{
  string pullParNames[NParameters] = {"Y","Z","SinPhi","DzDs","QPt"};
  string resParNames[NParameters] = {"Y","Z","SinPhi","DzDs","Pt"};

  float MinXRes [NParameters] = {-0.25f,-0.7f,-0.03f,-0.02f,-0.3f};
  float MaxXRes [NParameters];
    for(int i=0; i<NParameters; i++) MaxXRes[i] = -MinXRes[i];
  int   NBinsRes[NParameters] = {30.f,30.f,30.f,30.f,30.f};

  float MinXPull [NParameters] = {-7.f,-7.f,-7.f,-7.f,-7.f};
  float MaxXPull [NParameters];
    for(int i=0; i<NParameters; i++) MaxXPull[i] = -MinXPull[i];
  int   NBinsPull[NParameters] = {30.f,30.f,30.f,30.f,30.f};

  string NameTmp;
  for(int i=0; i<NParameters; i++)
  {
    NameTmp = "res"; NameTmp += resParNames[i];
    fResHistos[i] = TH1F(NameTmp.data(), NameTmp.data(), NBinsRes[i], MinXRes[i], MaxXRes[i]);

    NameTmp = "pull"; NameTmp += pullParNames[i];
    fPullHistos[i] = TH1F(NameTmp.data(), NameTmp.data(), NBinsPull[i], MinXPull[i], MaxXPull[i]);
  }
}

FitPerformance::~FitPerformance()
{
}

void FitPerformance::ReadData(fstream &InputTraksMC, const int &NTracks)
{
  mcX  = fill(0.f,NTracks);
  mcY  = fill(0.f,NTracks);
  mcZ  = fill(0.f,NTracks);
  mcPX = fill(0.f,NTracks);
  mcPY = fill(0.f,NTracks);
  mcPZ = fill(0.f,NTracks);
  mcQP = fill(0.f,NTracks);

  range<f32> mcXRange  = mcX.write_only_range();
  range<f32> mcYRange  = mcY.write_only_range();
  range<f32> mcZRange  = mcZ.write_only_range();
  range<f32> mcPXRange = mcPX.write_only_range();
  range<f32> mcPYRange = mcPY.write_only_range();
  range<f32> mcPZRange = mcPZ.write_only_range();
  range<f32> mcQPRange = mcQP.write_only_range();

  string tmpstring;
  int TrackCur = 0;
  for(int iTr=0; iTr<NTracks; iTr++)
  {
    InputTraksMC >> tmpstring >> TrackCur;
    if(TrackCur != iTr)
    {
      std::cout << "Reading MC Tracks, Error: file with data is corrupted." << std::endl;
      return;
    }
    InputTraksMC >> tmpstring >> mcXRange[iTr];
    InputTraksMC >> tmpstring >> mcYRange[iTr];
    InputTraksMC >> tmpstring >> mcZRange[iTr];
    InputTraksMC >> tmpstring >> mcQPRange[iTr];
    InputTraksMC >> tmpstring >> mcPXRange[iTr];
    InputTraksMC >> tmpstring >> mcPYRange[iTr];
    InputTraksMC >> tmpstring >> mcPZRange[iTr];
  }
}

void FitPerformance::FitQualityCheck(AliHLTTPCCAFitter &fitter)
{
  AliHLTTPCCATrackParamArBB InnerParam;
  fitter.GetTrackInnerParam(InnerParam);

  dense<boolean> FittedMask = fill(0, mcX.length());
  dense<i32> TrackID = fitter.GetTrackID();
  FittedMask = scatter(fill<boolean>(1,InnerParam.ArraySize), TrackID, FittedMask);

  mcX = pack(mcX,FittedMask);
  mcY = pack(mcY,FittedMask);
  mcZ = pack(mcZ,FittedMask);
  mcPX = pack(mcPX,FittedMask);
  mcPY = pack(mcPY,FittedMask);
  mcPZ = pack(mcPZ,FittedMask);
  mcQP = pack(mcQP,FittedMask);

  dense<f32> mcEx = mcPX*mcQP;
  dense<f32> mcEy = mcPY*mcQP;
  dense<f32> mcEz = mcPZ*mcQP;
  dense<f32> mcEt = arbb::sqrt( mcEx * mcEx + mcEy * mcEy );

  dense<f32> mcSinPhi = mcEy / mcEt;
  dense<f32> mcDzDs   = mcEz / mcEt;
  dense<f32> mcQPt    = mcQP / mcEt;

  dense<boolean> TransportedMask = fill(1,InnerParam.ArraySize);
  InnerParam.TransportToXWithMaterial( TransportedMask, mcX, fitter.cBz(), f32(0.999f) );

  dense<boolean> BackDirTracks = InnerParam.SignCosPhi()*mcEx < 0;

  mcSinPhi = select(BackDirTracks, -mcSinPhi, mcSinPhi);
  mcDzDs   = select(BackDirTracks, -mcDzDs,   mcDzDs);
  mcQPt    = select(BackDirTracks, -mcQPt,    mcQPt);

  dense<f32> mcPt = fill(1.f,mcQPt.length()) / arbb::abs( mcQPt );

  dense<f32> qPt = InnerParam.QPt();
  dense<f32> pt  = fill(1.f,InnerParam.ArraySize) / arbb::abs( qPt );

  dense<f32> dY      = InnerParam.Y()      - mcY;
  dense<f32> dZ      = InnerParam.Z()      - mcZ;
  dense<f32> dSinPhi = InnerParam.SinPhi() - mcSinPhi;
  dense<f32> dDzDs   = InnerParam.DzDs()   - mcDzDs;
  dense<f32> dQPt    = qPt                 - mcQPt;
  dense<f32> dPt     = pt                  - mcPt;
  dPt = dPt/mcPt;

  dense<f32> pY      = dY/arbb::sqrt(InnerParam.Err2Y());
  dense<f32> pZ      = dZ/arbb::sqrt(InnerParam.Err2Z());
  dense<f32> pSinPhi = dSinPhi/arbb::sqrt(InnerParam.Err2SinPhi());
  dense<f32> pDzDs   = dDzDs/arbb::sqrt(InnerParam.Err2DzDs());
  dense<f32> pQPt    = dQPt/arbb::sqrt(InnerParam.Err2QPt());

  for(int i=0; i<value(InnerParam.ArraySize); i++)
  {
    if(!(value(TransportedMask[i])) ) continue;

    fResHistos[0].Fill( value(dY[i]) );
    fResHistos[1].Fill( value(dZ[i]) );
    fResHistos[2].Fill( value(dSinPhi[i]) );
    fResHistos[3].Fill( value(dDzDs[i]) );
    if(fabs(value(qPt[i])) > 1.e-8) fResHistos[4].Fill( value(dPt[i]) );

    if ( value(InnerParam.Err2Y()[i]) > 0 )      fPullHistos[0].Fill( value( pY[i] ) );
    if ( value(InnerParam.Err2Z()[i]) > 0 )      fPullHistos[1].Fill( value( pZ[i] ) );
    if ( value(InnerParam.Err2SinPhi()[i]) > 0 ) fPullHistos[2].Fill( value( pSinPhi[i] ) );
    if ( value(InnerParam.Err2DzDs()[i]) > 0 )   fPullHistos[3].Fill( value( pDzDs[i] ) );
    if ( fabs(value(qPt[i])) > 1.e-7 && 
         value(InnerParam.Err2QPt()[i])>0 )      fPullHistos[4].Fill( value( pQPt[i] ) );
  }
}

FitPerformance &FitPerformance::Instance()
{
  // reference to static object
  static FitPerformance gFitPerformance;
  return gFitPerformance;
}

void FitPerformance::WriteDir2Current( TObject *obj )
{
  //* recursive function to write down all created in the file histos
  if ( !obj->IsFolder() ) obj->Write();
  else {
    TDirectory *cur = gDirectory;
    ((TDirectory*)obj)->cd();
    TList *listSub = ( ( TDirectory* )obj )->GetList();
    TIter it( listSub );
    while ( TObject *obj1 = it() ) WriteDir2Current( obj1 );
    cur->cd();
  }
}

void FitPerformance::WriteHistos()
{
  if(fOutputFile) WriteDir2Current(fOutputFile);
}

#endif //DO_TPCCATRACKER_EFF_PERFORMANCE