#include <IOAMagneticField.hxx>
#include <IFieldManager.hxx>
#include <TVector3.h>

#include <HEPUnits.hxx>
#include <ITrajectoryTracker.hxx>
#include <sys/stat.h>

using std::endl;
using std::cout;
//
// convenient overloads for printing things out
//

#define sign(x) (x > 0) ? 1 : ((x < 0) ? -1 : 0)
/*
ostream& operator COMET::TrajectoryTracker::<< ( ostream& theStream, TVector3& V )
{
  theStream<<" "<<V.X()<<" "<<V.Y()<<" "<<V.Z()<<" ";
  return theStream;
}
*/

int  COMET::ITrajectoryTracker::TrackParticle(TVector3 StartPosition, TVector3 EndPosition, TVector3 momentum, double charge, vector<TVector3>& points,int startPoint)
{
  /*
    Track a particle from StartPosition to EndPosition, filling array points as we go.
      Once we get close to endposition, stop and adjust all points to make a smooth path, unless we missed badly, in which case draw a straight line
  */




  TVector3 InitialMomentum=momentum;
  int nPoints=1;
  points[startPoint]=StartPosition;

  if(charge==0 || momentum.Mag()==0)
    {
      points[startPoint+1]=EndPosition;
      nPoints=2;
      return nPoints;
    }
  TVector3 position=StartPosition;
 

  bool imStraightAsADie=false;
  bool fieldAlwaysZero=true;
  bool exceededPoints=false;
 
  const int nPointsMax= 99999;
  float oldProximityXYZ=1e24; // tracks how close we are to the end point
  float proximityXYZ;
  int nPointClose=2;
  float Closest=oldProximityXYZ; // remembers the closest we managed to get
  TVector3 ClosePosition;
  int nstraight=0;
  int ntries=0;
  float trackLength[nPointsMax];
  float runLength=0.0;
  trackLength[0]=0.0;
  float straightLength=0.0;
  do
    {
     
          TVector3 BField=COMET::Bman().GetFieldValue(position);
      double step=50;
      const double   FieldScale = 1.0/unit::tesla;
      BField=FieldScale*BField;
      if(BField.Mag()>0 && BField.Mag()<0.19)step=1;
      if(momentum.Mag()<1000)step=10;
      /*
        In regions of low field strength, draw a straight line
      */
      
      if(BField.Mag()<0.001)
        {
          if(imStraightAsADie && nPoints>2 && nstraight <10 )
	  {
	    nPoints--; // if we are in a straight section, don't add on the intervening points
	    nstraight++;
	  }
	  if(nstraight==9)nPoints++;
          TVector3 delta;
          delta.SetX(step*momentum.X()/momentum.Mag());
          delta.SetY(step*momentum.Y()/momentum.Mag());
          delta.SetZ(step*momentum.Z()/momentum.Mag());
          position=position+delta;
          imStraightAsADie=true;
          straightLength=straightLength+delta.Mag();
        }
      else
        {
	  fieldAlwaysZero=false;
          StepHelix(position,momentum,charge);
          TVector3 unit = momentum.Unit();
          imStraightAsADie=false;
          nstraight=0;
        }
      // how close to end point are we ?
      proximityXYZ=(EndPosition-position).Mag();
      TVector3 difference=EndPosition-position;
     
      if(proximityXYZ>1e5){
        //

        break; // another emergency exit
      }
            if(proximityXYZ>oldProximityXYZ && oldProximityXYZ < 50 )
        {
      
          break; // if we WERE  close but now we are getting further away, leave the loop
        }
      if(proximityXYZ<Closest) // remember the closest point we managed to find
        {
          nPointClose=nPoints;
          Closest=proximityXYZ;
          ClosePosition=position;
	
        }

      if(nPoints>nPointsMax) // if we have tried for  a long time and failed, give up
        {
	
          nPoints=nPointClose+1; // go back to the nearest point we found
          exceededPoints=true;
          break; // abnormal exit
        }
        unsigned int iloc=startPoint+nPoints-1;
       
        if(iloc>points.size()-1){
          points.resize(points.size()+1000);
        }
      points[iloc]=position;
      runLength+=step;
      trackLength[nPoints]=runLength;
     
      nPoints++;
      if(nstraight>=10)nstraight=0;
      oldProximityXYZ=proximityXYZ;
     
    }while(ntries++<1e7);
  //
  /*
    If the field was always zero, join the dots and leave
  */
  if(fieldAlwaysZero){
  
      
      points[startPoint+1]=EndPosition;
      nPoints=2;
      //
      return nPoints;
  }

  /*
    Now we adjust all points in the track to get a smooth transition from start to end,provided we are close enough, however if we missed by
     a lot, more honest to just draw a straight line.

  */
   position=points[startPoint+nPoints-1];  //  refill position with the last point that is going to be drawn
   proximityXYZ=(EndPosition-position).Mag(); // compare to the target end position
   TVector3 difference=EndPosition-position;

  /*
      If we were closer before, go back there
  */

  if(Closest<proximityXYZ)
    {
   
      nPoints=nPointClose+1;
      points[startPoint+nPoints-1]=ClosePosition;
      position=points[startPoint+nPoints-1];  //  refill position with the last point that is going to be drawn
      position=ClosePosition;
      proximityXYZ=(EndPosition-position).Mag(); // compare to the target end position
      TVector3 difference=EndPosition-position;
    }
  if((nPoints>1) && ( proximityXYZ <250))
    {
      TVector3 delta = EndPosition-position;
 

      /*
        This loop adjusts each point so that the line is guaranteed to end at the
        next point . delta is the vector between where we ended and where we should have ended.
        We has a fraction of delta on to each point, the fraction is the proportion of the whole
        track length that this point represents. In cases with both a straight and curved section
        this fails so don't do it in those cases.
      */
      if(straightLength<500){
        for(int step=1;step<nPoints;++step){
        
          float fraction = trackLength[step]/trackLength[nPoints-1];
          TVector3 d = delta*fraction;
          
          points[startPoint+step]+=d;
	 
        }
      
      }
   
    }
  else
    {


      COMETWarn("COMET::TrajectoryTracker::TrackParticle:"<<endl<<" Failed to track helix from position "
        <<StartPosition.X()<<" "
        <<StartPosition.Y()<<" "
        <<StartPosition.Z()<<" "
        <<" to "
        <<EndPosition.X()<<" "
        <<EndPosition.Y()<<" "
        <<EndPosition.Z()<<"."
        <<" with initial momentum "
        <<InitialMomentum.X()<<" "
        <<InitialMomentum.Y()<<" "
        <<InitialMomentum.Z()<<endl
        <<" The closest approach was "<<Closest<<"mm."<<endl);
      points[startPoint+1]=EndPosition;
      nPoints=2;

    }

   
  return nPoints;
}
void  COMET::ITrajectoryTracker::StepHelix(TVector3& position, TVector3& momentum,double charge){
    /*
    step along a helical track starting at position with momentum.
    update position and momentum and return
    */
    //   COMETNamedTrace("EveDisplay"," Step Helix at "<<position);
    const double step=1;
    
    
  
    TVector3 BField=COMET::Bman().GetFieldValue(position);
    
    const double   FieldScale = 1.0/unit::tesla;
    BField=FieldScale*BField;
    double cosl = (momentum * BField)/(momentum.Mag()*BField.Mag())	;
    //                                                         calculate the cosine of the angle between the track and the direction of the b field
    
    double sinl = sqrt(1-cosl*cosl);                // and corresponding sine
   
    TRotation BToLab;
    BToLab.SetZAxis(BField);			 //    Find the rotation from B field to lab
    
    TRotation LabToB =BToLab.Inverse();              // and the on from lab to B field frame
    TVector3 pB=LabToB * momentum;                  // find the momentum vector in the B frame
    
    float signOfCharge=sign(charge);
    
    double   phi = pB.Phi()+signOfCharge*(TMath::Pi()/2.0); // phi of momentum in the B field
    double   pt = pB.Mag()*sinl;                   // transverse momentum
    double   R=pt/(0.3*BField.Mag());              // radius of curvature
    double   dr=step*sinl;                         // step in rho  ( around the circle which helix produces in the x,y plane of the B field
    double   dphi=signOfCharge*fabs(dr)/R;         // change in phi
  
    
    TRotation DeltaPhiMinus;                       // create a rotation by dphi
    DeltaPhiMinus.RotateZ(-dphi);
    pB = DeltaPhiMinus * pB ;                     // rotate the momentum vector  by delta phi  in the B field
    momentum = BToLab * pB;                       // find new momentum vector in the lab frame
    
    TVector3 Rvd(R*cos(phi),R*sin(phi),0.0);           // find position vector of particle in a frame whose centre is at the centre of rotation
    TVector3 Rvd2 = DeltaPhiMinus * Rvd;            // rotate the position vector Rvd by deltaPhi
    TVector3 Rdiff = Rvd2-Rvd;                     // find the change in position
    double deltaz=step*cosl;                       // calculate the distance traveled along the beam direction
    Rdiff.SetZ(deltaz);                          // add it to the difference vector
    TVector3 RdiffLab = BToLab * Rdiff ;           // rotate back to lab frame
    position=position+RdiffLab;                  // add back onto position, to find the new position - phew - we're done
   
}
int  COMET::ITrajectoryTracker::G4Track( COMET::IG4Trajectory& trajectory,vector<TVector3>& points)
{
  /*
     Create a set of points to represent a single Geant trajectory
  */

  
   int parcode                                   = trajectory.GetPDGEncoding();
   std::vector<COMET::IG4TrajectoryPoint> trajpvect= trajectory.GetTrajectoryPoints();
   if(!fFoundPDG)return false;


  float charge;
  
  if(fPDGDatabase->GetParticle(parcode)!=NULL)
  {
    charge=fPDGDatabase->GetParticle(parcode)->Charge()/3.0;
  }
  else
  {
    return -1;
  }
  
  if(points.size()<1000)points.resize(1000);	    
  TVector3 EndPosition,StartPosition,FirstPosition,momentum;
//  vector<TVector3> points(100000);
  bool first=true;
  /*
    Loop through all points on the track
  */
  unsigned int nPoint=0;

  
  for (std::vector<COMET::IG4TrajectoryPoint>::const_iterator trackPoint= trajpvect.begin();
       trackPoint != trajpvect.end(); trackPoint++ ) {
    EndPosition.SetX(trackPoint->GetPosition().X());
    EndPosition.SetY(trackPoint->GetPosition().Y());
    EndPosition.SetZ(trackPoint->GetPosition().Z());
 
   
    if(!first)
      {
	/*
	  Go and get a set of points from TrackParticle along the track
	*/
	if(charge!=0){
	  nPoint +=TrackParticle(StartPosition, EndPosition, momentum, charge,points,nPoint);	
	
	}
	else
	  {
	
	      if(nPoint>points.size()){
	        points.resize(points.size()+1000);	     
	      }
	   
	      points[++nPoint]=EndPosition;
	  }
      }
      else
      { 
        if(charge==0){
         
          points[0]=EndPosition;
        }
      }
    momentum=trackPoint->GetMomentum(); // the momentum for the next section is the momentum at the start so rememember the momentum here
    StartPosition=EndPosition;
    first=false;
  } // end of loop over points on track
  
  return nPoint;
}
COMET::ITrajectoryTracker::ITrajectoryTracker(){
    fFoundPDG=true;
    fPDGDatabase = TDatabasePDG::Instance();
    if(fPDGDatabase->GetParticle("gamma") != NULL)return;
    TString pdgtable( Form("%s/etc/pdg_table.txt", gSystem->Getenv("ROOTSYS")) );
    if(FileExists(pdgtable))
      fPDGDatabase->ReadPDGTable(pdgtable);
    else
      {
        pdgtable= Form("%s/pdg_table.txt", gSystem->Getenv("ROOTSYS"));
        if(FileExists(pdgtable))
          fPDGDatabase->ReadPDGTable(pdgtable);
        else
          {
            pdgtable= Form("%s/root/etc/pdg_table.txt", gSystem->Getenv("ROOTSYS"));
            if(FileExists(pdgtable))
              fPDGDatabase->ReadPDGTable(pdgtable);
            else
              {
                COMETWarn(" I failed to find the list of PDG codes stored in pdg_table.txt - it should be under ROOTSYS somewhere. I will assume all MC particles have charge 1.0.");
                fFoundPDG=false;
              }
          }
      }
}
// utility routine to check if a file exists
bool COMET::ITrajectoryTracker::FileExists(TString TstrFilename) {
  string strFilename= (string) TstrFilename;
  struct stat stFileInfo;
  bool blnReturn;
  int intStat;

  // Attempt to get the file attributes
  intStat = stat(strFilename.c_str(),&stFileInfo);
  if(intStat == 0) {
    // We were able to get the file attributes
    // so the file obviously exists.
    blnReturn = true;
  } else {
    // We were not able to get the file attributes.
    // This may mean that we don't have permission to
    // access the folder which contains this file. If you
    // need to do that level of checking, lookup the
    // return values of stat which will give you
    // more details on why stat failed.
    blnReturn = false;
  }

  return(blnReturn);
}