// PropZCyl.cpp

#include "PropZCyl.h"
#include <iostream>
#include "trfbase/PropStat.h"
#include "trfutil/TRFMath.h"
#include "trfcyl/SurfCylinder.h"
#include "trfzp/SurfZPlane.h"
#include "trfbase/VTrack.h"
#include "trfbase/ETrack.h"
#include "trfutil/trfstream.h"
#include "objstream/ObjData.hpp"
#include "objstream/ObjTable.hpp"
#include <cassert>

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

#ifndef DEFECT_NO_STDLIB_NAMESPACES
using std::sqrt;
using std::cos;
using std::sin;
using std::atan2;
using std::atan;
using std::asin;
using namespace trf;
#endif

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

namespace {

// Creator.

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

}
//**********************************************************************
// Assign track parameter indices.
enum {
  IX = SurfZPlane::IX,
  IY   = SurfZPlane::IY,
  IDXDZ = SurfZPlane::IDXDZ,
  IDYDZ = SurfZPlane::IDYDZ,
  IQP_Z  = SurfZPlane::IQP,

  IPHI = SurfCylinder::IPHI,
  IZ   = SurfCylinder::IZ,
  IALF = SurfCylinder::IALF,
  ITLM = SurfCylinder::ITLM,
  IQPT  = SurfCylinder::IQPT
};

//**********************************************************************
// helpers
//**********************************************************************

// Private class STCalcZ.
//
// An STCalcZ_ 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 STCalcZ {
private:
  bool _big_crv;
  double _st;
  double _dst_dphi21;
  double _dst_dcrv1;
  double _dst_dr1;
  double _cnst1,_cnst2;
public:
  // constructor
  STCalcZ();
  STCalcZ(double r1, double phi1, double alf1, double crv,
         double r2, double phi2, double alf2);
  double st() const { return _st; };
  double d_st_dalp1(double d_phi2_dalf1, double d_alf2_dalf1 ) const;
  double d_st_dcrv1(double d_phi2_dcrv1, double d_alf2_dcrv1 ) const;
  double d_st_dr1(  double d_phi2_dr1,   double d_alf2_dr1   ) const; 
  double _crv1;
};

STCalcZ::STCalcZ() { }

STCalcZ::STCalcZ(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( fabs(phi_dir_diff) <= PI );

  // Evaluate whether |C| is" big"
  double d2 = r1*r1 + r2*r2 - 2.0*r1*r2*cos(phi2-phi1);
  double d = d2>0.? sqrt(d2) : 0.;
  _big_crv = rmax*fabs(crv1) > 0.001 || fabs(r1-r2)<1.e-9 || d < 1e-13;
  if( fabs(crv1) < 1.e-10 ) _big_crv=false;

  // 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;
    double xsign = (crv1*_st == 0. ? 0.: fabs( (phi_dir_diff - _st*crv1) / (_st*crv1))  );   
    if ( crv1*_st ) {
      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 ( fabs(alf2) > fabs(alf1) ) sign = -1.0;
      if ( fabs(alf2) == fabs(alf1) ) {
        if ( fabs(alf2) < PI2 ) {
          if ( r2 < r1 ) sign = -1.0;
        } else {
          if ( r2 > r1 ) sign = -1.0;
        }
      }
    }

    if ( fabs(r2 - r1)<1e-12 && crv1*_st == 0. && fabs(alf1) < PI2 ) sign = -1.0;
    if ( fabs(r2 - r1)<1e-12 && crv1*_st == 0. && fabs(alf1) > PI2 ) sign = 1.0;

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

    // save derivatives
    _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 );
    if( fabs(_st) > 1e-11 ) {
      _dst_dphi21 = sign*(r1*r2*sin(phi2-phi1))*root/d;
      _dst_dr1= (1.+arg2/2.*(1+3./4.*arg2))/d*sign;
      _cnst1=r1-r2*cos(phi2-phi1);
      _cnst2=r1*r2*sin(phi2-phi1);
    }
    else {
      _dst_dphi21= r1;
      _dst_dr1= sign;
      _cnst1=0.;
      _cnst2=0.;
    }
  }

}

double STCalcZ::d_st_dalp1(double dphi2_dalf1, double dalf2_dalf1) const {
  if ( _big_crv ) return ( dphi2_dalf1 + dalf2_dalf1 - 1.0 ) / _crv1;
  else return _dst_dphi21 * dphi2_dalf1;
}

double STCalcZ::d_st_dcrv1(double dphi2_dcrv1, double dalf2_dcrv1) const {
  if ( _big_crv ) return ( dphi2_dcrv1 + dalf2_dcrv1 - _st ) / _crv1;
  else return _dst_dcrv1 + _dst_dphi21*dphi2_dcrv1;
}

double STCalcZ::d_st_dr1(double dphi2_dr1, double dalf2_dr1) const {
  if ( _big_crv ) return ( dphi2_dr1 + dalf2_dr1 ) / _crv1;
  else {
    if( fabs(_cnst2)<1e-12 && fabs(_cnst1)<1e-12 ) {
      return _dst_dr1*sqrt(1. + _dst_dphi21*_dst_dphi21*dphi2_dr1*dphi2_dr1);
    }
    else
      return _dst_dr1*(_cnst1+_cnst2*dphi2_dr1);
  }
}

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

// Private function to propagate a track without error
//
// The track parameters for a cylinder are:
// phi z alpha tan(lambda) curvature
// (NOT [. . . lambda .] as in TRF and earlier version of TRF++.)
//
// If pderiv is nonzero, return the derivative matrix there.
// On Cylinder:
// r (cm) is fixed 
// 0 - phi 
// 1 - z (cm)
// 2 - alp 
// 3 - tlam 
// 4 - q/pt   pt is transverse momentum of a track, q is its charge
// On ZPlane:
// 0 - x (cm)
// 1 - y (cm)
// 2 - dx/dz 
// 3 - dy/dz 
// 4 - q/p   p is momentum of a track, q is its charge
// If pderiv is nonzero, return the derivative matrix there.

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

  // construct return status
  PropStat pstat;

  // fetch the originating surface and vector
  const Surface& srf1 = *trv.get_surface();
  TrackVector vec1 = trv.get_vector();

  // Check origin is a Zplane.
  assert( srf1.get_pure_type() == SurfZPlane::get_static_type() );
  if ( srf1.get_pure_type( ) != SurfZPlane::get_static_type() )
    return pstat;
  const SurfZPlane& szp1 = (const SurfZPlane&) srf1;

  // Check destination is a cylinder.
  assert( srf.get_pure_type() == SurfCylinder::get_static_type() );
  if ( srf.get_pure_type( ) != SurfCylinder::get_static_type() )
    return pstat;
  const SurfCylinder& scy2 = (const SurfCylinder&) srf;


 // Fetch the z of the plane and the starting track vector.
 int iz  = SurfZPlane::ZPOS;
 double z = szp1.get_parameter(iz);

 TrackVector vec = trv.get_vector();
 double x = vec(IX);                  // v
 double y = vec(IY);                  // z
 double b = vec(IDXDZ);               // dv/du
 double a = vec(IDYDZ);               // dz/du
 double e = vec(IQP_Z);              // q/p

 // Fetch the radii and the starting track vector.
  int irad = SurfCylinder::RADIUS;
  double r2 = scy2.get_parameter(irad);
  
 int sign_dz = 0;
 if(trv.is_forward())  sign_dz =  1;
 if(trv.is_backward()) sign_dz = -1;
 if(sign_dz == 0){
  cerr << "PropZCyl::_vec_propagate: Unknown direction of a track "<< endl;
  abort();
 }

 // Calculate cylindrical coordinates
  
  double cnst1=x>0?0.:PI;
  double cnst2=PI;
  if(a>0.&&b>0.&&sign_dz>0.) cnst2=0.;
  if(a<0.&&b>0.&&sign_dz>0.) cnst2=0.;
  if(a>0.&&b<0.&&sign_dz<0.) cnst2=0.;
  if(a<0.&&b<0.&&sign_dz<0.) cnst2=0.;

  double sign_y= y>0 ? 1.:-1;
  double sign_a= a>0 ? 1.:-1;
  if(y==0) sign_y=0.;
  if(a==0) sign_a=0.;
  double atnxy=(x!=0.?atan(y/x):sign_y*PI/2.);
  double atnab=(b!=0.?atan(a/b):sign_a*PI/2.);

  double phi1= fmod2(atnxy+cnst1,TWOPI);
  double r1=sqrt(x*x+y*y);
  // 28jan01 dla
  // Handle r1 = 0 to avaoid crash later on.
  if ( r1 < 0.001 ) {
    return pstat;
  }
  double z1=z;
  double alp1= fmod2(atnab-phi1+cnst2,TWOPI);
  double tlm1= sign_dz/sqrt(a*a+b*b);
  double qpt1=e*sqrt((1+a*a+b*b)/(a*a+b*b));
  
  

  // Check alpha range.
  alp1 = fmod2( alp1, TWOPI );
  assert( fabs(alp1) <= PI );
  
  //if ( trv.is_forward() ) assert( fabs(alp1) <= PI2 );
  //else assert( fabs(alp1) > PI2 );

  // Calculate the cosine of lambda.
  //double clam1 = 1.0/sqrt(1+tlm1*tlm1);

  // Calculate curvature: C = _bfac*(q/p)/cos(lambda)
  // and its derivatives
  // assert( clam1 != 0.0 );
  // double dcrv1_dqp1 = bfac/clam1;
  // double crv1 = dcrv1_dqp1*qp1;
  // double dcrv1_dtlm1 = crv1*clam1*clam1*tlm1;

  // Calculate the curvature = _bfac*(q/pt)
  double dcrv1_dqpt1 = bfac;
  double crv1 = dcrv1_dqpt1*qpt1;
  //double dcrv1_dtlm1 = 0.0;

  // Evaluate the new track vector.
  // See dla log I-044

  // lambda and curvature do not change
  double tlm2 = tlm1;
  double crv2 = crv1;
  double qpt2 = qpt1;

  // We can evaluate sin(alp2), leaving two possibilities for alp2
  // 1st solution: alp21, phi21, phid21, tht21
  // 2nd solution: alp22, phi22, phid22, tht22
  // evaluate phi2 to choose
  double salp1 = sin( alp1 );
  double calp1 = cos( alp1 );
  double salp2 = r1/r2*salp1 + 0.5*crv1/r2*(r2*r2-r1*r1);
  // if salp2 > 1, track does not cross cylinder
  if ( fabs(salp2) > 1.0 ) return pstat;
  double alp21 = asin( salp2 );
  double alp22 = alp21>0 ? PI-alp21 : -PI-alp21;
  double calp21 = cos( alp21 );
  double calp22 = cos( alp22 );
  double phi20 = phi1 + atan2( salp1-r1*crv1, calp1 );
  double phi21 = phi20 - atan2( salp2-r2*crv2, calp21 );   // phi position
  double phi22 = phi20 - atan2( salp2-r2*crv2, calp22 );
  // Construct an sT object for each solution.
  STCalcZ sto1(r1,phi1,alp1,crv1,r2,phi21,alp21);
  STCalcZ sto2(r1,phi1,alp1,crv1,r2,phi22,alp22);
  // Check the two solutions are nonzero and have opposite sign 
  // or at least one is nonzero.

  // Choose the correct solution
  bool use_first_solution;
  switch (dir) {
  case Propagator::NEAREST:
    use_first_solution = abs(sto2.st()) > abs(sto1.st());
    break;
  case Propagator::FORWARD:
    use_first_solution = sto1.st() > 0.0;
    break;
  case Propagator::BACKWARD:
    use_first_solution = sto1.st() < 0.0;
    break;
  default:
    cerr << "PropCyl::_vec_propagate: Unknown direction." << endl;
    abort();
  }

  // Assign phi2, alp2 and sto2 for the chosen solution.
  double phi2, alp2;
  STCalcZ sto;
  double calp2;
  if ( use_first_solution ) {
    sto = sto1;
    phi2 = phi21;
    alp2 = alp21;
    calp2 = calp21;
   } else {
    sto = sto2;
    phi2 = phi22;
    alp2 = alp22;
    calp2 = calp22;
   }

  // fetch sT.
  double st = sto.st();

  // use sT to evaluate z2
  double z2 = z1 + tlm1*st;

  // Check alpha range.
  assert( fabs(alp2) <= PI );

  // put new values in vec
  vec(IPHI) = phi2;
  vec(IZ)   = z2;
  vec(IALF) = alp2;
  vec(ITLM) = tlm2;
  vec(IQPT) = qpt2;

  // Update trv
  trv.set_surface(SurfacePtr(srf.new_pure_surface()));
  trv.set_vector(vec);
  if ( fabs(alp2) <= PI2 ) trv.set_forward();
  else trv.set_backward();

  // Evaluate s.
  double s = st*sqrt(1.0+tlm2*tlm2);

  // Set the return status.
  pstat.set_path_distance(s);
  //st > 0 ? pstat.set_forward() : pstat.set_backward();

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

  // Calculate derivatives.
  // dphi1

  double dphi1_dx= -y/(x*x+y*y);
  double dphi1_dy= x/(x*x+y*y);

  // dz1

  //double dz1_dz=0.;

  // dalf1

  double dalp1_da= b/(a*a+b*b);
  double dalp1_db= -a/(a*a+b*b);
  double dalp1_dy= -dphi1_dy;
  double dalp1_dx= -dphi1_dx;

  // dr1

  double dr1_dx= x/sqrt(x*x+y*y);
  double dr1_dy= y/sqrt(x*x+y*y);

  // dtlm1

  double dtlm1_da= -sign_dz*a/sqrt(a*a+b*b)/(a*a+b*b);
  double dtlm1_db= -sign_dz*b/(a*a+b*b)/sqrt(a*a+b*b);

  // dcrv1

  double dqpt1_de= sqrt((1+a*a+b*b)/(a*a+b*b));
  double dqpt1_da= -a*e/sqrt((a*a+b*b)*(1+a*a+b*b))/(a*a+b*b);
  double dqpt1_db= -e*b/sqrt(1+a*a+b*b)/sqrt(a*a+b*b)/(a*a+b*b);
  
  double dcrv1_de= dqpt1_de*bfac;
  double dcrv1_da= dqpt1_da*bfac;
  double dcrv1_db= dqpt1_db*bfac;

  // alpha_2
  double da2da1 = r1*calp1/r2/calp2;
  double da2dc1 = (r2*r2-r1*r1)*0.5/r2/calp2;
  double da2dr1 = (salp1-crv2*r1)/r2/calp2;

  // phi2
  double rcsal1 = r1*crv1*salp1;
  double rcsal2 = r2*crv2*salp2;
  double den1 = 1.0 + r1*r1*crv1*crv1 - 2.0*rcsal1;
  double den2 = 1.0 + r2*r2*crv2*crv2 - 2.0*rcsal2;
  double dp2dp1 = 1.0;
  double dp2da1 = (1.0-rcsal1)/den1 - (1.0-rcsal2)/den2*da2da1;
  double dp2dc1 = -r1*calp1/den1 + r2*calp2/den2 
                  - (1.0-rcsal2)/den2*da2dc1;
  double dp2dr1= -crv1*calp1/den1-(1.0-rcsal2)*da2dr1/den2;

  // z2
  //double dz2dz1 = 1.0;
  double dz2dl1 = st;
  double dz2da1 = tlm1*sto.d_st_dalp1(dp2da1, da2da1);
  double dz2dc1 = tlm1*sto.d_st_dcrv1(dp2dc1, da2dc1);
  double dz2dr1 = tlm1*sto.d_st_dr1(  dp2dr1, da2dr1);


  // final derivatives

  // phi2
  double dphi2_dx=dp2dp1*dphi1_dx+dp2da1*dalp1_dx+dp2dr1*dr1_dx;
  double dphi2_dy=dp2dp1*dphi1_dy+dp2da1*dalp1_dy+dp2dr1*dr1_dy;
  double dphi2_db=dp2da1*dalp1_db+dp2dc1*dcrv1_db;
  double dphi2_da=dp2da1*dalp1_da+dp2dc1*dcrv1_da;
  double dphi2_de=dp2dc1*dcrv1_de;

  // alp2
  double dalp2_dx= da2da1*dalp1_dx+da2dr1*dr1_dx;
  double dalp2_dy= da2da1*dalp1_dy+da2dr1*dr1_dy;
  double dalp2_db= da2da1*dalp1_db+da2dc1*dcrv1_db;
  double dalp2_da= da2da1*dalp1_da+da2dc1*dcrv1_da;
  double dalp2_de= da2dc1*dcrv1_de;

  // crv2
  double dqpt2_da=dqpt1_da;
  double dqpt2_db=dqpt1_db;
  double dqpt2_de=dqpt1_de;

  // tlm2
  double dtlm2_da= dtlm1_da;
  double dtlm2_db= dtlm1_db;

  // z2
  double dz2_dx= dz2dr1*dr1_dx+dz2da1*dalp1_dx;
  double dz2_dy= dz2dr1*dr1_dy+dz2da1*dalp1_dy;
  double dz2_db= dz2da1*dalp1_db+dz2dl1*dtlm1_db+dz2dc1*dcrv1_db;
  double dz2_da= dz2da1*dalp1_da+dz2dl1*dtlm1_da+dz2dc1*dcrv1_da;
  double dz2_de= dz2dc1*dcrv1_de;


  // Build derivative matrix.
  TrackDerivative& deriv = *pderiv;
  deriv.fill( 0.0 );
  deriv(IPHI,IX) =dphi2_dx;
  deriv(IPHI,IY) =dphi2_dy;
  deriv(IPHI,IDXDZ)= dphi2_db;
  deriv(IPHI,IDYDZ)= dphi2_da;
  deriv(IPHI,IQP_Z)= dphi2_de;
  deriv(IZ,IX)   = dz2_dx;
  deriv(IZ,IY)  = dz2_dy;
  deriv(IZ,IDXDZ) =dz2_db;
  deriv(IZ,IDYDZ) =dz2_da;
  deriv(IZ,IQP_Z)=dz2_de; 
  deriv(IALF,IX) = dalp2_dx;
  deriv(IALF,IY) = dalp2_dy;
  deriv(IALF,IDXDZ) = dalp2_db;
  deriv(IALF,IDYDZ) = dalp2_da;
  deriv(IALF,IQP_Z) = dalp2_de;
  deriv(ITLM,IDXDZ) = dtlm2_db;
  deriv(ITLM,IDYDZ) = dtlm2_da;
  deriv(IQPT,IDXDZ) = dqpt2_db;
  deriv(IQPT,IDYDZ) = dqpt2_da;
  deriv(IQPT,IQP_Z) = dqpt2_de;
  /*
  if( abs(bfac) < 1e-10 ) {
    deriv(IQPT,IDXDZ) = 0.;
    deriv(IQPT,IDYDZ) = 0.;
    deriv(IQPT,IQP_Z) = 1.;
  }
  */
  return pstat;

}

//************************************************************************
// PropZCyl member functions.
//************************************************************************

// output stream

void PropZCyl::ostr( std::ostream& stream ) const
{ 
  stream << begin_object ;
  stream << "ZPlane-Cylinder propagation with constant "
           << get_bfield() << " Tesla field"; 
  stream << end_object;
}

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

// Constructor.

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

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

// Destructor.

PropZCyl::~PropZCyl() { }

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

// Clone.

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

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

// propagate a track without error in the specified direction
PropStat PropZCyl::
vec_dir_prop(VTrack& trv, const Surface& srf, PropDir dir,TrackDerivative* pder ) const {
  Propagator::reduce_direction(dir);
  double bfac = -BFAC * (*_bfield)();
  return vec_propagatezcyl_( bfac, trv, srf, dir, pder );
}


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

// return the magnetic field

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

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

// Return the creator.

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

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

// Write the object data.

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

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