/***********************************************************************
**                                                                    **
**          /----------------------------------------------\          **
**         |  Delphes, a framework for the fast simulation  |         **
**         |       of a  generic collider experiment        |         **
**          \----------------------------------------------/          **
**                                                                    **
**                                                                    **
**   This package uses:                                               **
**   ------------------                                               **
**   FastJet algorithm: Phys. Lett. B641 (2006) [hep-ph/0512210]      ** 
**   Hector: JINST 2:P09005 (2007) [physics.acc-ph:0707.1198v2]       ** 
**   FROG: [hep-ex/0901.2718v1]                                       **
**                                                                    **
** ------------------------------------------------------------------ **
**                                                                    **
**   Main authors:                                                    **
**   -------------                                                    **
**                                                                    **
**                Severine Ovyn                Xavier Rouby           ** 
**          severine.ovyn@uclouvain.be      xavier.rouby@cern         ** 
**                                                                    ** 
**         Center for Particle Physics and Phenomenology (CP3)        ** 
**                Universite catholique de Louvain (UCL)              **       
**                     Louvain-la-Neuve, Belgium                      **            
**                                                                    **
**                      Copyright (C) 2008-2009,                      **
**                        All rights reserved.                        **  
**                                                                    **
***********************************************************************/

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


//------------------------------------------------------------------------------
extern const float UNDEFINED;

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/100.; //[m]
   z_max = DET->TRACK_length/100.; //[m]
   B_x = DET->TRACK_bfield_x*tesla;
   B_y = DET->TRACK_bfield_y*tesla;
   B_z = DET->TRACK_bfield_z;

   loop_overflow_counter=0;
}






void TrackPropagation::bfield(TRootGenParticle *Part) {


  // initialisation, valid for z_max==0, R_max==0 and q==0
  Part->EtaCalo = Part->Eta;
  Part->PhiCalo = Part->Phi;//-atan2(Part->Px,Part->Py);

  // trivial cases
  if (!DET->FLAG_bfield ) return;

  double M; // GeV/c²
  //int q1  = ChargeVal(Part->PID) *eplus; // in units of 'e'
  if(Part->M < -999) { // unitialised!
     PdgParticle pdg_part = DET->PDGtable[Part->PID];
     q  = pdg_part.charge() *eplus; // in units of 'e'
     M  = pdg_part.mass(); // GeV/c²
  } else  { q = Part->Charge; M = Part->M; }

  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;}

  double X  = Part->X/1000.;//[m] 
  double Y  = Part->Y/1000.;//[m]
  double Z  = Part->Z/1000.;//[m]

  // out of tracking coverage?
  if(sqrt(X*X+Y*Y) > R_max){return;}
  if(fabs(Z) > z_max){return;}

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

     //in test mode, just run once
     if (loop_overflow_counter) return;
  
    // 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)
    double Px = Part->Px;                    // [GeV/c]
    double Py = Part->Py;
    double Pz = Part->Pz;
    double PT = Part->PT;
    double E  = Part->E;	             // [GeV]
    double Phi = UNDEFINED;

    float c_light = 2.99792458E+8;
    gammam = E*1E9/(c_light*c_light);               // gammam in [eV/c²]
    omega = q * B_z / (gammam);                     // omega is here in [ 89875518 / s]  
    //cout << "omega*gammam = B_z in BFieldProp.cc: " << fabs(omega*gammam) - B_z  << endl;
    r = PT / (omega * gammam)  *1E9/c_light;        // in [m]  

    // test mode ?
    bool test=false; if(test) loop_overflow_counter++;

    double delta= UNDEFINED;
    phi_0 = atan2(Py,Px); // [rad] in [-pi ; pi ]

    // 2. helix axis coordinates
      x_c = X + r*sin(phi_0);
      y_c = Y - r*cos(phi_0);
      R_c = sqrt( pow(x_c,2.) + pow(y_c,2.) );
      Phi_c = atan2(y_c,x_c);
      Phi = Phi_c;
      if(x_c<0) Phi += pi;

    // 3. time evaluation t = min(t_T, t_z)
    //    t_T : time to exit from the sides
    //    t_z : time to exit from the front or the back
      rr = sqrt( pow(R_c,2.) + pow(r,2.) );     // temp variable [m]
      t_T=0; 					//[ns]
      int sign_pz= (Pz >0) ? 1 : -1;
      if(Pz==0) t_z = 1E99;
      else t_z = gammam / (Pz*1E9/c_light) * (-Z + z_max*sign_pz );
      if( t_z <0) cout << "ERROR: t_z <0 !" << endl;

      if ( fabs(R_c - fabs(r)) > R_max || R_c + fabs(r)  < R_max ) t = t_z;
      else {
	 if(r==0) cout << "r ==0 !" << endl;
	 if(R_c==0) cout << "R_c ==0 !" << endl;
         if(r==0|| R_c ==0) t_T=1E99;
         else { 
		double asinrho = asin( (R_max + rr)*(R_max - rr) / (2*fabs(r)*R_c)  );
		delta = phi_0 - Phi;
		if(delta<-pi) delta += 2*pi;
		if(delta> pi) delta -= 2*pi;
		double t1 = (delta + asinrho) / omega;
		double t2 = (delta + pi - asinrho) / omega;
                double t3 = (delta + pi + asinrho) / omega;
                double t4 = (delta - asinrho) / omega;
                double t5 = (delta - pi - asinrho) / omega;
                double t6 = (delta - pi + asinrho) / omega;

		if(test) {
			cout << "t4 = " << t4 << "\t t5 = " << t5 << "\t t_6=" << t6 << "\t t_3 = " << t3 << endl;
			cout << "t1 = " << t1 << "\t t2 = " << t2 << "\t t_T=" << t_T << "\t t_z = " << t_z << endl;
			cout << "delta= " << delta << endl;
		}
		if(t1<0)t1=1E99;
		if(t2<0)t2=1E99;
		if(t3<0)t3=1E99;
		if(t4<0)t4=1E99;
		if(t5<0)t5=1E99;
		if(t6<0)t6=1E99;

		double t_Ta = min(t1,min(t2,t3));
                double t_Tb = min(t4,min(t5,t6));
	        t_T = min(t_Ta,t_Tb);
	        t = min(t_T,t_z); 
         }
      }

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

    // 5. position in terms of Theta(t), Phi(t), R(t), Eta(t)
      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 = UNDEFINED;
      }

     if(test) {
	cout << endl << endl;
	cout << "x0,y0,z0= " << X << ", " << Y << ", " << Z << endl;
	cout << "px0,py0,pz0= " << Px << ", " << Py << ", " << Pz << endl;
	cout << "r = " << r << "R_max = " << R_max << "\t phi_0=" << phi_0 << endl;
	cout << "gammam= " << gammam << "\t omega=" << omega << "\t PT = " << PT << endl;
	cout << "x_c = " << x_c << "\t y_c = " << y_c << "\t R_c = " << R_c << "\t Phi = " << Phi << endl;
	cout << "omega t = " << omega*t << "\t";
        cout << "cos(omega t -phi0)= " << cos(omega*t-phi_0) << "\t sin(omega t -phi0)= " << sin(omega*t-phi_0) << endl;
	cout << "t_T = " << t_T << "\t t_z = " << t_z << "\t r = " << r << endl;
	cout << "x_t = " << x_t << "\t y_t = " << y_t << "\t z_t = " << z_t << endl;
	cout << "R_t = " << R_t << "\t Phi_t = " << Phi_t << "\t";
	cout << "Theta_t = " << Theta_t << "\t Eta_t = " << Eta_t << endl;
     }

      
/*	Not needed here. but these formulae are correct -------
	// method1 (removed)
 	Px_t = - PT * sin(omega*t + phi_0); 
	Py_t =   PT * cos(omega*t + phi_0);

        // method2
 	Px_t =   PT * cos(phi_0 - omega*t);  
	Py_t =   PT * sin(phi_0 - omega*t);

        Pz_t =   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(M*M +p_t*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);
*/

	Part->EtaCalo = Eta_t;
	Part->PhiCalo = Phi_t;

// test zone ---
/*

        cout << "r = " << r << " et " << fabs(PT/(q*B_z)) << endl;
	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;
*/
	return;

  } else { // if B_x or B_y are non zero: longer computation
//cout << "bfield de loic\n";
  float Xvertex1 = Part->X;
  float Yvertex1 = Part->Y;
  float Zvertex1 = Part->Z;
  
  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; // see above
  double vx = px/M;
  double vy = py/M;
  double vz = pz/M;
  double qm = q/M;
 
//double v = sqrt(vx*vx + vy*vy + vz*vz)/3E8;
//cout << "v = " << v;
//double gamma = 1./sqrt(1-v*v);
//cout << "gamma = " << gamma << endl;

  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;
  float Theta=0; 
  if(x!=0 && y!=0 && z!=0) {
          Theta = atan2(sqrt(r2),z);
          Part->EtaCalo  = -log(tan(Theta/2.));
	  Part->PhiCalo = atan2(y,x);
	  //momentum.SetPtEtaPhiE(Part->PT,eta,phi,Part->E);
  }
 } // if b_x or b_y non zero
}