/*
 *   ---- Delphes ----
 *     A Fast Simulator for general purpose LHC detector
 *       S. Ovyn ~~~~ severine.ovyn@uclouvain.be
 *
 *     Center for Particle Physics and Phenomenology (CP3)
 *     Universite Catholique de Louvain (UCL)
 *     Louvain-la-Neuve, Belgium
 *             */

#include "BFieldProp.h"
#include<cmath>
using namespace std;


//------------------------------------------------------------------------------

TrackPropagation::TrackPropagation(){
   DET = new RESOLution();
   init();
}

TrackPropagation::TrackPropagation(const string& DetDatacard){
   DET = new RESOLution();
   DET->ReadDataCard(DetDatacard);
   init();	
}

TrackPropagation::TrackPropagation(const RESOLution* DetDatacard){ 
   DET= new RESOLution(*DetDatacard);
   init();
}

TrackPropagation::TrackPropagation(const TrackPropagation & tp){
   MAXITERATION = tp.MAXITERATION;
   DET   = new RESOLution(*(tp.DET));
   R_max = tp.R_max; 	z_max = tp.z_max;
   B_x   = tp.B_x;	B_y   = tp.B_y;		B_z   = tp.B_z;
   q     = tp.q;	phi_0 = tp.phi_0;	
   gammam= tp.gammam;	omega = tp.omega;
   r     = tp.r;	rr    = tp.rr;
   x_c   = tp.x_c;	y_c   = tp.y_c; 
   R_c   = tp.R_c; 	Phi_c = tp.Phi_c;
   t     = tp.t; 	t_z   = tp.t_z; 	t_T   = tp.t_T; 
   x_t   = tp.x_t;	y_t   = tp.y_t;		z_t   = tp.z_t;
   R_t   = tp.R_t;      Phi_t = tp.Phi_t;       
   Theta_t=tp.Theta_t;  Eta_t = tp.Eta_t;
   Px_t  = tp.Px_t;     Py_t  = tp.Py_t;	Pz_t  = tp.Pz_t;
   PT_t  = tp.PT_t;     p_t   = tp.p_t; 	E_t   = tp.E_t;
   loop_overflow_counter = tp.loop_overflow_counter;
}

TrackPropagation& TrackPropagation::operator=(const TrackPropagation & tp) {
   if(this==&tp) return *this;
   MAXITERATION = tp.MAXITERATION;
   DET   = new RESOLution(*(tp.DET));
   R_max = tp.R_max;    z_max = tp.z_max;
   B_x   = tp.B_x;      B_y   = tp.B_y;         B_z   = tp.B_z;
   q     = tp.q;        phi_0 = tp.phi_0;
   gammam= tp.gammam;   omega = tp.omega;
   r     = tp.r;        rr    = tp.rr;
   x_c   = tp.x_c;      y_c   = tp.y_c;
   R_c   = tp.R_c;      Phi_c = tp.Phi_c;
   t     = tp.t;        t_z   = tp.t_z;         t_T   = tp.t_T;
   x_t   = tp.x_t;      y_t   = tp.y_t;         z_t   = tp.z_t;
   R_t   = tp.R_t;      Phi_t = tp.Phi_t;
   Theta_t=tp.Theta_t;  Eta_t = tp.Eta_t;
   Px_t  = tp.Px_t;     Py_t  = tp.Py_t;        Pz_t  = tp.Pz_t;
   PT_t  = tp.PT_t;     p_t   = tp.p_t;         E_t   = tp.E_t;
   loop_overflow_counter = tp.loop_overflow_counter;
   return *this;
}



void TrackPropagation::init() {
 MAXITERATION = 10000; 
 q= UNDEFINED;  phi_0= UNDEFINED; gammam= UNDEFINED; omega=UNDEFINED; r=UNDEFINED;
 x_c=UNDEFINED; y_c=UNDEFINED;    R_c=UNDEFINED;     Phi_c=UNDEFINED;
 rr=UNDEFINED;  t=UNDEFINED;      t_z=UNDEFINED;     t_T=UNDEFINED;
 x_t=UNDEFINED; y_t=UNDEFINED;    z_t=UNDEFINED;
 R_t=UNDEFINED; Phi_t=UNDEFINED;  Theta_t=UNDEFINED; Eta_t=UNDEFINED;
 Px_t=UNDEFINED; Py_t=UNDEFINED;  Pz_t=UNDEFINED;    PT_t=UNDEFINED; p_t=UNDEFINED; E_t=UNDEFINED;

  // DET has been initialised in the constructors
  // magnetic field parameters
   R_max = DET->TRACK_radius;
   z_max = DET->TRACK_length/2.;
   B_x = DET->TRACK_bfield_x;
   B_y = DET->TRACK_bfield_y;
   B_z = DET->TRACK_bfield_z;

   loop_overflow_counter=0;
}



void TrackPropagation::Propagation(const TRootGenParticle *Part,TLorentzVector &momentum) {

  q  = Charge(Part->PID);
  if(q==0) return;

  if(R_max ==0) { cout << "ERROR: magnetic field has no lateral extention\n"; return;}
  if(z_max==0) { cout << "ERROR: magnetic field has no longitudinal extention\n"; return;}

  if (B_x== 0 && B_y== 0) { // faster if only B_z
    if (B_z==0) return; // nothing to do

    // initial conditions:
      // p_X0 = Part->Px, p_Y0 = Part->Py, p_Z0 = Part->Pz, p_T0 = Part->PT;
      // X_0 = Part->X, Y_0 = Part->Y, Z_0 = Part->Z;

    // 1.  initial transverse momentum p_{T0} : Part->PT
    //     initial transverse momentum direction \phi_0 = -atan(p_X0/p_Y0) 
    //     relativistic gamma : gamma = E/mc² ; gammam = gamma \times m
    //     giration frequency \omega = q/(gamma m) B_z
    //     helix radius r = p_T0 / (omega gamma m)
      phi_0 = -atan2(Part->Px,Part->Py);
      gammam = Part->E; // here c==1
      //cout << "gammam" << gammam << "\t gamma" << gammam/Part->M << endl;
      omega = q * B_z /gammam;
      r = Part->PT / (omega * gammam);

    // 2.  Helix parameters : center coordinates in transverse plane
    //     x_c = x_0 - r*cos(phi_0) and y_c = y_0 - r*sin(phi_0)
    //     R_c = \sqrt{x_c² + y_c²} and \Phi_c = atan{y_c/x_c}
      x_c = Part->X - r*cos(phi_0);  /// TEST !!
      y_c = Part->Y - r*sin(phi_0); 
      R_c = sqrt(pow(x_c,2.) + pow(y_c,2.) );
      Phi_c = atan2(y_c,x_c);

    // 3. time evaluation t = min(t_T, t_z)
    //    t_T : time to exit from the sides
    //    t_T= [ Phi_c - phi_0 + atan( (R_max^2 - (R_c^2 + r^2))/(2rR_c) )  ]/omega
    //    t_z : time to exit from the front or the back
    //    t_z = gamma * m /p_z0 \times (-z_0 + z_max * sign(p_z0))
      rr = sqrt( pow(R_c,2.) + pow(r,2.) ); // temp variable
      t_T=0;
      int sign_pz= (Part->Pz >0) ? 1 : -1;
      t_z = gammam / Part->Pz * (-Part->Z + z_max*sign_pz ) ;
      if ( fabs(R_c - r) > R_max || R_c + r  < R_max ) t = t_z;
      else {
         t_T = (Phi_c - phi_0 + atan2( (R_max + rr)*(R_max - rr) , 2*r*R_c ) ) / omega; 
         t = min(t_T,t_z);
      }

    // 4. position in terms of x(t), y(t), z(t)
    //    x(t) = x_c + r cos (omega t + phi_0)
    //    y(t) = y_c + r sin (omega t + phi_0)
    //    z(t) = z_0 + (p_Z0/gammam) t
      x_t = x_c + r * cos(omega * t + phi_0); 
      y_t = y_c + r * sin(omega * t + phi_0); 
      z_t = Part->Z + Part->Pz / gammam * t;

    // 5. position in terms of Theta(t), Phi(t), R(t), Eta(t)
    //    R(t) = sqrt(x(t)² + y(t)²)
    //    Phi(t) = atan(y(t)/x(t))
    //    Theta(t) = atan(R(t)/z(t))
    //    Eta(t) = -ln tan (Theta(t)/2)
      R_t   = sqrt( pow(x_t,2.) + pow(y_t,2.)  );
      Phi_t = atan2( y_t, x_t);
      if(R_t>0) {
	      Theta_t = acos( z_t / sqrt(z_t*z_t+ R_t*R_t));
	      Eta_t = - log(tan(Theta_t/2.));
      } else{
		Theta_t=0; Eta_t = 9999;
      }

	Px_t = - Part->PT * sin(omega*t + phi_0);
	Py_t =   Part->PT * cos(omega*t + phi_0);
        Pz_t =   Part->Pz;
        PT_t =   sqrt(Px_t*Px_t + Py_t*Py_t);
	p_t = sqrt(PT_t*PT_t + Pz_t*Pz_t);
	E_t=sqrt(Part->M*Part->M +p_t);
	//if(p_t != fabs(Pz_t) ) Eta_t = log( (p_t+Pz_t)/(p_t-Pz_t) )/2.;
	//if(p_t>0) Theta_t = acos(Pz_t/p_t);
        momentum.SetPxPyPzE(Px_t,Py_t,Pz_t,E_t);

// test zone ---
/*
	cout << cos(atan(R_t/z_t)) << "\t" << cos(Theta_t) << "\t" << cos(momentum.Theta()) << "\t" << Pz_t/temp_p << endl;
	double Eta_t1 = log( (E+Pz_t)/(E-Pz_t) )/2.;
	double Eta_t2 = log( (temp_p+Pz_t)/(temp_p-Pz_t) )/2.;
	if(0 &&  fabs(Eta_t -Eta_t2)>1e-310) {
		cout << "ERROR-BUG: Eta_t != Eta_t2\n";
		cout << "Eta_t= " << Eta_t << "\t Eta_t1= " << Eta_t1 << "\t Eta_t2= " << Eta_t2 << endl;
	}

        double R_t2  = sqrt( pow(R_c,2.) + pow(r,2.) + 2*r*R_c*cos(phi_0 + omega*t - Phi_c) ); // cross-check
        if(fabs(R_t - R_t2) > 1e-7) 
		cout << "ERROR-BUG: R_t != R_t2:  R_t=" << R_t << " R_t2=" << R_t2 << " R_t - R_t2 =" << R_t - R_t2 << endl;
	if( fabs(E - gammam) > 1e-3 ) {
		cout << "ERROR-BUG: energy is not conserved in src/BFieldProp.cc\n";
		cout << "E - momentum.E() = " <<  fabs(E - momentum.E()) <<  " gammam - E " << fabs(gammam -E) << endl; }
	if( fabs(PT_t - Part->PT) > 1e-10 ) {
		cout << "ERROR-BUG: PT is not conversed in src/BFieldProp.cc. ";
		cout << "(at " << 100*(PT_t - Part->PT) << "%)\n";
		}
	if(momentum.Pz() != Pz_t)
		cout << "ERROR-BUG: Pz is not conserved in src/BFieldProp.cc\n";

	double temp_p0=sqrt(Part->PT*Part->PT + Part->Pz*Part->Pz);
	if(fabs( (temp_p-temp_p0)*(temp_p+temp_p0) )>1e-10  ) {
		cout << "ERROR-BUG: momentum |vec{p}| is not conserved in src/BFieldProp.cc\n";
		cout << temp_p << "\t" << temp_p0 << endl;
	}

   // if x_c == y_c ==0 (set it by hand!), easy cross-check
   //cout << "tan(phi_p)= " << momentum.Py()/momentum.Px() << "\t -1/tan(phi_x)= " << -x_t/y_t << endl;
*/

  } else { // if B_x or B_y are non zero: longer computation

  float Xvertex1 = Part->X;
  float Yvertex1 = Part->Y;
  float Zvertex1 = Part->Z;
  
  //out of tracking coverage?
  if(sqrt(Xvertex1*Xvertex1+Yvertex1*Yvertex1) > R_max){return;}
  if(fabs(Zvertex1) > z_max){return;}
  
  double px = Part->Px / 0.003;
  double py = Part->Py / 0.003;
  double pz = Part->Pz / 0.003;
  double pt = Part->PT / 0.003; // sqrt(px*px+py*py);
  double p  = sqrt(pz*pz + pt*pt); //sqrt(px*px+py*py+pz*pz);
 
  double M  = Part->M;
  double vx = px/M;
  double vy = py/M;
  double vz = pz/M;
  double qm = q/M;

  double ax =  qm*(B_z*vy - B_y*vz);
  double ay =  qm*(B_x*vz - B_z*vx);
  double az =  qm*(B_y*vx - B_x*vy);
  double dt = 1/p;
  if(pt<266 && vz < 0.0012)       dt = fabs(0.001/vz); // ?????
 
  double xold=Xvertex1;         double x=xold;
  double yold=Yvertex1;         double y=yold;
  double zold=Zvertex1;         double z=zold;

  double VTold = pt/M; //=sqrt(vx*vx+vy*vy);
 
  unsigned int k = 0;
  double VTratio=0;
  double R_max2 = R_max*R_max;
  double r2=0; // will be x*x+y*y

  while(k < MAXITERATION){
        k++;

        vx += ax*dt;
        vy += ay*dt;
        vz += az*dt;

        VTratio = VTold/sqrt(vx*vx+vy*vy);
        vx *= VTratio;
        vy *= VTratio;

        ax = qm*(B_z*vy - B_y*vz);
        ay = qm*(B_x*vz - B_z*vx);
        az = qm*(B_y*vx - B_x*vy);

        x  += vx*dt;
        y  += vy*dt;
        z  += vz*dt;
	r2  = x*x + y*y;

       if( r2 > R_max2 ){ 
		x /= r2/R_max2; 
		y /= r2/R_max2; 
		break;
       }
       if( fabs(z)>z_max)break;

       xold = x;
       yold = y;
       zold = z;
  } // while loop

  if(k == MAXITERATION) loop_overflow_counter++;
  //cout << "too short loop in " << loop_overflow_counter << " cases" << endl;
 
  if(x!=0 && y!=0 && z!=0) {
          float Theta = atan2(sqrt(r2),z);
          double eta  = -log(tan(Theta/2.));
	  double phi = atan2(y,x);
	  momentum.SetPtEtaPhiE(Part->PT,eta,phi,Part->E);
  }
    
 } // if b_x or b_y non zero
}