#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"

MyTrajectoryBfield::MyTrajectoryBfield(const DMagneticFieldMap *bfield_in)
  : B(3), bfield(bfield_in), nparams(5) {
  //  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 x0, y0, pinv, sinThetaX, sinThetaY;
  x0 = param(1);
  y0 = param(2);
  pinv = param(3);
  sinThetaX = param(4);
  sinThetaY = param(5);
  MyTrajectoryBfield::swim(x0, y0, pinv, sinThetaX, sinThetaY);
  return;
}

void MyTrajectoryBfield::swim(double e, HepVector startingPoint, double p,
			     double theta, double phi) {
  double x0, y0, pinv, sinThetaX, sinThetaY;
  if (startingPoint(3) != 0.0) {
    cout << "starting point must be z = 0\n";
    int ierror = 3;
    throw ierror;
  }
  x0 = startingPoint(1);
  y0 = startingPoint(2);
  pinv = e/p;
  sinThetaX = sin(theta)*cos(phi);
  sinThetaY = sin(theta)*sin(phi);
  swim(x0, y0, pinv, sinThetaX, sinThetaY);
  return;
}

void MyTrajectoryBfield::swim(double x0, double y0, double pinv,
			     double sinThetaX, double sinThetaY) {
  //  cout << "starting helix swim\n";
  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;
  v_0(1) = beta*c*sinThetaX; // cm/ns
  v_0(2) = beta*c*sinThetaY;
  v_0(3) = beta*c*sqrt(1.0 - sinThetaX*sinThetaX - sinThetaY*sinThetaY);
  //  cout << "MyTrajectoryBfield::swim initial velocity:" << v_0;
  HepVector *thisVector;
  r_0(1) = x0;
  r_0(2) = y0;
  r_0(3) = 0.0;
  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++) {
    //B(1) = 0.0; B(2) = 0.0; B(3) = 4.0;
    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);
    //    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;
  }
  return;
}