#include <iostream>
#include <vector>
#include <CLHEP/Matrix/Matrix.h>
#include <CLHEP/Matrix/Vector.h>

using namespace std;

#include "HDGEOMETRY/DMagneticFieldMap.h"
#include "MyTrajectory.h"
#include "MyTrajectoryBfield.h"

#define TRACKING_RADIUS_MAX 60.0

MyTrajectoryBfield::MyTrajectoryBfield(const DMagneticFieldMap *bfield_in, int level)
  : MyTrajectory(level), B(3), bfield(bfield_in), nparams(5), delta(5, 0.001), debug_level(level) {
  //  cout << "MyTrajectoryBfield constructor called\n";
  //  cout << "nparams = " << nparams << endl;
  //  cout << "MyTrajectoryBfield::MyTrajectoryBfield B:" << B;
  return;
}

unsigned int MyTrajectoryBfield::getNumberOfParams() {
  return nparams;
}

void MyTrajectoryBfield::swim(const HepVector& param) {
  double xp0, z0, theta0, phi0, pinv;
  xp0 = param(1);
  z0 = param(2);
  theta0 = param(3);
  phi0 = param(4);
  pinv = param(5);
  MyTrajectoryBfield::swim(xp0, z0, theta0, phi0, pinv);
  return;
}

void MyTrajectoryBfield::swim(double xp0, double z0, double theta,
			     double phi, double pinv) {
  if (debug_level >= 3) cout << "starting B-field swim, xp0 = " << xp0 << " z0 = " << z0 << " theta = " << theta << " phi = " << phi << " pinv = " << pinv << endl;
  checkClear();
  double c = 30.0; // speed of light, cm/s
  double m=0.139, dt=0.01; // gev/c^2, ns
  HepVector r_0(3), r_1(3), r_2(3), k(3), v_0(3);
  HepMatrix A(3,3,0), I(3,3,1), M(3,3), onePlusA(3,3);
  int ierr;
  double beta, k_1;
  if (pinv != 0.0) {
    double p = 1.0/pinv;
    double E = sqrt(m*m + p*p);
    double gamma = E/m;
    beta = sqrt(1.0 - 1.0/(gamma*gamma));
    double e;
    if (pinv > 0.0) {e = 1.0;} else {e = -1.0;} // only allow singly
                                                  // charged particles
    k_1 = (e*1.6e-19)*(dt*1.0e-9)/(gamma*(m*1.78e-27));
    //    cout << "gamma = " << gamma << endl;
  } else {
    beta = 1.0;
    k_1 = 0.0;
  }
  //  cout << "beta = " << beta << " k_1 = " << k_1 << endl;
  double v = beta*c; // cm/ns
  double sinTheta = sin(theta);
  double cosTheta = cos(theta);
  double sinPhi = sin(phi);
  double cosPhi = cos(phi);
  v_0(1) = v*sinTheta*cosPhi;
  v_0(2) = v*sinTheta*sinPhi;
  v_0(3) = v*cosTheta;
  //  cout << "MyTrajectoryBfield::swim initial velocity:" << v_0;
  HepVector *thisVector;
  r_0(1) = xp0*sinPhi; // = xp0*cos(alpha) where alpha = phi - pi/2
  r_0(2) = -xp0*cosPhi; // = xp0*sin(alpha)
  r_0(3) = z0;
  thisVector = new HepVector(3);
  *thisVector = r_0;
  traj.push_back(thisVector);
  r_1 = r_0 + v_0*dt;
  thisVector = new HepVector(3);
  *thisVector = r_1;
  traj.push_back(thisVector);
  for (int istep = 0; istep < 2000; istep++) {
    bfield->GetField(r_1(1), r_1(2), r_1(3), B(1), B(2), B(3));
    //cout << "MyTrajectoryBfield::swim B vector:" << B;
    k = 0.5*k_1*B;
    //cout << "MyTrajectoryBfield::swim k vector:" << k;
    A(1,2) = -k(3);
    A(1,3) = k(2);
    A(2,1) = k(3);
    A(2,3) = -k(1);
    A(3,1) = -k(2);
    A(3,2) = k(1);
    onePlusA = I + A;
    M = onePlusA.inverse(ierr)*(I - A); // can we take advantage of anti-sym?
    r_2 = r_1 + M*(r_1 - r_0);
    if (debug_level >=3 ) {
      if (istep%300 == 0) cout << istep << ' ' << r_2(1) << ' '
			       << r_2(2) << ' ' << r_2(3) << endl;
    }
    thisVector = new HepVector(3);
    traj.push_back(thisVector);
    *thisVector = r_2;
    r_0 = r_1; r_1 = r_2;
    if (sqrt(r_0(1)*r_0(1) + r_0(2)*r_0(2)) > TRACKING_RADIUS_MAX) break;
  }
  return;
}