/*
 *   ---- 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 "interface/BFieldProp.h"
#include "TRandom.h"

#include <iostream>
#include <sstream>
#include <fstream>
#include <iomanip>

#include<cmath>


using namespace std;


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

TrackPropagation::TrackPropagation() {

 MAXITERATION = 20000;
 MINSEGLENGTH = 70;

}

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

  double q  = Charge(Part->PID);
  if(q==0)return;
  
  float Xvertex1 = Part->X;
  float Yvertex1 = Part->Y;
  float Zvertex1 = Part->Z;
  
  //out of trackibg coverage?
  if(sqrt(Xvertex1*Xvertex1+Yvertex1*Yvertex1) > TRACKING_RADIUS){return;}
  if(fabs(Zvertex1) > TRACKING_LENGTH){return;}
  
  float Px = Part->Px;
  float Py = Part->Py;
  float Pz = Part->Pz;
    
  double px = Px / 0.003;
  double py = Py / 0.003;
  double pz = Pz / 0.003;
  double pt = sqrt(px*px+py*py);
  double p  = sqrt(px*px+py*py+pz*pz);
 
  if(q!=0){
     double e  = Part->E  / 0.003;
//     double M  = sqrt(e*e - (px*px + py*py + pz*pz) );
     double M  = Part->M;
/*if(fabs(Part->PID)==11)
{
     cout<<"genMomentum.M() "<<genMomentum.M()<<"  "<<Part->M<<endl;
     cout<<"e*e - (px*px + py*py + pz*pz) "<<e*e - (px*px + py*py + pz*pz)<<endl;
}*/
     double vx = px/M;
     double vy = py/M;
     double vz = pz/M;

     double Bx = BFIELD_X;
     double By = BFIELD_Y;
     double Bz = BFIELD_Z;

     double ax =  (q/M)*(Bz*vy - By*vz);
     double ay =  (q/M)*(Bx*vz - Bz*vx);
     double az =  (q/M)*(By*vx - Bx*vy);
/*
cout<<"M "<<M<<" et le pid "<<Part->PID<<endl;
cout<<"ax "<<ax<<" ay "<<ay<<" az "<<az<<endl;
*/
     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 = sqrt(vx*vx+vy*vy);
 
     int k = 0;
     
     while(k < MAXITERATION){
        k++;

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

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

        ax = (q/M)*(Bz*vy - By*vz);
        ay = (q/M)*(Bx*vz - Bz*vx);
        az = (q/M)*(By*vx - Bx*vy);

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

       if( (x*x+y*y) > TRACKING_RADIUS*TRACKING_RADIUS ){ x /= (x*x+y*y)/(TRACKING_RADIUS*TRACKING_RADIUS); y /= (x*x+y*y)/(TRACKING_RADIUS*TRACKING_RADIUS); break;}
       if( fabs(z)>TRACKING_LENGTH)break;

       xold = x;
       yold = y;
       zold = z;
     }
 
      if(x!=0 && y!=0 && z!=0)
	{
	  double theta = fabs(atan(y/z));
	  double eta;
	  if(z > 0)eta = -log(tan(theta/2));
	  else eta = -(-log(tan(theta/2)));
	  double phi = atan2(y,x);
	  genMomentum.SetPtEtaPhiE(Part->PT,eta,phi,Part->E);
	}
     // cout<<"genMomentum final "<<genMomentum.Pt()<<endl;
    }
}