// PropCylZ.cpp

#include <iostream>
#include "PropCylZ.h"
#include "trfzp/PropZZ.h"
#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>

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

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

namespace {

// Creator.

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

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

// Assign track parameter indices.
enum {
  IPHI   = SurfCylinder::IPHI, 
  IZ_CYL = SurfCylinder::IZ, 
  IALF   = SurfCylinder::IALF, 
  ITLM   = SurfCylinder::ITLM, 
  IQPT   = SurfCylinder::IQPT,

  IX = SurfZPlane::IX,
  IY   = SurfZPlane::IY,
  IDXDZ = SurfZPlane::IDXDZ,
  IDYDZ = SurfZPlane::IDYDZ,
  IQP_Z  = SurfZPlane::IQP
};

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

// Private function to propagate a track without error
// The corresponding track parameters are:
// On ZPlane:
// 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
// On Cylinder:
// r (cm) is fixed
// 0 - phi
// 1 - z (cm)
// 2 - alpha = phi_dir - phi; tan(alpha) = r*dphi/dr
// 3 - sin(lambda) = dz/ds; tan(lambda) = dz/dsT
// 4 - q/pT (1/GeV/c) (pT is component of p parallel to cylinder)
// If pderiv is nonzero, return the derivative matrix there.

PropStat 
vec_transformcylz_( double B, VTrack& trv,
                    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 Cylinder.
 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;


 // Fetch the R of the cylinder and the starting track vector.
 int ir  = SurfCylinder::RADIUS;
 double Rcyl = scy1.get_parameter(ir);
 
 TrackVector vec = trv.get_vector();
 double c1 = vec(IPHI);                 // phi
 double c2 = vec(IZ_CYL);               // z
 double c3 = vec(IALF);                 // alpha
 double c4 = vec(ITLM);                 // tan(lambda)
 double c5 = vec(IQPT);                 // q/pt

 // check that dz != 0
 if(c4 == 0.) return pstat;

 double zpos = c2;
 double cos_c1 = cos(c1);
 double sin_c1 = sin(c1);
 double cos_dir = cos(c1+c3);
 double sin_dir = sin(c1+c3);
 double c4_hat2 = 1+c4*c4;
 double c4_hat  = sqrt(c4_hat2);

 double a1 = Rcyl*cos_c1;
 double a2 = Rcyl*sin_c1;
 double a3 = cos_dir/c4;
 double a4 = sin_dir/c4;
 double a5 = c5/c4_hat;


 // if tan(lambda)=dz/dst > 0 : dz>0; dzdst < 0 dz<0

 int sign_dz;
 if(c4 > 0) sign_dz =  1;
 if(c4 < 0) sign_dz = -1;

 vec(IX)     = a1;
 vec(IY)     = a2;
 vec(IDXDZ)  = a3;
 vec(IDYDZ)  = a4;
 vec(IQP_Z)  = a5;

 // Update trv
 SurfZPlane szp(zpos);
 trv.set_surface(SurfacePtr(szp.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();

 // Set the return status.
 pstat.set_same();

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

 double a34_hat = sign_dz*sqrt(a3*a3 + a4*a4);
 double a34_hat2 = a34_hat*a34_hat;

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

  // ddphi / dc
  
 double ddphi_dc2 = -B*a5*sqrt(a34_hat2+1.)*sign_dz;
 double ddphi_dc2_nob = -sqrt(a34_hat2+1.)*sign_dz;

  // da1 / dc

 double da1_dc1 = -Rcyl*sin_c1; 

  // da2 / dc

 double da2_dc1 =  Rcyl*cos_c1; 

  // da3 / dc

 double da3_dc1 = -sin_dir/c4; 
 double da3_dc3 = -sin_dir/c4;
 double da3_dc4 = -cos_dir/(c4*c4);

  // da4 / dc

 double da4_dc1 =  cos_dir/c4;
 double da4_dc3 =  cos_dir/c4;
 double da4_dc4 = -sin_dir/(c4*c4); 

  // da5 / dc

 double da5_dc4 =  -c5*c4/(c4_hat*c4_hat2);
 double da5_dc5 =   1./c4_hat;

  // da1n/ dc

 double da1n_dc1 = da1_dc1; 
 double da1n_dc2 = Rcos_phi*ddphi_dc2_nob;
 double da1n_dc3 = 0.;
 double da1n_dc4 = 0.;
 double da1n_dc5 = 0.;

  // da2n/ dc

 double da2n_dc1 = da2_dc1; 
 double da2n_dc2 = Rsin_phi*ddphi_dc2_nob;
 double da2n_dc3 = 0.;
 double da2n_dc4 = 0.;
 double da2n_dc5 = 0.;

  // da3n/ dc

 double da3n_dc1 = da3_dc1; 
 double da3n_dc2 = - a4*ddphi_dc2;
 double da3n_dc3 = da3_dc3;
 double da3n_dc4 = da3_dc4;
 double da3n_dc5 = 0.;
 
  // da4n/ dc

 double da4n_dc1 = da4_dc1; 
 double da4n_dc2 = a3*ddphi_dc2;
 double da4n_dc3 = da4_dc3;
 double da4n_dc4 = da4_dc4;
 double da4n_dc5 = 0.;

  // da5n 
  
 double da5n_dc1 = 0.;
 double da5n_dc2 = 0.;
 double da5n_dc3 = 0.; 
 double da5n_dc4 = da5_dc4;
 double da5n_dc5 = da5_dc5;

 TrackDerivative& deriv = *pderiv;

 deriv(IX,IPHI)      =da1n_dc1;
 deriv(IX,IZ_CYL)    =da1n_dc2;
 deriv(IX,IALF)      =da1n_dc3;
 deriv(IX,ITLM)      =da1n_dc4;
 deriv(IX,IQPT)      =da1n_dc5;
 deriv(IY,IPHI)      =da2n_dc1;
 deriv(IY,IZ_CYL)    =da2n_dc2;
 deriv(IY,IALF)      =da2n_dc3;
 deriv(IY,ITLM)      =da2n_dc4;
 deriv(IY,IQPT)      =da2n_dc5;
 deriv(IDXDZ,IPHI)   =da3n_dc1;
 deriv(IDXDZ,IZ_CYL) =da3n_dc2;
 deriv(IDXDZ,IALF)   =da3n_dc3;
 deriv(IDXDZ,ITLM)   =da3n_dc4;
 deriv(IDXDZ,IQPT)   =da3n_dc5;
 deriv(IDYDZ,IPHI)   =da4n_dc1;
 deriv(IDYDZ,IZ_CYL) =da4n_dc2;
 deriv(IDYDZ,IALF)   =da4n_dc3;
 deriv(IDYDZ,ITLM)   =da4n_dc4;
 deriv(IDYDZ,IQPT)   =da4n_dc5;
 deriv(IQP_Z,IPHI)   =da5n_dc1;
 deriv(IQP_Z,IZ_CYL) =da5n_dc2;
 deriv(IQP_Z,IALF)   =da5n_dc3;
 deriv(IQP_Z,ITLM)   =da5n_dc4;
 deriv(IQP_Z,IQPT)   =da5n_dc5;
 /*
 if( fabs(B) < 1e-11 ) {
   deriv(IQP_Z,IPHI)   =0.0;
   deriv(IQP_Z,IZ_CYL) =0.0;
   deriv(IQP_Z,IALF)   =0.0;
   deriv(IQP_Z,ITLM)   =0.0;
   deriv(IQP_Z,IQPT)   =1.0;
 }
 */

 return pstat;
}
//************************************************************************
// PropCylZ member functions.
//************************************************************************

// output stream

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

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

// Constructor.

PropCylZ::PropCylZ(ObjDoublePtr bfield) : _bfield(bfield),
                                          _propzz(bfield){
}

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

// Destructor.

PropCylZ::~PropCylZ() { }

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

// Clone.

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

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

// propagate a track without error in the specified direction
PropStat PropCylZ::vec_dir_prop( VTrack& trv, const Surface& srf,
				 PropDir dir,TrackDerivative* pder ) const {
 PropStat pstat;
 Propagator::reduce_direction(dir);
 // 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;

 VTrack trv0 = trv;
 TrackDerivative deriv1 , deriv2;

 double bfac = -BFAC * (*_bfield)();
 if ( pder != 0 ) pstat = vec_transformcylz_( bfac, trv0, dir, &deriv1 );
 else pstat = vec_transformcylz_( bfac, trv0, dir );

 if ( ! pstat.success() ) return pstat; 

 if ( pder != 0 ) pstat = _propzz.vec_dir_prop(trv0,srf,dir,&deriv2);
 else pstat = _propzz.vec_dir_prop(trv0,srf,dir);

 if ( pstat.success() ) 
   {
     trv = trv0;
     if ( pder != 0 )  *pder = TrackDerivative(deriv2*deriv1);
   }

 return pstat;
}

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

// return the magnetic field

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

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

// Return the creator.

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

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

// Write the object data.

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

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