#include<iostream>
#include<fstream>
#include<math.h>

//this routine transforms cartesian to molecular coordinates/momenta
//of a diatomic molecule
void transform_back(double* X, double* x, double* PX, double* px, double mA, double mB){

  double x1 = x[0]; double x2 = x[3]; //double px1 = px[0]; double px2 = px[3];
  double y1 = x[1]; double y2 = x[4]; //double py1 = px[1]; double py2 = px[4];
  double z1 = x[2]; double z2 = x[5]; //double pz1 = px[2]; double pz2 = px[5];

  double a = mA/(mA+mB);
  double b = mB/(mA+mB);


  double dx = x1-x2; double dy = y1-y2; double dz = z1-z2;
  //TRANSFORM COORDINATES
  X[0] = a*x1 + b*x2; X[1] = a*y1 + b*y2;  X[2] = a*z1 + b*z2;
  X[3] = sqrt(dx*dx+dy*dy+dz*dz); 
  double cos_theta = dz/X[3];
  X[4] = acos(cos_theta);
  //To avoid numerical problems
  double sin_phi = 0.; double cos_phi = 1.;
  X[5] = 0.;
  //DETERMINE PHI;
  //sin_theta = sqrt(1-cos(theta)*cos(theta)) because theta [0,pi]
  double sin_theta = sqrt(1. - cos_theta*cos_theta);
  double rfi = X[3]*sin_theta;
  if(sin_theta > 1e-10){
    sin_phi = dy/rfi;
    cos_phi = dx/rfi;
    X[5] = acos(cos_phi);
    if(sin_phi < 0.) X[5] = 2.*M_PI - X[5];
  }

  //X;Y;Z momenta
  PX[0] = px[0] + px[3]; //confirmed
  PX[1] = px[1] + px[4]; //confirmed
  PX[2] = px[2] + px[5]; //confirmed
  PX[3] = (b*px[0] - a*px[3])*cos_phi*sin_theta
        + (b*px[1] - a*px[4])*sin_phi*sin_theta
        + (b*px[2] - a*px[5])*cos_theta; //confirmed

  PX[4] = ( b*px[0] - a*px[3])*X[3]*cos_theta*cos_phi
        + ( b*px[1] - a*px[4])*X[3]*cos_theta*sin_phi  
        + (-b*px[2] + a*px[5])*rfi; //confirmed

  PX[5] = (-b*px[0] + a*px[3])*rfi*sin_phi
        + ( b*px[1] - a*px[4])*rfi*cos_phi;
}
//transforms position and momenta from molecular to cartesian coordinate system
void transform(double* X, double* x, double* PX, double* px, double mA, double mB){

  double sin_theta = sin(X[4]); double cos_theta = cos(X[4]);
  double sin_phi   = sin(X[5]); double cos_phi   = cos(X[5]);

  double a = mA/(mA+mB);
  double b = mB/(mA+mB);

  //COORDINATE TRANSFORMATION
  x[0] = X[0] + b*X[3]*sin_theta*cos_phi;
  x[3] = X[0] - a*X[3]*sin_theta*cos_phi;

  x[1] = X[1] + b*X[3]*sin_theta*sin_phi;
  x[4] = X[1] - a*X[3]*sin_theta*sin_phi;

  x[2] = X[2] + b*X[3]*cos_theta;
  x[5] = X[2] - a*X[3]*cos_theta;

  //MOMENTA TRANSFORMATION
  double work_val = 0.;
  if(sin_theta > 1e-10) work_val = 1./(X[3]*sin_theta);
  double fac1 = PX[3]*sin_theta*cos_phi 
              + PX[4]/X[3] * cos_theta*cos_phi 
              - PX[5]*work_val*sin_phi;

  px[0] = a*PX[0] + fac1;
  px[3] = b*PX[0] - fac1;

  fac1 = PX[3]*sin_theta*sin_phi
       + PX[4]/X[3] * cos_theta*sin_phi 
       + PX[5]*work_val*cos_phi;

  px[1] = a*PX[1] + fac1;
  px[4] = b*PX[1] - fac1;

  fac1  = PX[3]*cos_theta - (PX[4]/X[3]) * sin_theta;
 
  px[2] = a*PX[2] + fac1;
  px[5] = b*PX[2] - fac1;

}