// PropZXY.cpp

#include "PropZXY.h"
#include <iostream>
#include "trfxyp/PropXYXY.h"
#include "trfbase/PropStat.h"
#include "trfutil/TRFMath.h"
#include "trfxyp/SurfXYPlane.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::asin;
using namespace trf;
#endif

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

namespace {

// Creator.

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

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

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

  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 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
// If pderiv is nonzero, return the derivative matrix there.

PropStat 
vec_transformzxy_( double B, VTrack& trv, double phi_n,
                    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;


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

 double zpos_0 = szp1.get_parameter(izpos);
 
 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_Z);               // 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 << "PropZXY::_vec_propagate: Unknown direction of a track "<< endl;
  abort();
 }

 double sinphi_n = sin(phi_n);
 double cosphi_n = cos(phi_n);
 
 double  u =  a1*cosphi_n + a2*sinphi_n;
 if( u < 0. )
 {
  u = - u;
  sinphi_n = - sinphi_n;
  cosphi_n = - cosphi_n;
  phi_n = phi_n + PI;
  if( phi_n >= TWOPI) phi_n -= TWOPI;
 } 
 // check if du == 0 ( that is track moves parallel to the destination plane )
 double du_dz = a3*cosphi_n + a4*sinphi_n;
 if(du_dz == 0.) return pstat;
 
 double a_hat_phin = 1./du_dz;
 double a_hat_phin2 = a_hat_phin*a_hat_phin;

 // double  u =  a1*cosphi_n + a2*sinphi_n;
 double b1 = -a1*sinphi_n + a2*cosphi_n;
 double b2 = zpos_0;
 double b3 = (a4*cosphi_n - a3*sinphi_n)*a_hat_phin;
 double b4 = a_hat_phin;
 double b5 = a5;
 
 int sign_du;
 if(b4*sign_dz > 0) sign_du =  1;
 if(b4*sign_dz < 0) sign_du = -1;

 vec(IV)     = b1;
 vec(IZ)     = b2;
 vec(IDVDU)  = b3;
 vec(IDZDU)  = b4;
 vec(IQP_XY) = b5;

 // Update trv
 SurfXYPlane sxyp(u,phi_n);
 trv.set_surface(SurfacePtr(sxyp.new_pure_surface()));
 trv.set_vector(vec);

 // set new direction of the track
 if(sign_du ==  1) trv.set_forward();
 if(sign_du == -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;
 
 // check that q/p != 0
 // assert(b5 !=0. );

 double b34_hat = sqrt(1 + b3*b3 + b4*b4);
 double rsinphi =  (b5*B)*sign_du*b34_hat;

 // du_da

 double du_da1 = cosphi_n;
 double du_da2 = sinphi_n;

 // db1_da

 double db1_da1 = - sinphi_n;
 double db1_da2 =   cosphi_n;

 // db3_da

 double db3_da3 = -a4*a_hat_phin2;
 double db3_da4 =  a3*a_hat_phin2;

 // db4_da

 double db4_da3 = - cosphi_n*a_hat_phin2;
 double db4_da4 = - sinphi_n*a_hat_phin2;


 // db3_n_db

 double db3_n_du  = -(1. + b3*b3)*rsinphi;

 // db4_n_db

 double db4_n_du  = -b4*b3*rsinphi;

 // db5_n_db


 // db1_n_da

 double db1_n_da1 = -b3 * du_da1 + db1_da1;
 double db1_n_da2 = -b3 * du_da2 + db1_da2;
 double db1_n_da3 = 0.;
 double db1_n_da4 = 0.;
 double db1_n_da5 = 0.;

 // db2_n_da

 double db2_n_da1 = -b4 * du_da1;
 double db2_n_da2 = -b4 * du_da2;
 double db2_n_da3 = 0.;
 double db2_n_da4 = 0.;
 double db2_n_da5 = 0.;

 // db3_n_da

 double db3_n_da1 = db3_n_du * du_da1;
 double db3_n_da2 = db3_n_du * du_da2;
 double db3_n_da3 = db3_da3;
 double db3_n_da4 = db3_da4;
 double db3_n_da5 = 0.;

 // db4_n_da

 double db4_n_da1 = db4_n_du * du_da1;
 double db4_n_da2 = db4_n_du * du_da2;
 double db4_n_da3 = db4_da3;
 double db4_n_da4 = db4_da4;
 double db4_n_da5 = 0.;

 // db5_n_da

 double db5_n_da1 = 0.;
 double db5_n_da2 = 0.;
 double db5_n_da3 = 0.;
 double db5_n_da4 = 0.;
 double db5_n_da5 = 1.;

 TrackDerivative& deriv = *pderiv;

 deriv(IV,IX)        =db1_n_da1;
 deriv(IV,IY)        =db1_n_da2;
 deriv(IV,IDXDZ)     =db1_n_da3;
 deriv(IV,IDYDZ)     =db1_n_da4;
 deriv(IV,IQP_Z)     =db1_n_da5;
 deriv(IZ,IX)        =db2_n_da1;
 deriv(IZ,IY)        =db2_n_da2;
 deriv(IZ,IDXDZ)     =db2_n_da3;
 deriv(IZ,IDYDZ)     =db2_n_da4;
 deriv(IZ,IQP_Z)     =db2_n_da5;
 deriv(IDVDU,IX)     =db3_n_da1;
 deriv(IDVDU,IY)     =db3_n_da2;
 deriv(IDVDU,IDXDZ)  =db3_n_da3;
 deriv(IDVDU,IDYDZ)  =db3_n_da4;
 deriv(IDVDU,IQP_Z)  =db3_n_da5;
 deriv(IDZDU,IX)     =db4_n_da1;
 deriv(IDZDU,IY)     =db4_n_da2;
 deriv(IDZDU,IDXDZ)  =db4_n_da3;
 deriv(IDZDU,IDYDZ)  =db4_n_da4;
 deriv(IDZDU,IQP_Z)  =db4_n_da5;
 deriv(IQP_XY,IX)    =db5_n_da1;
 deriv(IQP_XY,IY)    =db5_n_da2;
 deriv(IQP_XY,IDXDZ) =db5_n_da3;
 deriv(IQP_XY,IDYDZ) =db5_n_da4;
 deriv(IQP_XY,IQP_Z) =db5_n_da5;

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

// output stream

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

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

// Constructor.

PropZXY::PropZXY(ObjDoublePtr bfield) : _bfield(bfield),
                                        _propxyxy(bfield){
}

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

// Destructor.

PropZXY::~PropZXY() { }

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

// Clone.

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

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

// propagate a track without error in the specified direction
PropStat PropZXY::vec_dir_prop( VTrack& trv, const Surface& srf,
				PropDir dir,TrackDerivative* pder) const {
 PropStat pstat; 
 Propagator::reduce_direction(dir);

 // Check destination is a XYPlane.
 assert( srf.get_pure_type() == SurfXYPlane::get_static_type() );
 if ( srf.get_pure_type( ) != SurfXYPlane::get_static_type() )
   return pstat;
 const SurfXYPlane& sxyp2 = (const SurfXYPlane&) srf;

   // Fetch phi of the destination plane

 int iphi  = SurfXYPlane::NORMPHI;
 double phi_n = sxyp2.get_parameter(iphi);
 VTrack trv0 = trv;
 TrackDerivative deriv1 , deriv2;
 
 double bfac = -BFAC * (*_bfield)();
 if ( pder != 0 ) pstat = vec_transformzxy_( bfac, trv0, phi_n, dir,&deriv1 );
 else pstat = vec_transformzxy_( bfac, trv0, phi_n, dir );

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

 if ( pder != 0 ) pstat = _propxyxy.vec_dir_prop(trv0,srf,dir,&deriv2);
 else  pstat = _propxyxy.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 PropZXY::get_bfield() const {
  return (*_bfield)();
}

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

// Return the creator.

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

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

// Write the object data.

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

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