// PropCyl.cpp

#include "PropCyl.h"
#include <iostream>
#include <cassert>
#include "objstream/ObjData.hpp"
#include "objstream/ObjTable.hpp"
#include "trfutil/trfstream.h"
#include "trfbase/PropStat.h"
#include "trfutil/TRFMath.h"
#include "SurfCylinder.h"
#include "trfbase/VTrack.h"
#include "trfbase/ETrack.h"
#include "SurfCylinderParameters.h"

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

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

using trf::PI;
using trf::TWOPI;
using trf::PI2;
using trf::fmod2;
using trf::Surface;
using trf::TrackVector;
using trf::TrackDerivative;
using trf::VTrack;
using trf::ETrack;
using trf::PropStat;
using trf::Propagator;
using trf::SurfCylinder;
using trf::PropCylField;
using trf::PropCyl;

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

namespace {

// Creator.

ObjPtr create(const ObjData& data) {
  assert( data.get_object_type() == "PropCyl" );
  PropCyl* pbare(0);
  if ( data.has("bfield_map") ) {
    PropCyl::PropCylFieldPtr pmap;
    data.get_share_pointer("bfield_map", pmap);
    pbare = new PropCyl(pmap);
  } else {
    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 << "PropCyl: can't get bfield parameter." << endl;
      abort();
    }
    pbare = new PropCyl(bfield);
  }
  assert( pbare != 0 );
  Ptr<PropCyl,LocalSharedDeletePolicy> pobj(pbare);
  pobj->set_direction(data);
  return pobj;
}

}

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

// Private class STCalc.
//
// An STCalc_ 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 STCalc {
private:
  bool _big_crv;
  double _st;
  double _dst_dphi21;
  double _dst_dcrv1;
public:
  // constructor
  STCalc();
  STCalc(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_phi2_dalf1, double d_alf2_dalf1 ) const;
  double d_st_dcrv1(double d_phi2_dcrv1, double d_alf2_dcrv1 ) const;
  double _crv1;
};

STCalc::STCalc() { }

STCalc::STCalc(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"
  _big_crv = rmax*fabs(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 dr = r1 - r2;
    double d = sqrt( dr*dr + 2.0*r1*r2*(1. - 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 ( _st*crv1 ) {
      double xsign = fabs( (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 ( 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;
        }
      }
    }

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

    // save derivatives
    if ( _st < 1e-10*r1 ) {
      _dst_dcrv1 = 0.0;
      _dst_dphi21 = 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_dphi21 = sign*(r1*r2*sin(phi2-phi1))*root/d;
    }

  }

}

double STCalc::d_st_dalf1(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 STCalc::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;
}

//************************************************************************
// PropCyl member functions.
//************************************************************************

// output stream

void PropCyl::ostr(ostream& stream) const {
  stream << begin_object;
  if ( _pmap == 0 ) {
    stream << "Cylinder propagation with constant "
           << get_bfield() << " Tesla field";
  } else {
    stream << "Cylinder propagation with map " << *_pmap;
  }
  stream << end_object;
}

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

// Constructor.

PropCyl::PropCyl(ObjDoublePtr bfield) : _bfield(bfield), _pmap(0) { }

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

// Constructor From field mapper..

PropCyl::PropCyl(PropCylFieldPtr pmap) : _bfield(0), _pmap(pmap) {
  assert( _pmap != 0 );
}

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

// Destructor.

PropCyl::~PropCyl() { }

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

// Clone.

Propagator* PropCyl::new_propagator() const {
  if ( _pmap == 0 ) {
    return new PropCyl( _bfield );
  } else {
    return new PropCyl(_pmap);
  }
}

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

// Return the creator.

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

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

// Write the object data.

ObjData PropCyl::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.
//
// The track parameters for a cylinder are:
// phi z alpha tan(lambda) curvature
// (NOT [. . . lambda .] as in TRF and earlier version of TRF++.)
//
// If pder is nonzero, return the derivative matrix there.

PropStat PropCyl::
vec_dir_prop(VTrack& trv, const Surface& srf, PropDir dir,
             TrackDerivative* pder) const {

  // 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 cylinder.
  //  if(srf1.get_pure_type() != SurfCylinder::get_static_type())
  //    assert( srf1.get_pure_type() == SurfCylinder::get_static_type() );
  if ( srf1.get_pure_type( ) != SurfCylinder::get_static_type() )
    return pstat;
  const SurfCylinder& scy1 = (const SurfCylinder&) 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;

  // Separate the move status from the direction.
  PropDirected::PropDir rdir(dir);
  bool move = Propagator::reduce_direction(rdir);

  // This flag indicates the new surface is only a short hop away.
  bool short_hop = false;

  // If surfaces are the same and move is not set, then we leave the
  // track where it is.
  if ( srf.pure_equal(srf1) && !move ) {
    if ( pder != 0 ) {
      TrackDerivative& deriv = *pder;
      deriv.fill( 0.0 );
      deriv(0,0) = 1.0;
      deriv(1,1) = 1.0;
      deriv(2,2) = 1.0;
      deriv(3,3) = 1.0;
      deriv(4,4) = 1.0;
    }
    pstat.set_same();
    return pstat;
  }

  // Fetch the radii.
  int irad = SurfCylinder::RADIUS;
  double r1 = scy1.get_parameter(irad);
  double r2 = scy2.get_parameter(irad);

  // Fetch the starting track vector.
  TrackVector vec = trv.get_vector();
  double phi1 = vec(IPHI);               // phi position
  double z1 = vec(IZ);                   // z
  double alf1 = vec(IALF);               // phi_dir - phi_pos
  double tlam1 = vec(ITLM);              // tan(lambda)
  double qpt1 = vec(IQPT);               // q/pT

  // If the radii are close, we short hop if the direction is appropriate.
  double rsum = r1 + r2;
  double rdif = r2 - r1;
  double ardif = fabs(rdif);
  double minrat = 1.0e-10;
  if ( ardif/rsum < minrat ) {
    switch (rdir) {
    case Propagator::NEAREST:
      if ( r1 != r2 ) {
        short_hop = true;
      }
      break;
    case Propagator::FORWARD:
      if ( r2 > r1 ) {
        short_hop = true;
        pstat.set_forward();
      }
      break;
    case Propagator::BACKWARD:
      if ( r2 < r1 ) {
        short_hop = true;
        pstat.set_backward();
      }
      break;
    default:
      cerr << "PropCyl::_vec_propagate: Unknown direction." << endl;
      abort();
    }
  }

  // Handle a short hop. 
  // This is a very small propagation distance that is presently handled
  // by leaving the track and error unchanged.  Perhaps we should replace
  // this with a simple linear approximation.
  // Note the usual propagation code will reject a path distance of zero
  // an it is important that we intercept s=0 or close enough to round to
  // zero.
  if ( short_hop ) {
    // Update trv
    trv.set_surface( SurfacePtr(srf.new_pure_surface()) );
    if ( fabs(alf1) <= PI2 ) trv.set_forward();
    else trv.set_backward();
    // Update derivative.
    if ( pder != 0 ) {
      TrackDerivative& deriv = *pder;
      deriv.fill( 0.0 );
      deriv(0,0) = 1.0;
      deriv(1,1) = 1.0;
      deriv(2,2) = 1.0;
      deriv(3,3) = 1.0;
      deriv(4,4) = 1.0;
    }
    // Calculate ds.
    // This formula gives the correct sign.
    double clam = 1.0/sqrt(1+tlam1*tlam1);
    assert( clam != 0.0 );
    double calf = cos(alf1);
    assert( calf != 0.0 );
    double s = rdif/clam/calf;
    pstat.set_path_distance(s);
    return pstat;
  }

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

  // Calculate the curvature = bfac*(q/pt) and its derivatives.
  double bfac;
  double dcrv1_dtlam1 = 0.0;
  double dcrv1_dz1 = 0.0;
  if(_pmap == 0)
    bfac = -BFAC*(*_bfield)();
  else {
    TrackVector mapder;
    bfac = -BFAC*_pmap->get_field(r1, r2, trv, rdir, &mapder);
    // We only handle the derivative w.r.t. z and tan(lambda).
    assert( mapder(IPHI) == 0.0 );
    assert( mapder(IALF) == 0.0 );
    assert( mapder(IQPT) == 0.0 );
    dcrv1_dz1 += mapder(IZ)*qpt1;
    dcrv1_dtlam1 += mapder(ITLM)*qpt1;
  }
  double dcrv1_dqpt1 = bfac;
  double crv1 = bfac*qpt1;

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

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

  // lambda and curvature do not change
  double tlam2 = tlam1;
  double crv2 = crv1;
  double qpt2 = qpt1;

  // We can evaluate sin(alf2), leaving two possibilities for alf2
  // 1st solution: alf21, phi21, phid21, tht21
  // 2nd solution: alf22, phi22, phid22, tht22
  // evaluate phi2 to choose
  double salf1 = sin( alf1 );
  double calf1 = cos( alf1 );
  double salf2 = r1/r2*salf1 + 0.5*crv1/r2*(r2*r2-r1*r1);
  // If salf2 is close to 1 or -1, set it to that value.
  double diff = fabs( fabs(salf2) - 1.0 );
  if ( diff < 1.e-10 ) salf2 = salf2>0 ? 1.0 : -1.0;
  // if salf2 > 1, track does not cross cylinder
  if ( fabs(salf2) > 1.0 ) return pstat;
  double alf21 = asin( salf2 );
  double alf22 = alf21>0 ? PI-alf21 : -PI-alf21;
  double calf21 = cos( alf21 );
  double calf22 = cos( alf22 );
  double phi20 = phi1 + atan2( salf1-r1*crv1, calf1 );
  double phi21 = phi20 - atan2( salf2-r2*crv2, calf21 );   // phi position
  double phi22 = phi20 - atan2( salf2-r2*crv2, calf22 );
  // Construct an sT object for each solution.
  STCalc sto1(r1,phi1,alf1,crv1,r2,phi21,alf21);
  STCalc sto2(r1,phi1,alf1,crv1,r2,phi22,alf22);
  // Check the two solutions are nonzero and have opposite sign 
  // or at least one is nonzero.

  // Extract the two sT-values.
  double st1 = sto1.st();
  double st2 = sto2.st();
  double ast1 = abs(st1);
  double ast2 = abs(st2);
  
  // If both solutions have the same sign, then the most distant one
  // should actually go the other way along a full cicle minus the
  // calculated value of s.  We discard this solution by setting it
  // to zero.
  if ( st1*st2 > 0.0 ) {
    if ( ast1 > ast2 ) {
      st1 = 0.0;
    } else {
      st2 = 0.0;
    }
  }
      
  // Choose the correct solution
  bool use_first_solution;
  switch (rdir) {
  case Propagator::NEAREST:
    use_first_solution = st1!=0 && (st2==0.0 || ast1<ast2);
    break;
  case Propagator::FORWARD:
    if ( st1 > 0.0 ) {
      use_first_solution = true;
    } else if ( st2 > 0.0 ) {
      use_first_solution = false;
    } else {
      return pstat;
    }
    break;
  case Propagator::BACKWARD:
    if ( st1 < 0.0 ) {
      use_first_solution = true;
    } else if ( st2 < 0.0 ) {
      use_first_solution = false;
    } else {
      return pstat;
    }
    break;
  default:
    cerr << "PropCyl::_vec_propagate: Unknown direction." << endl;
    abort();
  }

  // Assign phi2, alf2 and sto2 for the chosen solution.
  double phi2, alf2;
  STCalc sto;
  double calf2;
  if ( use_first_solution ) {
    sto = sto1;
    phi2 = phi21;
    alf2 = alf21;
    calf2 = calf21;
   } else {
    sto = sto2;
    phi2 = phi22;
    alf2 = alf22;
    calf2 = calf22;
   }

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

  if ( st == 0.0 ) {
    return pstat;
  }

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

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

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

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

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

  // exit now if user did not ask for error matrix.
  if ( pder == 0 ) return pstat;

  // Calculate derivatives.

  // alpha_2
  double da2da1 = r1*calf1/r2/calf2;
  double da2dc1 = (r2*r2-r1*r1)*0.5/r2/calf2;

  // phi2
  double rcsal1 = r1*crv1*salf1;
  double rcsal2 = r2*crv2*salf2;
  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*calf1/den1 + r2*calf2/den2 
                  - (1.0-rcsal2)/den2*da2dc1;

  // z2
  double dz2dz1 = 1.0;
  double dz2dl1 = st;
  double dz2da1 = tlam1*sto.d_st_dalf1(dp2da1, da2da1);
  double dz2dc1 = tlam1*sto.d_st_dcrv1(dp2dc1, da2dc1);

  // Build derivative matrix.
  TrackDerivative& deriv = *pder;
  deriv.fill( 0.0 );
  deriv(IPHI,IPHI) = dp2dp1;
  deriv(IPHI,IZ)   = dp2dc1*dcrv1_dz1;
  deriv(IPHI,IALF) = dp2da1;
  deriv(IPHI,ITLM) = dp2dc1*dcrv1_dtlam1;
  deriv(IPHI,IQPT) = dp2dc1*dcrv1_dqpt1;
  deriv(IZ,IZ)   = dz2dz1 + dz2dc1*dcrv1_dz1;
  deriv(IZ,IALF) = dz2da1;
  deriv(IZ,ITLM) = dz2dl1 + dz2dc1*dcrv1_dtlam1;
  deriv(IZ,IQPT) = dz2dc1*dcrv1_dqpt1;
  deriv(IALF,IZ)   = da2dc1*dcrv1_dz1;
  deriv(IALF,IALF) = da2da1;
  deriv(IALF,ITLM) = da2dc1*dcrv1_dtlam1;
  deriv(IALF,IQPT) = da2dc1*dcrv1_dqpt1;
  deriv(ITLM,ITLM) = 1.0;
  deriv(IQPT,IQPT) = 1.0;

  return pstat;

}

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

// return the magnetic field

double PropCyl::get_bfield() const {
  if(_pmap) {
    return 0.;
  }
  else
    return (*_bfield)();
}

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