// CylTrackNtuple.cpp

#include "CylTrackNtuple.h"
#include "trfutil/TupleManager.h"
#include "trfcyl/SurfCylinder.h"
#include "trfbase/ETrack.h"
#include "trffit/HTrack.h"
#include "trffit/TrackMatch.h"
#include "trffit/RTrack.h"
#include "trfutil/TRFMath.h"

using trf::VTrack;
using trf::HTrack;
using trf::McRecoTrackNtuple;
using trf::SurfCylinder;
using trf::CylTrackNtuple;

// Assign track parameter indices.
enum {
  IPHI = SurfCylinder::IPHI,
  IZ   = SurfCylinder::IZ,
  IALF = SurfCylinder::IALF,
  ITLM = SurfCylinder::ITLM,
  IQPT  = SurfCylinder::IQPT
};

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

// constructor

CylTrackNtuple::CylTrackNtuple(bool mc, bool reco, bool mc2)
: McRecoTrackNtuple(mc,reco,mc2) {

  // We do not yet support final MC tracks.

  // define column names
  add_column_int("stat");
  add_column_int("evt");
  add_column_float("mr");
  add_column_float("mphi");
  add_column_float("mz");
  add_column_float("malf");
  add_column_float("mtlm");
  add_column_float("mqpt");
  add_column_float("rr");
  add_column_float("rphi");
  add_column_float("rz");
  add_column_float("ralf");
  add_column_float("rtlm");
  add_column_float("rqpt");
  add_column_float("epp");
  add_column_float("ezz");
  add_column_float("eaa");
  add_column_float("ell");
  add_column_float("eqq");
  add_column_int("nhit");
  add_column_int("nmeas");
  add_column_float("chsq");
  add_column_float("mchsq");

  if ( mc2 ) {
    add_column_float("malf2");
    add_column_float("mz2");
    add_column_float("mphid2");
    add_column_float("mtlm2");
    add_column_float("mqpt2");
    add_column_float("mchsq2");
  }

  // surface for checking type
  _psrf = new SurfCylinder(1.0);

}

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

// destructor

CylTrackNtuple::~CylTrackNtuple() { }

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

// Fill MC part of an ntuple.

void CylTrackNtuple::_fill_mc(const VTrack& trv, double mphi) {

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

  fill_float( "mr", mcsrf.get_parameter(0) );
  fill_float( "mphi", mphi );
  fill_float( "mz", trv.get_vector()(IZ) );
  fill_float( "malf", trv.get_vector()(IALF) );
  fill_float( "mtlm", trv.get_vector()(ITLM) );
  fill_float( "mqpt", trv.get_vector()(IQPT) );

}

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

// Fill the second MC part of an ntuple.

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

  assert( mc2() );

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

  fill_float( "malf2", trv.get_vector()(IALF) );
  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 CylTrackNtuple::_fill_reco(const HTrack& trh, double rphi) {

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

  fill_float( "rr", fitsrf.get_parameter(0) );
  fill_float( "rphi", rphi );
  fill_float( "rz", trh.get_track().get_vector()(IZ) );
  fill_float( "ralf", trh.get_track().get_vector()(IALF) );
  fill_float( "rtlm", trh.get_track().get_vector()(ITLM) );
  fill_float( "rqpt", trh.get_track().get_vector()(IQPT) );
  fill_float( "epp", trh.get_track().get_error()(IPHI,IPHI) );
  fill_float( "ezz", trh.get_track().get_error()(IZ,IZ) );
  fill_float( "eaa", trh.get_track().get_error()(IALF,IALF) );
  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 CylTrackNtuple::
fill(const VTrack& trv, const HTrack& trh, const VTrack& trv2) {
//  assert(false);

  // 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()(IPHI);
  mphid = fmod1( mphid, TWOPI );

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

  // Shift mc2 phid to be close to MC phid.
  double mphid2 = trv2.get_vector()(IPHI);
  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 CylTrackNtuple::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 phi to range (0,PI).
  double mphi = trv.get_vector()(IPHI);
  mphi = fmod1( mphi, TWOPI );

  // Shift reco phi to be close to MC phi.
  double rphi = trh.get_track().get_vector()(IPHI);
  double dphi = fmod2( rphi-mphi, TWOPI );
  rphi = mphi + dphi;

  // 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(mphi));
  _fill_reco(trh,float(rphi));
  fill_float( "mchsq", float(match) );
  if ( auto_store() ) store_row();

}

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

// fill ntuple from an MC-MC pair

void CylTrackNtuple::
fill(const VTrack& trv, const VTrack& trv2) {
//  assert(false);

  // 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()(IPHI);
  mphid = fmod1( mphid, TWOPI );

  // Shift mc2 phid to be close to MC phid.
  double mphid2 = trv2.get_vector()(IPHI);
  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 CylTrackNtuple::fill(const VTrack& trv) {

  // Shift phi to range (0,PI).
  double mphi = trv.get_vector()(0);
  mphi = fmod1( mphi, TWOPI );

  fill_int( "stat", 1 );
  fill_int( "evt",  _evt );
  _fill_mc(trv,float(mphi));
  if ( auto_store() ) store_row();

}

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

// Fill ntuple from an unmatched reco track.

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

  // Shift phi to range (0,PI).
  double rphi = trh.get_track().get_vector()(0);
  rphi = fmod1( rphi, TWOPI );

  fill_int( "stat", 1 );
  fill_int( "evt",  _evt );
  _fill_reco(trh,float(rphi));
  if ( auto_store() ) store_row();

}

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