// AddFitZPlane_XY2_Z2.cpp

#include "AddFitZPlane_XY2_Z2.h"
#include "trffit/HTrack.h"
#include "trfzp/HitZPlane2.h"
#include "trfxyp/HitXYPlane2.h"
#include "trfutil/TRFMath.h"
#include "trfutil/TupleManager.h"
#include "trfutil/EventID.h"
#include "trfbase/Propagator.h"
#include "trfbase/PropStat.h"
#include "trfutil/trfstream.h"
#include "objstream/ObjData.hpp"
#include "objstream/ObjTable.hpp"
#include "trfutil/RootFindLinear.h"
#include "trfutil/Circle.h"

using std::endl;
using std::cerr;

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

using namespace trfutil;
using namespace trf;

// Assign track parameter indices.

enum {
  IX     = SurfZPlane::IX,
  IY     = SurfZPlane::IY,
  IDXDZ  = SurfZPlane::IDXDZ,
  IDYDZ  = SurfZPlane::IDYDZ,
  IQP    = SurfZPlane::IQP
};

//**********************************************************************
// Static data.
//**********************************************************************

bool AddFitZPlane_XY2_Z2::_ltuple = false;

char* AddFitZPlane_XY2_Z2::_tuple_name = "AddFitZPlane_XY2_Z2";

//**********************************************************************
// Free functions.
//**********************************************************************

namespace {

// Creator.

ObjPtr create(const ObjData& data) {
  assert( data.get_object_type() == "AddFitZPlane_XY2_Z2" );

  ObjDoublePtr bfield;
  ObjData::Type dtype = data.get_type("bfield");
  if(dtype == ObjData::DOUBLE)
    bfield = new ObjDouble(data.get_double("bfield"));
  else if(dtype == ObjData::SHARE_PTR)
    bfield.assign_with_dynamic_cast(data.get_share_pointer("bfield"));
  else {
    cerr << "AddFitZPlane_XY2_Z2: can't get bfield parameter." << endl;
    abort();
  }

  ObjDoublePtr ptmin;
  dtype = data.get_type("ptmin");
  if(dtype == ObjData::DOUBLE)
    ptmin = new ObjDouble(data.get_double("ptmin"));
  else if(dtype == ObjData::SHARE_PTR)
    ptmin.assign_with_dynamic_cast(data.get_share_pointer("ptmin"));
  else {
    cerr << "AddFitZPlane_XY2_Z2: can't get ptmin parameter." << endl;
    abort();
  }
  ObjDoublePtr vx(new ObjDouble(0.));
  if( data.has("vx") ) {
    dtype = data.get_type("vx");
    if(dtype == ObjData::DOUBLE)
      vx = new ObjDouble(data.get_double("vx"));
    else if(dtype == ObjData::SHARE_PTR)
      vx.assign_with_dynamic_cast(data.get_share_pointer("vx"));
    else {
      cerr << "ClusterFilterDCA: can't get vx parameter." << endl;
      abort();
    }
  }

  ObjDoublePtr vy(new ObjDouble(0.));
  if( data.has("vy") ) {
    dtype = data.get_type("vy");
    if(dtype == ObjData::DOUBLE)
      vy = new ObjDouble(data.get_double("vy"));
    else if(dtype == ObjData::SHARE_PTR)
      vy.assign_with_dynamic_cast(data.get_share_pointer("vy"));
    else {
      cerr << "ClusterFilterDCA: can't get vy parameter." << endl;
      abort();
    }
  }

  PropagatorPtr prop;
  data.get_share_pointer("propagator",prop);
  SurfacePtr psrf;
  data.get_share_pointer("surface",psrf);
  AddFitterPtr pfit;
  data.get_share_pointer("fitter",pfit);
  std::vector<double> error;
  data.get_double_list("error",error);
  assert(error.size()==15);
  TrackError err;
  int n=0;
  for(int i=0;i<5;++i)
    for(int j=0;j<=i;++j)
      err(i,j)=error[n++];
  return ObjPtr( new AddFitZPlane_XY2_Z2(bfield,pfit,prop,psrf,err,(*ptmin)(),(*vx)(),(*vy)()) );
}

}

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

// Root finder for asin(c1*x) - asin(c2*x) - dphi .

namespace {
class RootFindCurv : public RootFindLinear {
private:
  double _c1;
  double _c2;
  double _dphi;
public:
  RootFindCurv(double c1, double c2, double dphi)
  : _c1(c1), _c2(c2), _dphi(dphi) { };
  StatusDouble evaluate(double x) const {
    double xfail = 9999.99;
    double arg1 = _c1*x;
    double arg2 = _c2*x;
    if ( abs(arg1) > 1.0 ) return StatusDouble(1,xfail);
    if ( abs(arg2) > 1.0 ) return StatusDouble(2,xfail);
    double val = asin(arg1) - asin(arg2) - _dphi;
    return StatusDouble(0,val);
  };
};
}

//**********************************************************************
// Methods.
//**********************************************************************

// ouput stream

void AddFitZPlane_XY2_Z2::ostr(std::ostream& stream) const {
  stream << begin_object ;
  stream << "Add fitter for one XYPlane(v,z) and one ZPlane(x,y) measurements."
  << new_line;
  stream << "B-field is " << get_bfield() << " Tesla." << new_line;
  stream << "Internal fitter is " << *_fit << new_line;
  stream << "Internal propagator is " << *_prop<<new_line;
  stream << "Starting surface is " << *_srf<<new_line;
  stream << "Starting error matrix is " << _err << new_line;
  stream << "Minimum pt is " << get_ptmin() << new_line;
  stream << end_object;
}

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

// fill tuple and exit
int AddFitZPlane_XY2_Z2::fill_tuple(ReturnStatus status) const {
  if ( tuple_is_active() ) {
    _ptup->capture("return",status);
    _ptup->storeCapturedData();
    _ptup->clearData();
  }
  return status;
}

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

// Define the ntuple.

void AddFitZPlane_XY2_Z2::define_tuple() {
  assert( _ptup == 0 );
  bool defined = TupleManager::is_tuple(_tuple_name);
  _ptup = &TupleManager::tuple(_tuple_name);
  if ( ! defined ) {
    float xdefault = -9999.99;
    _ptup->columnDirect("evid",-1);
    _ptup->columnDirect("qpt",xdefault);
    _ptup->columnDirect("return",-999);
  }
}

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

// constructor

AddFitZPlane_XY2_Z2::AddFitZPlane_XY2_Z2(const ObjDoublePtr& bfield,
const AddFitterPtr& fit, const PropagatorPtr& prop, const SurfacePtr& srf,
                                         const TrackError& err,double ptmin,
					 double xv,double yv)
: _bfield(bfield),_fit(fit),_prop(prop),_srf(srf),_err(err),_qpt_max(1.0/ptmin),_xv(xv),_yv(yv)  {
  if ( _ltuple ) define_tuple();
}

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

// destructor

AddFitZPlane_XY2_Z2::~AddFitZPlane_XY2_Z2() { }

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

// add the hit

int AddFitZPlane_XY2_Z2::add_hit(HTrack& trh, const HitPtr& phit2) const {

  // Add event number to tuple.
  if ( tuple_is_active() ) {
    _ptup->capture("evid",EventID::get_global_event_number());
  }

   // check type
  if ( phit2->get_type() != HitZPlane2::get_static_type() )
    return fill_tuple(WRONG_TYPE);

 // Fetch track vector

  ETrack trv = trh.get_track();

  // check that new hit and track on the same surface
  if ( !phit2->get_surface().pure_equal(*trv.get_surface()) )
    return fill_tuple(DIFFERENT_SURFACES);

  // check hit count
  if ( trh.get_hits().size() != 1 ) return fill_tuple(WRONG_HIT_COUNT);

  // check chi-square
  // assert( trh.get_chi_square() == 0.0 );

  // check the type of the first hit
  const HitXYPlane2* hit1 = dynamic_cast<const HitXYPlane2*>(&*(trh.get_hits().front()));
  const HitZPlane2* hit2 = dynamic_cast<const HitZPlane2*>(&*phit2);
  if ( ! hit1 || ! hit2 )  return fill_tuple(WRONG_TYPE_EXISTING);

  bool zero_bfield = false;
  if( fabs(get_bfield()) < 1e-7 ) 
    zero_bfield = true;

  // Fetch the track vector and error.
  TrackVector vec = trv.get_vector();

  // define ZPlane indices

  int izpos  = SurfZPlane::ZPOS;

  // define XYPlane indices

  int iphi  = SurfXYPlane::NORMPHI;
  int idist = SurfXYPlane::DISTNORM;

  // Fetch u,v,z and plane phi1 for the first hit.
  const Surface& srf1 = hit1->get_surface();
  const ClusXYPlane2& clu1 = hit1->get_full_cluster();
  double phi1 = srf1.get_parameter(iphi);
  double u1   = srf1.get_parameter(idist);
  double v1   = clu1.get_v();
  double z1   = clu1.get_z();
  double r1=u1;

  // Fetch x,y,z for the second (new) hit.
  const Surface& srf2 = hit2->get_surface();
  const ClusZPlane2& clu2 = hit2->get_full_cluster();
  double z2 = srf2.get_parameter(izpos);
  double y2 = clu2.get_y();
  double x2 = clu2.get_x();

 
  // check that z2-z1 doesn't equal to 0
  double dz = z2 - z1;
  if ( dz == 0. ) return fill_tuple(DZ_IS_ZERO);

  // Calculate vector (u1,v1,z1) in coordinates of the ZPlane

  double cos_phi1 = cos(phi1);
  double sin_phi1 = sin(phi1);

  double x1   = u1*cos_phi1 - v1*sin_phi1;
  double y1   = v1*cos_phi1 + u1*sin_phi1;

  // Calculate track parameters.
  // We assume track is a straight line
  // from (x1,y1,z1) to (x2,y2,z2)

  double dx = x2 - x1;
  double dy = y2 - y1;
  double dxdz2 = dx/dz;
  double dydz2 = dy/dz;

  if ( dx == 0. && dy == 0. ) return fill_tuple(DU_IS_ZERO);

  double bfac = BFAC*(*_bfield)();

  if( !zero_bfield ) {
    /*
    // check that curvature is small enough
    double r_min = 0.5*sqrt(dx*dx + dy*dy);
    double a34_hat2 = dxdz2*dxdz2+dydz2*dydz2;
    double pt_p = 1./sqrt(1. + 1./a34_hat2);
    double a5_max = fabs(pt_p/(r_min*bfac));
    if ( fabs(vec(IQP)) > a5_max ) return fill_tuple(LARGE_CURVATURE);
    */
    Circle cl(x1,y1,x2,y2,_xv,_yv);
    if( !cl.is_dphi_ok() ) return fill_tuple(DPHI_TOO_LARGE);
    if(fabs(cl.get_curvature()/bfac)>_qpt_max)return fill_tuple(LARGE_CURVATURE);
  }

  // check curvature and dhpi
  
  vec(IX)    = x2;
  vec(IY)    = y2;
  vec(IDXDZ) = dxdz2;
  vec(IDYDZ) = dydz2;
  //  vec(IQP_Z) stays the same

  trv.set_vector(vec);
  
  // set new direction of the track assuming backward in time propagation
  (dz<0.)?trv.set_forward():trv.set_backward();

  //propagate it to a starting surface
  
  PropStat pstat = _prop->vec_dir_prop(trv,*_srf,Propagator::FORWARD);
  if ( ! pstat.success() ) return fill_tuple(PROP_FAILED);

  // construct and fit a new track
  ETrack tre( trv , _err );
  trv.is_forward() ?  tre.set_forward() : tre.set_backward() ;

  double chsq = 0.0;
 
  // add hit to a list of hits

  HitList::const_iterator ihit;
  HitList hits = trh.get_hits();
  hits.push_back(phit2);

  // Loop over hits and add each in turn to the track.
  
  for ( ihit=hits.begin(); ihit!=hits.end(); ++ihit ) {
    const Surface& srf = (*ihit)->get_surface();
    PropStat pstat =  _prop->err_dir_prop(tre,srf,Propagator::BACKWARD);
    if ( ! pstat.success() ) return fill_tuple(PROP_FAILED);
    int stat = _fit->add_hit_fit(tre,chsq,*ihit);
    if ( stat )  return fill_tuple(ENCLOSED_FITTER_ERROR);
  }

  // Fit was successful: update fit and add new hit.
  chsq=0.;
  trh.add_hit(phit2);
  trh.set_fit(tre,chsq);
  return fill_tuple(OK);
}
//**********************************************************************

// Return the creator.

ObjCreator AddFitZPlane_XY2_Z2::get_creator() {
  return create;
}

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

// Write the object data.

ObjData AddFitZPlane_XY2_Z2::write_data() const {
  ObjData data( get_type_name() );
  data.add_double( "bfield", get_bfield());
  data.add_double( "ptmin", get_ptmin());
  data.add_double( "vx",get_vx());
  data.add_double( "vy",get_vy());
  data.add_share_pointer( "propagator", _prop );
  data.add_share_pointer( "fitter", _fit );
  data.add_share_pointer( "surface", _srf );
  std::vector<double> error;
  for(int i=0;i<5;++i)
    for(int j=0;j<=i;++j)
      error.push_back(_err(i,j));
  data.add_double_list("error",error);
  return data;
}

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