// DCATrackNtuple.cpp

#include "DCATrackNtuple.h"
#include "trfdca/SurfDCA.h"
#include "trffit/HTrack.h"
#include "trfbase/ETrack.h"
#include "trfutil/TupleManager.h"
#include "trfutil/TRFMath.h"

using namespace trf;

//**********************************************************************

namespace {

// Assign track parameter indices.
enum {
  IRSIGNED = SurfDCA::IRSIGNED,
  IZ       = SurfDCA::IZ,
  IPHID    = SurfDCA::IPHID,
  ITLM     = SurfDCA::ITLM,
  IQPT     = SurfDCA::IQPT
};

}  // end unnamed namespoace

//**********************************************************************

// constructor

DCATrackNtuple::DCATrackNtuple(bool mc_, bool reco_, bool mc2_)
: McRecoTrackNtuple(mc_,reco_,mc2_) {

  // define columns
  add_column_int("evt");
  if ( mc() && reco() ) {
    add_column_int("stat");
  }
  if ( mc() ) {
    add_column_float("mrs");
    add_column_float("mz");
    add_column_float("mphid");
    add_column_float("mtlm");
    add_column_float("mqpt");
  }
  if ( reco() ) {
    add_column_float("rrs");
    add_column_float("rz");
    add_column_float("rphid");
    add_column_float("rtlm");
    add_column_float("rqpt");
    add_column_float("err");
    add_column_float("ezz");
    add_column_float("epp");
    add_column_float("ell");
    add_column_float("eqq");
    add_column_int("nhit");
    add_column_int("nmeas");
    add_column_float("chsq");
  }
  if ( mc() && reco() ) {
    add_column_float("mchsq");
  }
  if ( mc2() ) {
    add_column_float("mrs2");
    add_column_float("mz2");
    add_column_float("mphid2");
    add_column_float("mtlm2");
    add_column_float("mqpt2");
    if ( reco() ) {
      add_column_float("mchsq2");
    }
  }

  // surface for checking type
  _psrf = new SurfDCA();

}

//**********************************************************************

// destructor

DCATrackNtuple::~DCATrackNtuple() { }

//**********************************************************************

// Fill MC part of an ntuple.

void DCATrackNtuple::_fill_mc(const VTrack& trv, double mphid) {

  assert( mc() );

  // Check the MC track is at DCA.
  const Surface& mcsrf = *trv.get_surface();
  assert( _psrf->get_type() == mcsrf.get_type() );

  fill_float( "mrs", trv.get_vector()(IRSIGNED) );
  fill_float( "mz", trv.get_vector()(IZ) );
  fill_float( "mphid", mphid );
  fill_float( "mtlm",  trv.get_vector()(ITLM) );
  fill_float( "mqpt",  trv.get_vector()(IQPT) );

}

//**********************************************************************

// Fill the second MC part of an ntuple.

void DCATrackNtuple::_fill_mc2(const VTrack& trv, double mphid) {

  assert( mc2() );

  // Check the MC track is at DCA.
  const Surface& mcsrf = *trv.get_surface();
  assert( _psrf->get_type() == mcsrf.get_type() );

  fill_float( "mrs2", trv.get_vector()(IRSIGNED) );
  fill_float( "mz2",  trv.get_vector()(IZ) );
  fill_float( "mphid2", mphid );
  fill_float( "mtlm2", trv.get_vector()(ITLM) );
  fill_float( "mqpt2", trv.get_vector()(IQPT) );

}

//**********************************************************************

// Fill reco part of an ntuple.
// Phi is provided separately so caller can set the phi-range.

void DCATrackNtuple::_fill_reco(const HTrack& trh, double rphid) {

  assert( reco() );

  // Check the track is at DCA.
  const Surface& fitsrf = *trh.get_track().get_surface();
  assert( _psrf->get_type() == fitsrf.get_type() );

  fill_float( "rrs",   float(trh.get_track().get_vector()(IRSIGNED)) );
  fill_float( "rz",    float(trh.get_track().get_vector()(IZ))   );
  fill_float( "rphid", float(rphid)                              );
  fill_float( "rtlm",  float(trh.get_track().get_vector()(ITLM)) );
  fill_float( "rqpt",  float(trh.get_track().get_vector()(IQPT)) );

  fill_float( "err", trh.get_track().get_error()(IRSIGNED,IRSIGNED) );
  fill_float( "ezz", trh.get_track().get_error()(IZ,IZ) );
  fill_float( "epp", trh.get_track().get_error()(IPHID,IPHID) );
  fill_float( "ell", trh.get_track().get_error()(ITLM,ITLM) );
  fill_float( "eqq", trh.get_track().get_error()(IQPT,IQPT) );
  fill_int( "nhit", trh.get_hits().size() );
  int nmeas = 0;
  const HitList& hits = trh.get_hits();
  HitList::const_iterator it;
  for (it = hits.begin(); it!=hits.end();++it){
    nmeas += (*it)->size();
  }
  fill_int( "nmeas", nmeas );
  fill_float( "chsq", trh.get_chi_square() );

}

//**********************************************************************

// fill ntuple from triplet

void DCATrackNtuple::
fill(const VTrack& trv, const HTrack& trh, const VTrack& trv2) {

  // Check the Monte Carlo and fit tracks are at the same surface.
  const Surface& mcsrf = *trv.get_surface();
  const Surface& mcsrf2 = *trv.get_surface();
  const Surface& fitsrf = *trh.get_track().get_surface();
  assert( mcsrf.pure_equal(fitsrf) );
  assert( mcsrf.pure_equal(mcsrf2) );

  // Shift phid to range (0,2*PI).
  double mphid = trv.get_vector()(IPHID);
  mphid = fmod1( mphid, TWOPI );

  // Shift reco phid to be close to MC phid.
  double rphid = trh.get_track().get_vector()(IPHID);
  double dphid = fmod2( rphid-mphid, TWOPI );
  rphid = mphid + dphid;

  // Shift mc2 phid to be close to MC phid.
  double mphid2 = trv2.get_vector()(IPHID);
  double dphid2 = fmod2( mphid2-mphid, TWOPI );
  mphid2 = mphid + dphid2;

  // Evaluate the match chi-square.
  double match = chisq_diff( trh.get_track(), trv );
  double match2 = chisq_diff( trh.get_track(), trv2 );

  fill_int( "stat", 0 );
  fill_int( "evt", _evt );
  _fill_mc(trv,float(mphid));
  _fill_reco(trh,float(rphid));
  _fill_mc2(trv2,float(mphid2));
  fill_float( "mchsq", float(match) );
  fill_float( "mchsq2", float(match2) );
  if ( auto_store() ) store_row();

}

//**********************************************************************

// fill ntuple from an MC-reco pair

void DCATrackNtuple::fill(const VTrack& trv, const HTrack& trh) {

  // Check the Monte Carlo and fit tracks are at the same surface.
  const Surface& mcsrf = *trv.get_surface();
  const Surface& fitsrf = *trh.get_track().get_surface();
  assert( mcsrf.pure_equal(fitsrf) );

  // Shift phid to range (0,2*PI).
  double mphid = trv.get_vector()(IPHID);
  mphid = fmod1( mphid, TWOPI );

  // Shift reco phid to be close to MC phid.
  double rphid = trh.get_track().get_vector()(IPHID);
  double dphid = fmod2( rphid-mphid, TWOPI );
  rphid = mphid + dphid;

  // Evaluate the match chi-square.
  double match = chisq_diff( trh.get_track(), trv );

  fill_int( "stat", 0 );
  fill_int( "evt", _evt );
  _fill_mc(trv,float(mphid));
  _fill_reco(trh,float(rphid));
  fill_float( "mchsq", float(match) );
  if ( auto_store() ) store_row();

}

//**********************************************************************

// fill ntuple from an MC-MC pair

void DCATrackNtuple::
fill(const VTrack& trv, const VTrack& trv2) {

  // Check the Monte Carlo tracks are at the same surface.
  const Surface& mcsrf = *trv.get_surface();
  const Surface& mcsrf2 = *trv.get_surface();
  assert( mcsrf.pure_equal(mcsrf2) );

  // Shift phid to range (0,2*PI).
  double mphid = trv.get_vector()(IPHID);
  mphid = fmod1( mphid, TWOPI );

  // Shift mc2 phid to be close to MC phid.
  double mphid2 = trv2.get_vector()(IPHID);
  double dphid2 = fmod2( mphid2-mphid, TWOPI );
  mphid2 = mphid + dphid2;

  fill_int( "evt", _evt );
  if ( reco() ) {
    fill_int( "stat", 1 );
  }
  _fill_mc(trv,float(mphid));
  _fill_mc2(trv2,float(mphid2));
  if ( auto_store() ) store_row();

}

//**********************************************************************

// Fill ntuple from an unmatched MC track.

void DCATrackNtuple::fill(const VTrack& trv) {

  // Shift phid to range (0,PI).
  double mphid = trv.get_vector()(IPHID);
  mphid = fmod1( mphid, TWOPI );

  fill_int( "evt", _evt );
  if ( reco() ) {
    fill_int( "stat", 1 );
  }
  _fill_mc(trv,float(mphid));
  if ( auto_store() ) store_row();

}

//**********************************************************************

// Fill ntuple from an unmatched reco track.

void DCATrackNtuple::fill(const HTrack& trh) {

  // Shift phid to range (0,PI).
  double rphid = trh.get_track().get_vector()(IPHID);
  rphid = fmod1( rphid, TWOPI );

  fill_int( "evt", _evt );
  if ( mc() ) {
    fill_int( "stat", 2 );
  }
  _fill_reco(trh,float(rphid));
  if ( auto_store() ) store_row();

}

//**********************************************************************