// PropXYCyl.cpp

#include "PropXYCyl.h"
#include <iostream>
#include "trfbase/PropStat.h"
#include "trfutil/TRFMath.h"
#include "trfcyl/SurfCylinder.h"
#include "trfxyp/SurfXYPlane.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() == "PropXYCyl" );
  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 << "PropXYCyl: can't get bfield parameter." << endl;
    abort();
  }
  return ObjPtr( new PropXYCyl(bfield) );
}

}
//**********************************************************************

// Assign track parameter indices.
enum {
  IV = SurfXYPlane::IV,
  IZC   = SurfXYPlane::IZ,
  IDVDU = SurfXYPlane::IDVDU,
  IDZDU = SurfXYPlane::IDZDU,
  IQP_XY  = SurfXYPlane::IQP,

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

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

// Private class STCalcXY.
//
// An STCalcXY_ 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 STCalcXY {
private:
  bool _big_crv;
  double _st;
  double _dst_dphi21;
  double _dst_dcrv1;
  double _dst_dr1;
  double _cnst1,_cnst2;
public:
  // constructor
  STCalcXY();
  STCalcXY(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;
};

STCalcXY::STCalcXY() { }

STCalcXY::STCalcXY(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 STCalcXY::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 STCalcXY::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 STCalcXY::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 XYPlane:
// u (cm) is fixed 
// 0 - v (cm)
// 1 - z (cm)
// 2 - dv/du 
// 3 - dz/du 
// 4 - q/p   p is momentum of a track, q is its charge

PropStat 
vec_propagatexycyl_( 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 XYPlane.
  assert( srf1.get_pure_type() == SurfXYPlane::get_static_type() );
  if ( srf1.get_pure_type( ) != SurfXYPlane::get_static_type() )
    return pstat;
  const SurfXYPlane& sxyp1 = (const SurfXYPlane&) 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 dist and phi of the plane and the starting track vector.
 int iphi  = SurfXYPlane::NORMPHI;
 int idist = SurfXYPlane::DISTNORM;
 double phi = sxyp1.get_parameter(iphi);
 double    r = sxyp1.get_parameter(idist);

 TrackVector vec = trv.get_vector();
 double v = vec(IV);                  // v
 double z = vec(IZC);                  // z
 double b = vec(IDVDU);               // dv/du
 double a = vec(IDZDU);               // dz/du
 double e = vec(IQP_XY);              // q/p

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

 // Calculate cylindrical coordinates
  
  double cnst=sign_du>0?0.:PI;
  double atn=atan(v/r);
  assert(r!=0.);

  double phi1= fmod2(phi+atn,TWOPI);
  double r1=sqrt(r*r+v*v);
  double z1=z;
  double alp1= fmod2(atan(b)-atn+cnst,TWOPI);
  double tlm1= a*sign_du/sqrt(1+b*b);
  double qpt1=e*sqrt((1+a*a+b*b)/(1+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.
  STCalcXY sto1(r1,phi1,alp1,crv1,r2,phi21,alp21);
  STCalcXY 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;
  STCalcXY 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();

  // Set the return status.
  double s = st*sqrt(1.0+tlm2*tlm2);
  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_dv= r/(r*r+v*v);

  // dz1

  double dz1_dz=1.;

  // dalf1

  double dalp1_db= 1./(1.+b*b);
  double dalp1_dv= -r/(r*r+v*v);

  // dr1

  double dr1_dv= v/sqrt(v*v+r*r);

  // dtlm1

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

  // dcrv1

  double dqpt1_de= sqrt((1+a*a+b*b)/(1+b*b));
  double dqpt1_da= a*e/sqrt((1+b*b)*(1+a*a+b*b));
  double dqpt1_db= -e*b*a*a/sqrt(1+a*a+b*b)/sqrt(1+b*b)/(1+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_dv=dp2dp1*dphi1_dv+dp2da1*dalp1_dv+dp2dr1*dr1_dv;
  double dphi2_db=dp2da1*dalp1_db+dp2dc1*dcrv1_db;
  double dphi2_da=dp2dc1*dcrv1_da;
  double dphi2_de=dp2dc1*dcrv1_de;

  // alp2
  double dalp2_dv= da2da1*dalp1_dv+da2dr1*dr1_dv;
  double dalp2_db= da2da1*dalp1_db+da2dc1*dcrv1_db;
  double dalp2_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_dz= dz2dz1*dz1_dz;
  double dz2_dv= dz2dr1*dr1_dv+dz2da1*dalp1_dv;
  double dz2_db= dz2da1*dalp1_db+dz2dl1*dtlm1_db+dz2dc1*dcrv1_db;
  double dz2_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,IV) =dphi2_dv;
  deriv(IPHI,IDVDU)= dphi2_db;
  deriv(IPHI,IDZDU)= dphi2_da;
  deriv(IPHI,IQP_XY)= dphi2_de;
  deriv(IZ,IV)   = dz2_dv;
  deriv(IZ,IZC)  = dz2_dz;
  deriv(IZ,IDVDU) =dz2_db;
  deriv(IZ,IDZDU) =dz2_da;
  deriv(IZ,IQP_XY)=dz2_de; 
  deriv(IALF,IV) = dalp2_dv;
  deriv(IALF,IDVDU) = dalp2_db;
  deriv(IALF,IDZDU) = dalp2_da;
  deriv(IALF,IQP_XY) = dalp2_de;
  deriv(ITLM,IDVDU) = dtlm2_db;
  deriv(ITLM,IDZDU) = dtlm2_da;
  deriv(IQPT,IDVDU) = dqpt2_db;
  deriv(IQPT,IDZDU) = dqpt2_da;
  deriv(IQPT,IQP_XY) = dqpt2_de;
  /*
  if( abs(bfac) < 1e-10 ) {
    deriv(IQPT,IDVDU) = 0.;
    deriv(IQPT,IDZDU) = 0.;
    deriv(IQPT,IQP_XY) = 1.;
  }
  */
  return pstat;

}

//************************************************************************
// PropXYCyl member functions.
//************************************************************************

// output stream

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

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

// Constructor.

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

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

// Destructor.

PropXYCyl::~PropXYCyl() { }

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

// Clone.

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

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

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


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

// return the magnetic field

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

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

// Return the creator.

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

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

// Write the object data.

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

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