#include "hparticleRK.h"
#include "hkalgeomtools.h"
#include "TMath.h"
#include <math.h>
#include <iostream>
#include <iomanip>
#include <fstream>

using namespace std;


ClassImp(HParticleRKPlane)
ClassImp(HParticleRK)

//_HADES_CLASS_DESCRIPTION
///////////////////////////////////////////////////////////////////////////////
//
// HParticleRKPlane
//
// A class that implements a planar surface of infinite dimension. The plane
// is defined either by a point on the plane and its normal vector
//
// HParticleRK
//
// This is the track propagator for the RK tracks. Propagation is done
// with the Runge-Kutta method 4th ord, addaptive stepsize in SEC coord. sys.
// The track state is x,y,z,px,py,pz [MeV/c,mm]
// Coordinate systems have to be used constistenly. SEC. coord. sys
// is used for definitions of normal+center vectors of the planes
// or PCA (vertex), filed map and the track states.
//
// Before doing track propagation, a pointer to the magnetic field map must
// be set using the function
// setFieldMap(HMdcTrackGField *fpField, Double_t fpol). fpol keeps the field
// map scaling + polarity of the field
//
// The main function is propagateToPlane() that will propagate the track
// to a target plane and calculate the track state vector at that plane and
// the track length from starting point. It will return a boolean indicating whether
// problems were encountered during the Runge-Kutta stepping.
//
// All the points and B-field vectors may be stored in two arrays for
// debugging purposes or graphics output if that option has been enabled with
// setFillPointsArrays(Bool_t fill,Bool_t owner).
// The functions getPointsTrack() and getFieldTrack() will return these arrays.
// Each array element is a TVector3 object. If owner = kTRUE, the point
// Arrays will be deleted by each call the progagateXXX(). If owner=kFALSE
// The arrays will store points of continous calls to propagateXXX() and
// The user has to call HParticleRK::clearPointsArrays() when finished to
// clean up the memory.
///////////////////////////////////////////////////////////////////////////////
		
HParticleRKPlane::HParticleRKPlane(const TVector3 &center, const TVector3 &normal)
    : TObject()
{
    // Creates a plane defined by a point on the plane and its normal vector.
    // Hesse Normal Form: center * normal - d = 0
    // center: a point on the plane.
    // normal: normal vector of the plane.

    setPlane(center, normal);
}

HParticleRKPlane::~HParticleRKPlane() {
}

Double_t HParticleRKPlane::distanceToPlane(const TVector3 &point) const
{
    // Calculates the distance of vector pos to the plane.

    return HKalGeomTools::distancePointToPlane(point, getCenter(), getNormal());
}

Bool_t HParticleRKPlane::findIntersection(TVector3 &pointIntersect, const TVector3 &pos, const TVector3 &dir) const
{
    // Finds the intersection point of a straight line with this plane.
    // pointIntersect: the intersection point (return value).
    // t:              line parameter to reach the plane (return value)
    // pos:            a point on the straight line.
    // dir:            direction of the straight line.

    return HKalGeomTools::findIntersectionLinePlane(pointIntersect, pos, dir, getCenter(), getNormal());
}

Bool_t HParticleRKPlane::isOnSurface(const TVector3 &point) const
{
    // Checks if point is on the plane.

    return HKalGeomTools::isPointOnPlane(point, getCenter(), getNormal());
}

void HParticleRKPlane::print(Option_t* opt) const
{
    // Print center and normal vectors of this plane.
    // opt: print options (not used)

    cout<<"**** Properties of plane: ****"<<endl;
    cout<<"Center of layer: "<<endl;
    getCenter().Print();
    cout<<"Normal of layer:"<<endl;
    getNormal().Print();
    cout<<"**** End print of plane ****"<<endl;
}

Double_t HParticleRKPlane::signedDistanceToPlane(const TVector3 &point) const {
    // Calculates the distance of vector pos to the plane.
    // The signed distance is positive if the origin of the coordinate system and the
    // test point are on opposite sides of the plane and negative if they are on the same side.

    return HKalGeomTools::signedDistancePointToPlane(point, getCenter(), getNormal());
}

void HParticleRKPlane::transform(const TRotation &transMat) {
    // Transforms the plane using a rotation matrix.
    vCenter.Transform(transMat);
    vNormal.Transform(transMat);
}

Bool_t HParticleRKPlane::setPlane(const TVector3 &origin, const TVector3 &normal) {
    // Set the plane using a point on the plane and the plane's normal vector.

    vCenter = origin;
    vNormal = normal.Unit();
    // Make sure the normal vector points from (0,0,0) to the plane.
    if(vCenter * vNormal < 0.) {
        vNormal *= -1.;
    }
    return kTRUE;
}

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


HParticleRK::HParticleRK()
{
  // constructor
  fieldMap       = 0;
  fieldScaleFact = 0.;

  // Eventually these parameters should be stored in a parameter container!
  initialStepSize = 50.;
  stepSizeDec     = 0.75;
  stepSizeInc     = 1.25;
  maxStepSize     = 200.;
  minStepSize     = 10.;
  minPrecision    = 2.e-04;
  maxPrecision    = 2.e-05;
  maxNumSteps     = 1000;
  maxDist         = 2.5;

  bFillPointArrays = kFALSE;
  bOwnArrays       = kTRUE;
}

Bool_t HParticleRK::propagateToPlane(const Int_t chrg,
				     TVectorD &stateTo, const TVectorD &stateFrom,
				     const HParticleRKPlane &planeTo,
				     Double_t& length, Bool_t propDir)
{

    // Propagate a particle trajectory to a plane using a 4th order Runge-Kutta.
    // Returns true if no errors were encountered during track propagation.
    //
    // Output:
    // stateTo:    Track state vector (x,y,z,px,py,pz) at target plane. [sec. coord. sys]
    //
    // Input:
    // stateFrom:  Track state vector at (x,y,tx,ty,q/p) at source plane.
    // planeTo:    Plane to propagate to.
    // propDir:    propagation direction kIterForward or kIterBackward.
    // length :    track length from start to end of the propagation

    //#define debugRK
    trackLength = 0;

    TVector3 posAt(stateFrom(0),stateFrom(1),stateFrom(2));
    TVector3 mom  (stateFrom(3),stateFrom(4),stateFrom(5));


    Double_t p = mom.Mag();
    p *= chrg;

    Double_t sd = planeTo.signedDistanceToPlane(posAt);
    if((sd > 0. && propDir == kIterForward) || (sd < 0. && propDir == kIterBackward)) {
	if(TMath::Abs(sd) > 1.e-3) {
	    Error("propagateToPlane()", Form("Track was already past the target plane before propagation by %f.", sd));
	    return kFALSE;
	}
    }

    jstep             = 0;
    Double_t d        = 0.;
    Double_t step     = initialStepSize;
    Double_t nextStep = step;
    Double_t stepFac  = 1.;


    TVector3 isect;

    TVector3 upos = posAt;
    TVector3 udir = mom;
    udir          = udir.Unit();

    //------------------------------------------------------------------------
    // for debugging + graphics purpose
    if(bFillPointArrays) {
	if(bOwnArrays) clearPointsArrays();

	TVector3 b;
	calcField_mm(upos, b, fieldScaleFact);
	pointsTrack.Add(new TVector3(upos));
	fieldTrack .Add(new TVector3(b));
    }
    //------------------------------------------------------------------------

    if(propDir == kIterBackward ){  // reverse dir and mom for backward propagation
	udir *= -1;
	p    *= -1;
    }
#ifdef debugRK
    cout<<"###################################################################################"<<endl;
    cout<<"ini pos "<<setw(10)<< upos[0]<<" "<<setw(10)<< upos[1]<<" "<<setw(10)<< upos[2]
	<<"ini dir "<<setw(10)<< udir[0]<<" "<<setw(10)<< udir[1]<<" "<<setw(10)<< udir[2]
	<<" p "<<setw(10)<<p<<" chrg "<<chrg<<endl;
#endif
    do {
	stepFac = rkgtrk(step,p,upos,udir);
	planeTo.findIntersection(isect,upos,udir);
	d = HKalGeomTools::distance2Points(upos,isect);

#ifdef debugRK
	cout<<"step "<<setw(3)<<jstep
	    <<" pos "<<setw(10)<< upos[0]<<" "<<setw(10)<< upos[1]<<" "<<setw(10)<< upos[2]
	    <<" dir "<<setw(10)<< udir[0]<<" "<<setw(10)<< udir[1]<<" "<<setw(10)<< udir[2]
	    <<" inter "<<setw(10)<< udir[0]<<" "<<setw(10)<< udir[1]<<" "<<setw(10)<< udir[2]
	    <<" d "<<setw(10)<<d<<" p "<<setw(10)<<p<<" pmag "<<setw(10)<<udir.Mag()<<" chrg "<<chrg<<" length "<<setw(10)<<trackLength<<endl;
#endif
	//------------------------------------------------------------------------
	// for debugging + graphics purpose
	if(bFillPointArrays) {
            TVector3 b;
	    calcField_mm(upos, b, fieldScaleFact);
	    pointsTrack.Add(new TVector3(upos));
	    fieldTrack .Add(new TVector3(b));
	}
	//------------------------------------------------------------------------

	if (step < nextStep) {  // last step was rest step to chamber
	    if (stepFac < 1.) nextStep = step * stepFac;
	} else nextStep *= stepFac;
	if (d > nextStep || d < maxDist) step = nextStep;
	else step = d;
    } while (d >= maxDist && jstep < maxNumSteps);

    sd = planeTo.signedDistanceToPlane(upos);

    if((sd > 0. && propDir == kIterForward) || (sd < 0. && propDir == kIterBackward)) {

#ifdef debugRK
	if(TMath::Abs(sd) > 1.e-3) {
	    Error("propagateToPlane()", Form("Track not on the target plane after propagation by %f.", sd));
	}
#endif
	trackLength -= sd;
	planeTo.findIntersection(isect,upos,udir);
	upos = isect; // make sure track is on plane
    }

    trackLength -= sd;
    planeTo.findIntersection(isect,upos,udir);
    upos = isect; // make sure track is on plane

    p *= chrg;

    if(propDir == kIterBackward ){  // reverse dir and mom for backward propagation
	udir *= -1;
	p    *= -1;

    }

    udir  = udir.Unit();
    udir *= p;

    //------------------------------------------------------------------------
    // for debugging + graphics purpose
    if(bFillPointArrays) {
	TVector3 b;
	calcField_mm(upos, b, fieldScaleFact);
	pointsTrack.Add(new TVector3(upos));
	fieldTrack .Add(new TVector3(b));
    }
    //------------------------------------------------------------------------

    stateTo(0) = upos(0);
    stateTo(1) = upos(1);
    stateTo(2) = upos(2);
    stateTo(3) = udir(0);
    stateTo(4) = udir(1);
    stateTo(5) = udir(2);
    length = trackLength;

#ifdef debugRK
    cout<<"final "<<setw(3)<<jstep
	<<" pos "<<setw(10)<< upos[0]<<" "<<setw(10)<< upos[1]<<" "<<setw(10)<< upos[2]
	<<" dir "<<setw(10)<< udir[0]<<" "<<setw(10)<< udir[1]<<" "<<setw(10)<< udir[2]
	<<" inter "<<setw(10)<< udir[0]<<" "<<setw(10)<< udir[1]<<" "<<setw(10)<< udir[2]
	<<" p "<<setw(10)<<p<<" pmag "<<setw(10)<<udir.Mag()<<" chrg "<<chrg<<" length "<<setw(10)<<trackLength<<endl;
#endif

    return kTRUE;

}

Bool_t HParticleRK::propagateToPCA(const Int_t chrg,
				   TVectorD &stateTo, const TVectorD &stateFrom,
				   const TVector3& pointTo,
				   Double_t& length, Bool_t propDir)
{

    // Propagate a particle trajectory to PCA of vertex using a 4th order Runge-Kutta (sec. coord. sys)
    // Returns true if no errors were encountered during track propagation.
    //
    // Output:
    // stateTo:    Track state vector (x,y,z,px,py,pz) at target plane. [sec. coord. sys]
    //
    // Input:
    // stateFrom:  Track state vector at (x,y,z,px,py,pz) at PCA.
    // pointTo:    Point to propagate to.
    // propDir:    propagation direction kIterForward or kIterBackward.
    // length :    track length from start to end of the propagation


//#define debugRK
    trackLength = 0;

    TVector3 posAt(stateFrom(0),stateFrom(1),stateFrom(2));
    TVector3 mom  (stateFrom(3),stateFrom(4),stateFrom(5));


    Double_t p = mom.Mag();
    p *= chrg;

    // pca stepping stuff
    TVector3 pca,point2;
    Int_t    pcaFlag    = 0;
    Double_t precDCA    = 1.e-3;
    Bool_t   doNextStep = kTRUE;
    TVector3 uposBefore;
    TVector3 udirBefore;
    Double_t trackLengthBefore = 0;


    // standard stepping stuff
    jstep               = 0;
    Double_t d          = 0.;
    Double_t step       = initialStepSize;
    Double_t nextStep   = step;
    Double_t stepFac    = 1.;

    TVector3 isect;

    TVector3 upos = posAt;
    TVector3 udir = mom;
    udir          = udir.Unit();


    point2 = upos + 10. * udir;
    Double_t sd = HKalGeomTools::distancePointToLine(pca, pcaFlag, pointTo, upos, point2);
    if((sd > 0. && propDir == kIterForward) || (sd < 0. && propDir == kIterBackward)) {
	if(TMath::Abs(sd) > 1.e-3) {
	    Error("propagateToPlane()", Form("Track was already past PCA  before propagation by %f.", sd));
	    return kFALSE;
	}
	else {
	    Warning("propagateToPlane()", Form("Track was already past PCA before propagation by %f.", sd));
	    return kTRUE;
	}
    }
    Double_t sdOld = sd;

    // initial distance between sv and pca
    d = HKalGeomTools::distance2Points(posAt, pca);
    if(d < step) {
        step     = d;
        nextStep = step;
    }

    //------------------------------------------------------------------------
    // for debugging + graphics purpose
    if(bFillPointArrays) {
	if(bOwnArrays) clearPointsArrays();

	TVector3 b;
	calcField_mm(upos, b, fieldScaleFact);
	pointsTrack.Add(new TVector3(upos));
	fieldTrack .Add(new TVector3(b));
    }
    //------------------------------------------------------------------------

    if(propDir == kIterBackward ){  // reverse dir and mom for backward propagation
	udir *= -1;
	p    *= -1;
    }
#ifdef debugRK
    cout<<"###################################################################################"<<endl;
    cout<<"ini pos "<<setw(10)<< upos[0]<<" "<<setw(10)<< upos[1]<<" "<<setw(10)<< upos[2]
	<<"ini dir "<<setw(10)<< udir[0]<<" "<<setw(10)<< udir[1]<<" "<<setw(10)<< udir[2]
	<<" p "<<setw(10)<<p<<" chrg "<<chrg<<endl;
#endif
    while (d >= maxDist && doNextStep && jstep < maxNumSteps && step)
    {
	uposBefore = upos;
	udirBefore = udir;
        trackLengthBefore = trackLength;

	stepFac = rkgtrk(step,p,upos,udir);
	point2 = upos + udir;
        sd = HKalGeomTools::distancePointToLine(pca, pcaFlag,pointTo,upos, point2);
	d  = HKalGeomTools::distance2Points(upos,isect);

	Double_t prec = 1.e-6;
      
        if((d + prec < step) || (sd - sdOld > precDCA && d < 50)) {
            // Track went past target plane during propagation. Repeating last step with reduced step size.
            step       *= stepSizeDec;
            nextStep   *= stepSizeDec;

            // rewind last step
	    upos = uposBefore;
            udir = udirBefore;
	    trackLength = trackLengthBefore;
            continue;
        }

	sdOld = sd;

#ifdef debugRK
	cout<<"step "<<setw(3)<<jstep
	    <<" pos "<<setw(10)<< upos[0]<<" "<<setw(10)<< upos[1]<<" "<<setw(10)<< upos[2]
	    <<" dir "<<setw(10)<< udir[0]<<" "<<setw(10)<< udir[1]<<" "<<setw(10)<< udir[2]
	    <<" inter "<<setw(10)<< udir[0]<<" "<<setw(10)<< udir[1]<<" "<<setw(10)<< udir[2]
	    <<" d "<<setw(10)<<d<<" p "<<setw(10)<<p<<" pmag "<<setw(10)<<udir.Mag()<<" chrg "<<chrg<<" length "<<setw(10)<<trackLength<<endl;
#endif
	//------------------------------------------------------------------------
	// for debugging + graphics purpose
	if(bFillPointArrays) {
	    TVector3 b;
	    calcField_mm(upos, b, fieldScaleFact);
	    pointsTrack.Add(new TVector3(upos));
	    fieldTrack .Add(new TVector3(b));
	}
	//------------------------------------------------------------------------

	point2 = upos + udir;
        sd = HKalGeomTools::distancePointToLine(pca, pcaFlag,pointTo,upos, point2);
	d  = HKalGeomTools::distance2Points(upos,isect);


	if (step < nextStep) {  // last step was rest step to chamber
	    if (stepFac < 1.) nextStep = step * stepFac;
	} else nextStep *= stepFac;
	if (d > nextStep || d < maxDist) step = nextStep;
	else step = d;

        if(propDir == kIterForward) {
            if(upos.z() < pca.z()) { doNextStep = kTRUE; }
            else { doNextStep = kFALSE; }
        } else {
            if(upos.z() > pca.z()) { doNextStep = kTRUE; }
            else { doNextStep = kFALSE; }
        }
    }

    // stepping done, make some checks
    // and provide final output

    point2 = upos + udir;
    sd = HKalGeomTools::distancePointToLine(pca, pcaFlag,pointTo,upos, point2);

    if((sd > 0. && propDir == kIterForward) || (sd < 0. && propDir == kIterBackward)) {

#ifdef debugRK
	if(TMath::Abs(sd) > 1.e-3) {
	    Error("propagateToPCA()", Form("Track not on PCA after propagation by %f.", sd));
	}
#endif
    }

    trackLength -= sd;
    upos = pca; // make sure track is on plane

    p *= chrg;

    if(propDir == kIterBackward ){  // reverse dir and mom for backward propagation
	udir *= -1;
	p    *= -1;
    }

    udir  = udir.Unit();
    udir *= p;

    //------------------------------------------------------------------------
    // for debugging + graphics purpose
    if(bFillPointArrays) {
	TVector3 b;
	calcField_mm(upos, b, fieldScaleFact);
	pointsTrack.Add(new TVector3(upos));
	fieldTrack .Add(new TVector3(b));
    }
    //------------------------------------------------------------------------

    stateTo(0) = upos(0);
    stateTo(1) = upos(1);
    stateTo(2) = upos(2);
    stateTo(3) = udir(0);
    stateTo(4) = udir(1);
    stateTo(5) = udir(2);
    length = trackLength;

#ifdef debugRK
    cout<<"final "<<setw(3)<<jstep
	<<" pos "<<setw(10)<< upos[0]<<" "<<setw(10)<< upos[1]<<" "<<setw(10)<< upos[2]
	<<" dir "<<setw(10)<< udir[0]<<" "<<setw(10)<< udir[1]<<" "<<setw(10)<< udir[2]
	<<" inter "<<setw(10)<< udir[0]<<" "<<setw(10)<< udir[1]<<" "<<setw(10)<< udir[2]
	<<" p "<<setw(10)<<p<<" pmag "<<setw(10)<<udir.Mag()<<" chrg "<<chrg<<" length "<<setw(10)<<trackLength<<endl;
#endif

    return kTRUE;

}

void HParticleRK::calcField_mm(const TVector3& xv, TVector3& btos, Double_t fpol) const
{
    // Calculate B-field (T) in x,y,z (mm).
    //
    // Output:
    // btos: field vector at position xv in Tesla.
    //
    // Input:
    // xv:   return field at this position vector. Coordinates must be in mm.
    // fpol: field scaling factor * polarity.

    static Double_t abtos[3];
    static Double_t axv  [3];

    axv[0] = xv.X() * 0.1;  // convert to cm
    axv[1] = xv.Y() * 0.1;
    axv[2] = xv.Z() * 0.1;

    fieldMap->calcField(axv,abtos,fpol);

    btos.SetX(abtos[0]);
    btos.SetY(abtos[1]);
    btos.SetZ(abtos[2]);
}


Double_t HParticleRK::rkgtrk(Double_t totStep, Double_t p, TVector3& upos, TVector3& udir)
{

    // one step of track tracing from the position upos in the direction of udir

    TVector3 k1,k2,k3,k4;
    TVector3 b;
    TVector3 u1pos, u1dir, u2pos, u2dir;
    Double_t hstep,qstep;

    Double_t est    = 0.;
    Double_t step    = totStep;
    Double_t stepFac = 1.;

    calcField_mm(upos,b,fieldScaleFact);

    rkeqfw(b,p,udir,k1);

    do {
	jstep ++;
	hstep = step * 0.5;
	qstep = step * step *0.125;

	for(Int_t i = 0; i < 3; i ++) {
	    u1pos[i] = upos[i] + hstep * udir[i] + qstep * k1[i];
	    u1dir[i] = udir[i] + hstep * k1[i];
	}

	calcField_mm(u1pos,b,fieldScaleFact);
	rkeqfw(b,p,u1dir,k2);

	for(Int_t i = 0; i < 3; i ++) {
	    u1dir[i] = udir[i] + hstep * k2[i];
	}
	rkeqfw(b,p,u1dir,k3);

	for(Int_t i = 0; i < 3; i ++) {
	    u2pos[i] = upos[i] + step * udir[i] + 2. * hstep * hstep * k3[i];
	    u2dir[i] = udir[i] + step * k3[i];
	}

        calcField_mm(u2pos,b,fieldScaleFact);
	rkeqfw(b,p,u2dir,k4);

	est = 0.;
	for(Int_t i = 0; i < 3; i ++) {
	    est += fabs(k1[i] + k4[i] - k2[i] - k3[i]) * hstep;
	}

	if ( est < minPrecision || step <= minStepSize) {
	    break;
	} else {
	    stepFac *= stepSizeDec;
	    step    *= stepSizeDec;
	}
    } while (jstep < maxNumSteps);

    for(Int_t i = 0; i < 3; i++) {
	upos[i] = upos[i] + step * udir[i] + step * step * (k1[i] + k2[i] + k3[i])/6.;
	udir[i] = udir[i] + step * (k1[i] + 2. * k2[i] + 2. * k3[i] + k4[i])/6.;
    }

    trackLength += step;  // calculate track length

    if ( est < maxPrecision && step < maxStepSize) {
	 stepFac *= stepSizeInc;
    }

    return stepFac;
}

void HParticleRK::rkeqfw(TVector3& b, Double_t p,TVector3& udir, TVector3& duds)
{
    // calculates 2nd derivative
    Double_t  con = TMath::C()*10e-10;
    duds[0] = con * (udir[1] * b[2] - udir[2] * b[1])/p;
    duds[1] = con * (udir[2] * b[0] - udir[0] * b[2])/p;
    duds[2] = con * (udir[0] * b[1] - udir[1] * b[0])/p;
}


TVector3 HParticleRK::calcFieldIntegral() const {
    // Calculates the field integral (Tm) by numerical integration
    // along the stepping points :  | sum over i  x(i+1)-x(i) x B(i) |
    // where x and B are position and field vectors.
    // setFillPointArrays(kTRUE) has to set before call propagate
    // to fill the point and field arrays which are needed for
    // the calculation.

    TVector3 bdl;
    if(!bFillPointArrays) {
        Warning("calcFieldIntegral()", "Stepping point arrays not filled. Cannot calculate field integral.");
        return bdl;
    }
    Int_t n = pointsTrack.GetEntries() - 1;
    for(Int_t i = 0; i < n ; i++){
        TVector3& x1 = (*(TVector3*)pointsTrack.At(i    ));
        TVector3& x2 = (*(TVector3*)pointsTrack.At(i + 1));
        TVector3& B  = (*(TVector3*)fieldTrack .At(i    ));
        TVector3 l   = (x2 - x1);
        bdl += l.Cross(B);
    }
    return bdl * 0.001; // mm -> m
}