// PropCylDCA.cpp

#include "PropCylDCA.h"
#include <iostream>
#include <cassert>
#include <cmath>
#include "objstream/ObjData.hpp"
#include "objstream/ObjTable.hpp"
#include "trfutil/trfstream.h"
#include "trfutil/TRFMath.h"
#include "trfbase/TrackVector.h"
#include "trfbase/VTrack.h"
#include "trfbase/ETrack.h"
#include "trfbase/PropStat.h"
#include "trfcyl/SurfCylinder.h"
#include "SurfDCA.h"

#ifndef DEFECT_CMATH_NOT_STD
using std::sqrt;
using std::cos;
using std::sin;
using std::atan2;
using std::asin;
using std::abs;
#endif
using std::cerr;
using std::endl;

using namespace trf;

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

namespace {

// Creator.

ObjPtr create(const ObjData& data) {
  assert( data.get_object_type() == "PropCylDCA" );
  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 << "PropCylDCA: can't get bfield parameter." << endl;
    abort();
  }
  return ObjPtr( new PropCylDCA(bfield) );
}

}

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

// Assign track parameter indices

enum {
  IPHI     = SurfCylinder::IPHI,
  IZ_CYL   = SurfCylinder::IZ,
  IALF     = SurfCylinder::IALF,
  ITLM_CYL = SurfCylinder::ITLM,
  IQPT_CYL = SurfCylinder::IQPT,

  IRSIGNED = SurfDCA::IRSIGNED,
  IZ_DCA   = SurfDCA::IZ,
  IPHID    = SurfDCA::IPHID,
  ITLM_DCA = SurfDCA::ITLM,
  IQPT_DCA = SurfDCA::IQPT
};

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

// Private class ST_CylDCA.
//
// An ST_CylDCA_ object calculates sT (the signed transverse path length)
// and its derivatives w.r.t. alf1 and crv1.  It is constructed from
// the starting (r1, phi1, alf1, crv1) and final track parameters
// (r2, phi2, alf2) assuming these are consistent.  Methods are
// provided to retrieve sT and the two derivatives.

class ST_CylDCA {
private:
  bool _big_crv;
  double _st;
  double _dst_dr2;
  double _dst_dcrv1;
  double _dst_dphi2;
  double  _dst_dr1;
  double _dst_dphi2_or;
public:
  // constructor
  ST_CylDCA();
  ST_CylDCA(double r1, double phi1, double alf1, double crv,
         double r2, double phi2, double alf2);
  double st() const { return _st; };
  double d_st_dalf1( double d_r2_dalf1, double d_phi2_dalf1 ) const;
  double d_st_dalf1_or( double r1,double d_r2_dalf1_or, double d_phi2_dalf1_m1_or ) const;
  double d_st_dcrv1( double d_r2_dcrv1, double d_phi2_dcrv1 ) const;
  double d_st_dr1() const {return _dst_dr1;}
  double _crv1;
};

ST_CylDCA::ST_CylDCA() { }

ST_CylDCA::ST_CylDCA(double r1, double phi1, double alf1, double crv1,
               double r2, double phi2, double alf2) {

  _crv1 = crv1;
  assert( r1 >= 0.0 );
  assert( r2 >= 0.0 );
  double rmax = r1+r2;

  // Calculate the change in xy direction
  double phi_dir_diff = fmod2(phi2+alf2-phi1-alf1,TWOPI);
  assert( abs(phi_dir_diff) <= PI );

  // Evaluate whether |C| is" big"
  _big_crv = rmax*abs(crv1) > 0.001;

  // If the curvature is big we can use
  // sT = (phi_dir2 - phi_dir1)/crv1
  if ( _big_crv ) {
    assert( crv1 != 0.0 );
    _st = phi_dir_diff/crv1;
  }

  // Otherwise, we calculate the straight-line distance
  // between the points and use an approximate correction
  // for the (small) curvature.
  else {

    // evaluate the distance
    double d = sqrt( r1*r1 + r2*r2 - 2.0*r1*r2*cos(phi2-phi1) );
    double arg = 0.5*d*crv1;
    double arg2 = arg*arg;
    float st_minus_d = d*arg2*( 1.0/6.0 + 3.0/40.0*arg2 );
    _st = d + st_minus_d;

    // evaluate the sign
    // We define a metric xsign = abs( (dphid-d*C)/(d*C) ).
    // Because sT*C = dphid and d = abs(sT):
    // xsign = 0 for sT > 0
    // xsign = 2 for sT < 0
    // Numerical roundoff will smear these predictions.
    double sign = 0.0;
    if ( crv1*_st ) {
      double xsign = abs( (phi_dir_diff - _st*crv1) / (_st*crv1) );
      if ( xsign < 0.5 ) sign = 1.0;
      if ( xsign > 1.5  &&  xsign < 3.0 ) sign = -1.0;
    }
    // If the above is indeterminate, assume zero curvature.
    // In this case abs(alpha) decreases monotonically
    // with sT.  Track passing through origin has alpha = 0 on one
    // side and alpha = +/-pi on the other.  If both points are on
    // the same side, we use dr/ds > 0 for |alpha|<pi/2.
    if ( ! sign ) {
      sign = 1.0;
      if ( abs(alf2) > abs(alf1) ) sign = -1.0;
      if ( abs(alf2) == abs(alf1) ) {
        if ( abs(alf2) < PI2 ) {
          if ( r2 < r1 ) sign = -1.0;
        } else {
          if ( r2 > r1 ) sign = -1.0;
        }
      }
    }

    // Correct _st using the above sign.
    assert( abs(sign) == 1.0 );
    _st = sign*_st;

    // save derivatives
    if ( is_zero(d) ) {
      _dst_dcrv1 = 0.0;
      _dst_dphi2 = sign*r1;
      _dst_dphi2_or = sign;
      _dst_dr2 =  0.0;
    } else {
      _dst_dcrv1 = sign*d*d*arg*( 1.0/6.0 + 3.0/20.0*arg2);
      double root = (1.0 + 0.5*arg*arg + 3.0/8.0*arg*arg*arg*arg );
      _dst_dphi2_or = sign*(r2*sin(phi2-phi1))*root/d;
      _dst_dphi2 = r1*_dst_dphi2_or;
      _dst_dr2 =   sign*(r2-r1*cos(phi2-phi1))*root/d;
    }

  }
  _dst_dr1 = -cos(alf1)/(1.+r1*crv1*r1*crv1-2.*sin(alf1)*r1*crv1);
}

double ST_CylDCA::d_st_dalf1_or(double r1,double dr2_dalf1_or, double dphi2_dalf1_m1_or) const {
  if ( _big_crv ) return ( dphi2_dalf1_m1_or ) / _crv1;
  else return  _dst_dr2 * dr2_dalf1_or + _dst_dphi2_or * (dphi2_dalf1_m1_or*r1+1.);
}

double ST_CylDCA::d_st_dalf1(double dr2_dalf1, double dphi2_dalf1) const {
  if ( _big_crv ) return ( dphi2_dalf1 - 1. ) / _crv1;
  else return  _dst_dr2 * dr2_dalf1 + _dst_dphi2 * dphi2_dalf1;
}

double ST_CylDCA::d_st_dcrv1(double dr2_dcrv1, double dphi2_dcrv1) const {
  if ( _big_crv ) return ( dphi2_dcrv1 - _st ) / _crv1;
  else return  _dst_dr2 * dr2_dcrv1 + _dst_dphi2 * dphi2_dcrv1+ _dst_dcrv1;
}


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

PropStat
cyl_dca_propagate_(       double            _bfac,
                          VTrack&             trv, 
                    const Surface&            srf,
                    const Propagator::PropDir dir, 
                          TrackDerivative* pderiv =0 )  {

  // construct return status
  PropStat pstat;

  // fetch the originating Surface and check it is a Cylinder
  const Surface& srf_scy = *trv.get_surface();
  assert( srf_scy.get_pure_type() == SurfCylinder::get_static_type() );
  if ( srf_scy.get_pure_type( ) != SurfCylinder::get_static_type() )
    return pstat;

  // check that the destination surface is a DCA surface
  assert( srf.get_pure_type() == SurfDCA::get_static_type() );
  if ( srf.get_pure_type( ) != SurfDCA::get_static_type() )
    return pstat;
  const SurfDCA& srf_dca = (const SurfDCA&) srf;

  // original coordinates are marked with prime ( phip,alfp,r1p)
  // coordinates centered at (xv,yv) are phi,alf,r1

  // fetch the originating radius
  double r1p    = srf_scy.get_parameter(SurfCylinder::RADIUS);
  double xv = srf_dca.get_x();
  double yv = srf_dca.get_y();


  // fetch the originating TrackVector
  TrackVector vec_scy = trv.get_vector();
  double phi1p  = vec_scy(SurfCylinder::IPHI);      // phi_position
  double z1    = vec_scy(SurfCylinder::IZ);        // z
  double alf1p  = vec_scy(SurfCylinder::IALF);      // phi_dir - phi_pos
  double tlam1 = vec_scy(SurfCylinder::ITLM);      // tan(lambda)
  double qpt1  = vec_scy(SurfCylinder::IQPT);      // q/pT

  double cosphi1p=cos(phi1p);
  double sinphi1p=sin(phi1p);

  double x1=r1p*cosphi1p-xv;
  double y1=r1p*sinphi1p-yv;

  double phi1 = atan2(y1,x1);       // phi position in (xv,yv) coord system
  double r1 = sqrt(x1*x1+y1*y1);    // r1 in (xv,yv) coord system
  double alf1 = alf1p + phi1p - phi1; // alf1 in (xv,yv) coord system

  // calculate salf1, calf1 and crv1
  double salf1 = sin( alf1 );
  double calf1 = cos( alf1 );  
  double  crv1 = _bfac * qpt1;

  // calculate r2 and alf2 of the destination DCA Surface

  double r2;
  const double sign_alf1 = alf1 > 0.0 ? 1.0 : -1.0;
  double sign_alf2 = 1.0;

  double del = (r1*crv1 - 2.0*salf1);
  double rcdel = r1*crv1*del;
  double sroot = sqrt(1.0 + rcdel);
  double sroot_minus_one = sroot - 1.0;
  if ( rcdel < 1.e-6 ) 
    sroot_minus_one = 0.5*rcdel - 0.125*rcdel*rcdel + 0.0625*rcdel*rcdel*rcdel;

  if ( fabs(del)<1e-14 ) {
    sign_alf2 = sign_alf1;
    r2 = 0.0;
  } else if ( is_zero(crv1) ) {
    sign_alf2 = sign_alf1;
    r2 = r1 * abs(salf1);
  } else {
    sign_alf2 = crv1*r1 > 2.0*salf1 ? -1.0 : 1.0;
    r2 = -sign_alf2*sroot_minus_one/crv1;
    assert( r2 > 0.0 );
  }

  double alf2 = sign_alf2 * PI2;

  // calculate phi2 of the destination DCA Surface
  double sign_crv = 1.;

  // calculate tlam2, qpt2 and crv2 of the destination DCA Surface
  double tlam2 = tlam1;
  double qpt2  = qpt1;
  double crv2  = crv1;


  double salf2 = sign_alf2 > 0 ? 1.: (sign_alf2 < 0 ? -1.: 0. );
  //double calf2 = 0.;
  double cnst = salf2-r2*crv2 > 0. ? PI2: (sign_alf2 == 0. ? 0. : -PI2);
  double phi2 = phi1 + atan2( salf1-r1*crv1, calf1 ) 
                     - cnst;
  if(crv1 == 0. ) phi2= phi1 + alf1 - cnst;

  double phid2 = phi2 + sign_alf2 * PI2;
  phid2 = fmod2( phid2, TWOPI );

  // construct an object of ST_CylDCA
  ST_CylDCA sto(r1,phi1,alf1,crv1,r2,phi2,alf2);
  // fetch sT
  double st = sto.st();
  double s = st*sqrt(1+tlam1*tlam1);

  // calculate z2 of the destination DCA Surface
  double z2 = z1 + st * tlam1;

  // construct the destination TrackVector
  TrackVector vec_dca;
  vec_dca(IRSIGNED) = sign_alf2 * r2; 
  vec_dca(IZ_DCA)   = z2;
  vec_dca(IPHID)    = phid2;
  vec_dca(ITLM_DCA) = tlam2;
  vec_dca(IQPT_DCA) = qpt2;

  // set the surface of trv to the destination DCA surface
  trv.set_surface( SurfacePtr(srf_dca.new_pure_surface()) );

  // set the vector of trv to the destination TrackVector (DCA coord.)
  trv.set_vector(vec_dca);

  // set the direction of trv to the default value for VTrack on DCA
  trv.set_forward();

  // set the return status
  pstat.set_path_distance(s);

  // exit now if user did not ask for error matrix
  if ( ! pderiv ) return pstat;


  // calculate derivatives

  double dr2_dalf1_or = calf1/(sign_alf2-crv1*r2);
  double dr2_dalf1 = r1*dr2_dalf1_or;
  double dr2_dcrv1 = (r2*r2-r1*r1)/(2.0*(sign_alf2-crv1*r2));
  double dr2_dr1 = -sign_alf2/sroot*(r1*crv1 - salf1);

  double dphi2_dalf1 = sign_crv*(1.0-r1*crv1*salf1)/(sroot*sroot);
  double dphi2_dalf1_m1_or = sign_crv*(-crv1*salf1-crv1*del)/(sroot*sroot);
  double dphi2_dcrv1 = -sign_crv*(r1*calf1)/(sroot*sroot);
  double dphi2_dr1 = -crv1*calf1/(1.+r1*crv1*r1*crv1-2.*salf1*r1*crv1);

  double dst_dalf1 = sto.d_st_dalf1(dr2_dalf1, dphi2_dalf1);
  double dst_dalf1_or = sto.d_st_dalf1_or(r1,dr2_dalf1_or, dphi2_dalf1_m1_or);
  double dst_dcrv1 = sto.d_st_dcrv1(dr2_dcrv1, dphi2_dcrv1);
  double dst_dr1 = sto.d_st_dr1();

  // derivatives between phi,alf and phip (prime) alfp
  double dphi1_dphi1p_tr = r1p*cos(phi1p-phi1);
  double dalf1_dphi1p_tr = r1 - dphi1_dphi1p_tr;
  double dalf1_dalf1p = 1.;
  double dr1_dphi1p = r1p*sin(phi1-phi1p);


  // derivatives to original coordinates

  double dr2_dphi1p = dr2_dalf1_or*dalf1_dphi1p_tr + dr2_dr1*dr1_dphi1p;
  double dr2_dalf1p = dr2_dalf1 * dalf1_dalf1p;

  double dst_dphi1p = dst_dalf1_or*dalf1_dphi1p_tr + dst_dr1*dr1_dphi1p;
  double dst_dalf1p = dst_dalf1*dalf1_dalf1p;

  double dphi2_dphi1p = -dphi1_dphi1p_tr*dphi2_dalf1_m1_or + dphi2_dr1*dr1_dphi1p + dphi2_dalf1;
  double dphi2_dalf1p = dphi2_dalf1*dalf1_dalf1p;

  // build the derivative matrix

  TrackDerivative& deriv = *pderiv;
  deriv.fill( 0.0 );
  
  deriv(IRSIGNED, IPHI)     = sign_alf2 * dr2_dphi1p;
  deriv(IRSIGNED, IZ_CYL)   = 0.0;
  deriv(IRSIGNED, IALF)     = sign_alf2 * dr2_dalf1p;
  deriv(IRSIGNED, ITLM_CYL) = 0.0;
  deriv(IRSIGNED, IQPT_CYL) = sign_alf2 * _bfac * dr2_dcrv1;
 
  deriv(IZ_DCA,   IPHI)     = tlam1 * dst_dphi1p;
  deriv(IZ_DCA,   IZ_CYL)   = 1.0;
  deriv(IZ_DCA,   IALF)     = tlam1 * dst_dalf1p;
  deriv(IZ_DCA,   ITLM_CYL) = st;
  deriv(IZ_DCA,   IQPT_CYL) = tlam1 * _bfac * dst_dcrv1;

  deriv(IPHID,    IPHI)     = dphi2_dphi1p;
  deriv(IPHID,    IZ_CYL)   = 0.0;
  deriv(IPHID,    IALF)     = dphi2_dalf1p;
  deriv(IPHID,    ITLM_CYL) = 0.0;
  deriv(IPHID,    IQPT_CYL) = _bfac * dphi2_dcrv1;

  deriv(ITLM_DCA, IPHI)     = 0.0;
  deriv(ITLM_DCA, IZ_CYL)   = 0.0;
  deriv(ITLM_DCA, IALF)     = 0.0;
  deriv(ITLM_DCA, ITLM_CYL) = 1.0;
  deriv(ITLM_DCA, IQPT_CYL) = 0.0;

  deriv(IQPT_DCA, IPHI)     = 0.0;
  deriv(IQPT_DCA, IZ_CYL)   = 0.0;
  deriv(IQPT_DCA, IALF)     = 0.0;
  deriv(IQPT_DCA, ITLM_CYL) = 0.0;
  deriv(IQPT_DCA, IQPT_CYL) = 1.0;

  return pstat;

}

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

// Constructor

PropCylDCA::PropCylDCA(ObjDoublePtr bfield) :
   _bfield(bfield) { }

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

// Destructor

PropCylDCA::~PropCylDCA() { }

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

// Clone

Propagator* PropCylDCA::new_propagator() const {
  return new PropCylDCA( _bfield );
}

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

// Return the creator.

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

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

// Write the object data.

ObjData PropCylDCA::write_data() const {
  ObjData data( get_type_name() );
  data.add_double( "bfield", get_bfield() );
  return data;
}

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

// propagate a track without error in the specified direction

PropStat PropCylDCA::
vec_dir_prop( VTrack& trv, const Surface& srf, PropDir dir,
              TrackDerivative* pderiv ) const {
  double bfac = -BFAC*(*_bfield)();
  return cyl_dca_propagate_( bfac, trv, srf, dir, pderiv );
}

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

// propagate a track with error in the specified direction

// PropStat PropCylDCA::
// err_dir_prop( ETrack& tre, const Surface& srf, PropDir dir ) const {
//   TrackDerivative deriv;
//   PropStat pstat =  cyl_dca_propagate_( _bfac, tre, srf, dir, &deriv );
//   if ( ! pstat.success() ) return pstat;
//   TrackError err_cyl = tre.get_error();
//   TrackError err_dca = err_cyl.Xform(deriv);
//   tre.set_error(err_dca);
//   return pstat;
// }

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

// return the magnetic field

double PropCylDCA::get_bfield() const {
  return (*_bfield)();
}

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

// output stream

void PropCylDCA::ostr( ostream& stream ) const {
  stream << begin_object;
  stream << "Propagation from Cylinder to DCA with constant "
         << get_bfield() << " Tesla field";
  stream << end_object;
}

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