// PropZZ.cpp

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

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::asin;
#endif

using namespace trf;

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

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

namespace {

// Creator.

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

PropStat _zero_bfield(VTrack& trv,const SurfZPlane& szp1,const SurfZPlane& szp2,Propagator::PropDir dir,TrackDerivative* pderiv);


}

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

// Private function to propagate a track without error
// The corresponding track parameters are:
// z (cm) is fixed 
// 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_propagatezz_( double B, VTrack& trv, const Surface& srf, 
                  const Propagator::PropDir dir1,
                  TrackDerivative* pderiv =0 ) {

 // construct return status
 PropStat pstat;
 Propagator::PropDir dir(dir1);
 bool move = Propagator::reduce_direction(dir);

 // fetch the originating surface and vector
 const SurfacePtr psrf1 = trv.get_surface();
 // TrackVector vec1 = trv.get_vector();

 // Check origin is a ZPlane.
 assert( psrf1->get_pure_type() == SurfZPlane::get_static_type() );
 if ( psrf1->get_pure_type( ) != SurfZPlane::get_static_type() )
   return pstat;
 const SurfZPlane& szp1 = (const SurfZPlane&) *psrf1;

  // Check destination is a ZPlane.
 assert( srf.get_pure_type() == SurfZPlane::get_static_type() );
 if ( srf.get_pure_type( ) != SurfZPlane::get_static_type() )
   return pstat;
 const SurfZPlane& szp2 = (const SurfZPlane&) srf;

  // If surfaces are the same, we can return now.
 if ( srf.pure_equal(*psrf1) ) {
   // if XXX_MOVE was requested , return fail because track doesn't cross
   // Z place in more than one place
   if ( move ) return PropStat();
   if(pderiv)
     {
       TrackDerivative& deriv = *pderiv;
       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;
 }

 if( fabs(B) < 1e-7 ) return _zero_bfield(trv,szp1,szp2,dir1,pderiv);

  // Fetch the zpos's of the planes and the starting track vector.
 int izpos  = SurfZPlane::ZPOS;

 double zpos_0 = szp1.get_parameter(izpos);
 double zpos_n = szp2.get_parameter(izpos);

 // Calculate difference in zpos

 double dz = zpos_n - zpos_0;

 TrackVector vec = trv.get_vector();
 double a1 = vec(IX);                  // x
 double a2 = vec(IY);                  // y
 double a3 = vec(IDXDZ);               // dx/dz
 double a4 = vec(IDYDZ);               // dy/dz
 double a5 = vec(IQP);                 // q/p

 int sign_dz = 0;
 if(trv.is_forward())  sign_dz =  1;
 if(trv.is_backward()) sign_dz = -1;
 if(sign_dz == 0){
  cerr << "PropZZ::_vec_propagate: Unknown direction of a track "<< endl;
  abort();
 }

 // if delta_z*dz/dt >0 and backward - fail
 // if delta_z*dz/dt <0 and forward  - fail

 int flag_forward;
 switch (dir) {
  case Propagator::NEAREST:
   if(sign_dz*dz>0.) flag_forward = 1;
   else flag_forward = -1;

  break;

  case Propagator::FORWARD:

    //   z-z propagation failed 
   if(sign_dz*dz<0.) return pstat;
   flag_forward =  1;

  break;
  case Propagator::BACKWARD:

    //   z-z propagation failed 
   if(sign_dz*dz>0.) return pstat;
   flag_forward = -1;

  break;
  default:
   cerr << "PropZZ::_vec_propagate: Unknown direction." << endl;
   abort();
 }

 double a34_hat = sign_dz*sqrt(a3*a3 + a4*a4);
 double a34_hat2 = a34_hat*a34_hat;
 double dphi = B*dz*a5*sqrt(a34_hat2+1.)*sign_dz;
 double cos_dphi = cos(dphi);
 double sin_dphi = sin(dphi);


 double Rcos_phi = 1./(B*a5)*a3/sqrt(1.+a34_hat2)*sign_dz;
 double Rsin_phi = 1./(B*a5)*a4/sqrt(1.+a34_hat2)*sign_dz; 

 double a1n = a1 + Rsin_phi*(cos_dphi - 1) + Rcos_phi*sin_dphi; 
 double a2n = a2 - Rcos_phi*(cos_dphi - 1) + Rsin_phi*sin_dphi;
 double a3n =   a3*cos_dphi - a4*sin_dphi;
 double a4n =   a3*sin_dphi + a4*cos_dphi;
 double a5n = a5;

 vec(IX)    = a1n;
 vec(IY)    = a2n;
 vec(IDXDZ) = a3n;
 vec(IDYDZ) = a4n;
 vec(IQP)   = a5n;

 // Update trv
 trv.set_surface(SurfacePtr(srf.new_pure_surface()));
 trv.set_vector(vec);

 // set new direction of the track

 if(sign_dz ==  1) trv.set_forward();
 if(sign_dz == -1) trv.set_backward();

 // Calculate the path distance s.
 double ctlamsq = a3*a3 + a4*a4;
 double slam_inv = sign_dz*sqrt(1.0+ctlamsq);
 double ds = dz*slam_inv;
 
 // Set the return status.
 //(flag_forward==1)?pstat.set_forward():pstat.set_backward();
 pstat.set_path_distance(ds);

 // Check the direction of propgation is consistent with the old
 // calculation.
 if ( flag_forward == 1 ) assert( pstat.forward() );
 else assert( pstat.backward() );

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

 double da34_hat_da3;
 double da34_hat_da4;
 if(fabs(a3)>=fabs(a4))
 {
  int sign3=1;
  if(a3<0) sign3 = -1;
  if(a3 !=0.){
   da34_hat_da3 = sign_dz*sign3/sqrt(1.+(a4/a3)*(a4/a3) );
   da34_hat_da4 = sign_dz*(a4/fabs(a3))/sqrt(1.+(a4/a3)*(a4/a3) );
  }
  else {
   da34_hat_da3 = 0.;
   da34_hat_da4 = 0.;
  }
 }
 if(fabs(a4)>fabs(a3))
 {
  int sign4=1;
  if(a4<0) sign4 = -1;
  if(a4 !=0.){
  da34_hat_da3 = sign_dz*(a3/fabs(a4))/sqrt(1.+(a3/a4)*(a3/a4) );
  da34_hat_da4 = sign_dz*sign4/sqrt(1.+(a3/a4)*(a3/a4) );
  }
  else {
   da34_hat_da3 = 0.;
   da34_hat_da4 = 0.;
  }
 }

  // ddphi

 double ddphi_da3 = B*dz*a5*a34_hat*sign_dz/sqrt(1.+a34_hat2)*da34_hat_da3;
 double ddphi_da4 = B*dz*a5*a34_hat*sign_dz/sqrt(1.+a34_hat2)*da34_hat_da4;
 double ddphi_da5 = B*dz*sqrt(a34_hat2+1.)*sign_dz;

  // dRsin_phi


 double dRsin_phi_da3 = -Rsin_phi*a34_hat/(1.+a34_hat2)*da34_hat_da3;
 double dRsin_phi_da4 = -Rsin_phi*a34_hat/(1.+a34_hat2)*da34_hat_da4 +
                            sign_dz/(B*a5)/sqrt(1.+a34_hat2);
 double dRsin_phi_da5 = -Rsin_phi/a5;

  // dRcos_phi


 double dRcos_phi_da3 = -Rcos_phi*a34_hat/(1.+a34_hat2)*da34_hat_da3 +
                            sign_dz/(B*a5)/sqrt(1.+a34_hat2);
 double dRcos_phi_da4 = -Rcos_phi*a34_hat/(1.+a34_hat2)*da34_hat_da4;
 double dRcos_phi_da5 = -Rcos_phi/a5;

  // da1n first two are simple because dR,dphi _da1,_da2 = 0. 

 double da1n_da1 = 1.;
 double da1n_da2 = 0.;
 double da1n_da3 = dRsin_phi_da3*(cos_dphi-1.) + dRcos_phi_da3*sin_dphi 
                 - Rsin_phi*sin_dphi*ddphi_da3 + Rcos_phi*cos_dphi*ddphi_da3;
 double da1n_da4 = dRsin_phi_da4*(cos_dphi-1.) + dRcos_phi_da4*sin_dphi 
                 - Rsin_phi*sin_dphi*ddphi_da4 + Rcos_phi*cos_dphi*ddphi_da4;
 double da1n_da5 = dRsin_phi_da5*(cos_dphi-1.) + dRcos_phi_da5*sin_dphi 
                 - Rsin_phi*sin_dphi*ddphi_da5 + Rcos_phi*cos_dphi*ddphi_da5;

  // da2n first two are simple because dR,dphi _da1,_da2 = 0. 

 double da2n_da1 = 0.;
 double da2n_da2 = 1.;
 double da2n_da3 = -dRcos_phi_da3*(cos_dphi-1.) + dRsin_phi_da3*sin_dphi
                  + Rcos_phi*sin_dphi*ddphi_da3 + Rsin_phi*cos_dphi*ddphi_da3;
 double da2n_da4 = -dRcos_phi_da4*(cos_dphi-1.) + dRsin_phi_da4*sin_dphi
                  + Rcos_phi*sin_dphi*ddphi_da4 + Rsin_phi*cos_dphi*ddphi_da4;
 double da2n_da5 = -dRcos_phi_da5*(cos_dphi-1.) + dRsin_phi_da5*sin_dphi
                  + Rcos_phi*sin_dphi*ddphi_da5 + Rsin_phi*cos_dphi*ddphi_da5;
  
  // da3n first two are simple because dphi _da1,_da2 = 0.

 double da3n_da1 = 0.;
 double da3n_da2 = 0.;
 double da3n_da3 = cos_dphi - a3*sin_dphi*ddphi_da3 - a4*cos_dphi*ddphi_da3;
 double da3n_da4 = - sin_dphi - a3*sin_dphi*ddphi_da4 - a4*cos_dphi*ddphi_da4;
 double da3n_da5 = - a3*sin_dphi*ddphi_da5 - a4*cos_dphi*ddphi_da5;

  // da4n first two are simple because dphi _da1,_da2 = 0.

 double da4n_da1 = 0.;
 double da4n_da2 = 0.; 
 double da4n_da3 =   sin_dphi + a3*cos_dphi*ddphi_da3 - a4*sin_dphi*ddphi_da3;
 double da4n_da4 =   cos_dphi + a3*cos_dphi*ddphi_da4 - a4*sin_dphi*ddphi_da4;
 double da4n_da5 =   a3*cos_dphi*ddphi_da5 - a4*sin_dphi*ddphi_da5;

  // da5n 
  
 double da5n_da1 = 0.;
 double da5n_da2 = 0.;
 double da5n_da3 = 0.; 
 double da5n_da4 = 0.;
 double da5n_da5 = 1.;

 TrackDerivative& deriv = *pderiv;

 deriv(IX,IX)       =da1n_da1;
 deriv(IX,IY)       =da1n_da2;
 deriv(IX,IDXDZ)    =da1n_da3;
 deriv(IX,IDYDZ)    =da1n_da4;
 deriv(IX,IQP)      =da1n_da5;
 deriv(IY,IX)       =da2n_da1;
 deriv(IY,IY)       =da2n_da2;
 deriv(IY,IDXDZ)    =da2n_da3;
 deriv(IY,IDYDZ)    =da2n_da4;
 deriv(IY,IQP)      =da2n_da5;
 deriv(IDXDZ,IX)    =da3n_da1;
 deriv(IDXDZ,IY)    =da3n_da2;
 deriv(IDXDZ,IDXDZ) =da3n_da3;
 deriv(IDXDZ,IDYDZ) =da3n_da4;
 deriv(IDXDZ,IQP)   =da3n_da5;
 deriv(IDYDZ,IX)    =da4n_da1;
 deriv(IDYDZ,IY)    =da4n_da2;
 deriv(IDYDZ,IDXDZ) =da4n_da3;
 deriv(IDYDZ,IDYDZ) =da4n_da4;
 deriv(IDYDZ,IQP)   =da4n_da5;
 deriv(IQP,IX)      =da5n_da1;
 deriv(IQP,IY)      =da5n_da2;
 deriv(IQP,IDXDZ)   =da5n_da3;
 deriv(IQP,IDYDZ)   =da5n_da4;
 deriv(IQP,IQP)     =da5n_da5;

 return pstat;

}

namespace {

PropStat _zero_bfield( VTrack& trv,const SurfZPlane& szp1,
                       const SurfZPlane& szp2,Propagator::PropDir dir1,
                      TrackDerivative* pderiv) {
  PropStat pstat;
  Propagator::PropDir dir(dir1);
  bool move = Propagator::reduce_direction(dir);
  bool same = szp2.pure_equal(szp1);

  // There is only one solution. Can't XXX_MOVE
  if ( same && move ) return pstat;

  if ( same ) {
    if(pderiv)
      {
	TrackDerivative& deriv = *pderiv;
	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;
  }

  TrackVector vec = trv.get_vector();
  double x0 = vec(IX);
  double y0 = vec(IY);
  double dxdz0 = vec(IDXDZ);
  double dydz0 = vec(IDYDZ);

  double dz0 =1.;
  if( trv.is_backward() ) dz0 = -1.;

  double z0 = szp1.get_parameter(SurfZPlane::ZPOS);

  double z1 = szp2.get_parameter(SurfZPlane::ZPOS);

  double a = dxdz0*dz0;
  double b = dydz0*dz0;
  double c = dz0;

  double ap = 0.;
  double bp = 0.;
  double cp = 1.0;

  double xp = 0.;
  double yp = 0.;
  double zp = z1;

  double denom = a*ap + b*bp + c*cp;

  if( denom == 0. ) return pstat;

  double S = ( (xp-x0)*ap + (yp-y0)*bp + (zp-z0)*cp )/denom;

  double x1 = x0 + S*a;
  double y1 = y0 + S*b;

  bool forward = S > 0. ? true : false;

  if( dir == Propagator::FORWARD && !forward ) return pstat;
  if( dir == Propagator::BACKWARD && forward ) return pstat;


  double dxdz1 = (x1-x0)/(z1-z0);
  double dydz1 = (y1-y0)/(z1-z0);

  vec(IX) = x1;
  vec(IY) = y1;
  vec(IDXDZ) = dxdz1;
  vec(IDYDZ) = dydz1;
 
  trv.set_surface(SurfacePtr(szp2.new_pure_surface()));
  trv.set_vector(vec);
  if( dz0 >0 ) trv.set_forward();
  else trv.set_backward();

  double ds = S*sqrt(a*a+b*b+c*c);

  pstat.set_path_distance(ds);

  if( pderiv == 0 ) return pstat;


  // S= (zp-z0)/dz0
  // a(dxdz dz0) b(dydz dz0) c(dz0)
 

  double da_dxdz0 = dz0;
  double db_dydz0 = dz0;

  double dx1dx0 =1.;
  double dy1dy0 =1.;

  double dx1_dxdz0 = S*da_dxdz0;
  double dy1_dydz0 = S*db_dydz0;

  double dxdz1_dx0 = (dx1dx0 - 1.)/(z1-z0);
  double dydz1_dy0 = (dy1dy0 - 1.)/(z1-z0);

  double dxdz1_dxdz0 = dx1_dxdz0/(z1-z0);
  double dydz1_dydz0 = dy1_dydz0/(z1-z0);

  TrackDerivative& deriv = *pderiv;
  deriv.fill( 0.0 );
  
  deriv(IX,IX) = dx1dx0;
  deriv(IX,IDXDZ) = dx1_dxdz0;

  deriv(IY,IY) = dy1dy0;
  deriv(IY,IDYDZ) = dy1_dydz0;

  deriv(IDXDZ,IX) = dxdz1_dx0;
  deriv(IDXDZ,IDXDZ) = dxdz1_dxdz0;

  deriv(IDYDZ,IY) = dydz1_dy0;
  deriv(IDYDZ,IDYDZ) = dydz1_dydz0;

  deriv(IQP,IQP) = 1.0;

  return pstat;

 }
}

//************************************************************************
// PropZZ member functions.
//************************************************************************

// output stream

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

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

// Constructor.

PropZZ::PropZZ(ObjDoublePtr bfield) : _bfield(bfield) {

}

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

// Destructor.

PropZZ::~PropZZ() { }

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

// Clone.

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

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

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


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

// return the magnetic field

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

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

// Return the creator.

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

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

// Write the object data.

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

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