// PropZXY.cpp

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

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() == "PropXYZ" );
  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 << "PropXYZ: can't get bfield parameter." << endl;
    abort();
  }
  return ObjPtr( new PropXYZ(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_transformxyz_( 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 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;


 // Fetch the dist and phi of the plane and the starting track vector.
 int iphi  = SurfXYPlane::NORMPHI;
 int idist = SurfXYPlane::DISTNORM;
 double sphi = sxyp1.get_parameter(iphi);
 double    u = sxyp1.get_parameter(idist);
 TrackVector vec = trv.get_vector();
 double b1 = vec(IV);                  // v
 double b2 = vec(IZ);                  // z
 double b3 = vec(IDVDU);               // dv/du
 double b4 = vec(IDZDU);               // dz/du
 double b5 = vec(IQP_XY);              // q/p

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

 double zpos = b2;
 double cos_sphi = cos(sphi);
 double sin_sphi = sin(sphi);

 double a1 = u*cos_sphi - b1*sin_sphi;
 double a2 = b1*cos_sphi + u*sin_sphi; 
 double a3 = cos_sphi/b4 - b3/b4*sin_sphi;
 double a4 = b3/b4*cos_sphi + sin_sphi/b4;
 double a5 = b5;


 // if dz/du*du/dt > 0 and forward : dz>0; dz/du*du/dt < 0 dz<0
 // if dz/du*du/dt < 0 and backward : dz>0; dz/du*du/dt > 0 dz<0  

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

 int sign_dz;
 if(b4*sign_du > 0) sign_dz =  1;
 if(b4*sign_du < 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 / db
  
 double ddphi_db2 = -B*a5*sqrt(a34_hat2+1.)*sign_dz;
 double ddphi_db2_nob = -sqrt(a34_hat2+1.)*sign_dz;

  // da3 / db

 double da3_db3 = -sin_sphi/b4; 
 double da3_db4 = (-cos_sphi + b3*sin_sphi)/(b4*b4);

  // da4 / db

 double da4_db3 =  cos_sphi/b4;
 double da4_db4 = (- sin_sphi - b3*cos_sphi)/(b4*b4); 
 
  // da1/ db

 double da1n_db1 = -sin_sphi; 
 double da1n_db2 = Rcos_phi*ddphi_db2_nob;
 double da1n_db3 = 0.;
 double da1n_db4 = 0.;
 double da1n_db5 = 0.;

  // da2/ db

 double da2n_db1 = cos_sphi; 
 double da2n_db2 = Rsin_phi*ddphi_db2_nob;
 double da2n_db3 = 0.;
 double da2n_db4 = 0.;
 double da2n_db5 = 0.;

  // da3/ db

 double da3n_db1 = 0.; 
 double da3n_db2 = -a4*ddphi_db2;
 double da3n_db3 = da3_db3;
 double da3n_db4 = da3_db4;
 double da3n_db5 = 0.;
 
  // da4/ db

 double da4n_db1 = 0.; 
 double da4n_db2 = a3*ddphi_db2;
 double da4n_db3 = da4_db3;
 double da4n_db4 = da4_db4;
 double da4n_db5 = 0.;

  // da5n 
  
 double da5n_db1 = 0.;
 double da5n_db2 = 0.;
 double da5n_db3 = 0.; 
 double da5n_db4 = 0.;
 double da5n_db5 = 1.;

 TrackDerivative& deriv = *pderiv;

 deriv(IX,IV)        =da1n_db1;
 deriv(IX,IZ)        =da1n_db2;
 deriv(IX,IDVDU)     =da1n_db3;
 deriv(IX,IDZDU)     =da1n_db4;
 deriv(IX,IQP_XY)    =da1n_db5;
 deriv(IY,IV)        =da2n_db1;
 deriv(IY,IZ)        =da2n_db2;
 deriv(IY,IDVDU)     =da2n_db3;
 deriv(IY,IDZDU)     =da2n_db4;
 deriv(IY,IQP_XY)    =da2n_db5;
 deriv(IDXDZ,IV)     =da3n_db1;
 deriv(IDXDZ,IZ)     =da3n_db2;
 deriv(IDXDZ,IDVDU)  =da3n_db3;
 deriv(IDXDZ,IDZDU)  =da3n_db4;
 deriv(IDXDZ,IQP_XY) =da3n_db5;
 deriv(IDYDZ,IV)     =da4n_db1;
 deriv(IDYDZ,IZ)     =da4n_db2;
 deriv(IDYDZ,IDVDU)  =da4n_db3;
 deriv(IDYDZ,IDZDU)  =da4n_db4;
 deriv(IDYDZ,IQP_XY) =da4n_db5;
 deriv(IQP_Z,IV)     =da5n_db1;
 deriv(IQP_Z,IZ)     =da5n_db2;
 deriv(IQP_Z,IDVDU)  =da5n_db3;
 deriv(IQP_Z,IDZDU)  =da5n_db4;
 deriv(IQP_Z,IQP_XY) =da5n_db5;

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

// output stream

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

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

// Constructor.

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

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

// Destructor.

PropXYZ::~PropXYZ() { }

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

// Clone.

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

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

// propagate a track without error in the specified direction
PropStat PropXYZ::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;

   // Fetch phi of the destination plane
 
 VTrack trv0 = trv;
 TrackDerivative deriv1 , deriv2;

 double bfac = -BFAC * (*_bfield)();
 if ( pder != 0 )  pstat = vec_transformxyz_( bfac, trv0, dir,&deriv1 );
 else pstat = vec_transformxyz_( 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 PropXYZ::get_bfield() const {
  return (*_bfield)();
}

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

// Return the creator.

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

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

// Write the object data.

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

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