/*
 * axisfinder.cpp
 *
 * Purpose: Finds the settings of a gonimeter that aligns
 *          the crystal and lab coordinate systems.
 *
 * Richard Jones, University of Connecticut
 * November 2, 2006
 *
 * Notes:
 * 1. The lab coordinate system is defined by the following rules.
 *       a. the beam axis is along the z axis
 *       b. the origin is at the center of the X-ray scattering target
 *       c. X-ray scattering is in the yz plane
 *       d. the x,y,z axes form a right-handed orthogonal system
 *
 * 2. The goniometer mount system is defined to coincide with the lab
 *    frame when the goniometer settings are all zero.  Goniometer settings
 *    are recorded at the maximum of the rocking curve when the spot is
 *    aligned in the yz laboratory plane.
 *
 * 3. The crystal axes are defined by the following rules.
 *       a. the user defines which h,k,l vectors are doing the scattering
 *       b. the indices h,k,l belong to a right-handed orthogonal system
 *       c. the basic unit cell of the crystal is cubic
 *
 * 4. A goniometer setting is defined by the following three angles,
 *    whose sense is defined by an active rotation of the mount.
 *       theta - rotation around the lab x axis, direction z->y, [-pi/2,pi/2]
 *       chi - rotation around the z' axis, direction y->x, [-pi,pi]
 *       phi - rotation around the y'' axis, direction z->x, [-pi,pi]
 *
 * 5. Inputs are two sets of maps {h,k,l -> theta,chi,phi} for scattering
 *    in a monochromatic X-ray beam of known energy k.
 *       k (real) - X-ray beam energy (GeV)
 *       hkl (integer[3]) - Miller indices of scattering plane
 *       theta,chi,phi (double) - goniometer setting at peak (degrees)
 *
 * 6. Outputs are as follows.
 *       twotheta,theta,chi,psi (real) output goniometer setting (degrees)
 *    By reporting the results as gonimeter settings, the program makes
 *    it easy for the user to compute goniometer offsets to align with
 *    other arbitrary crystal planes.
 *
 * 7. Ambiguities
 *    There is an ambiguity in aligning a crystal for X-ray scattering
 *    from a particular h,k,l plane: any rotation within the plane
 *    preserves the scattering condition.  In general this does not
 *    correspond to one of the goniometer axes, but rather a mixture
 *    of the three called a "wobble".  Frequently the geometry of the
 *    goniometer prevents its operation over the full range in theta,
 *    chi,phi so it is useful to exploit the wobble to find an allowed
 *    goniometer setting that can reach a particular set of planes.  The
 *    user is able to enter arbitrary values for the wobble to search for
 *    the most convenient choice.
 */

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

//#define use_diamond 1
#define use_silicon 1

#if use_diamond
double aLattice=3.5597e-10;  // diamond lattice constant (m)
char crystal[]="diamond";
#elif use_silicon
double aLattice=5.43095e-10; // silicon lattice constant (m)
char crystal[]="silicon";
#else
# error "You must select either silicon or diamond"
#endif

double hbarc=1.97326e-7;     // hbar*c in eV.m
double kX=15e3;              // X-ray beam energy (eV)

struct gset_ {       // angles in radians
   double theta;
   double chi;
   double phi;
};

struct miller_ {
   int h;
   int k;
   int l;
};

struct xpeak_ {
   struct miller_ miller;
   struct gset_ goni;
};

typedef struct miller_ miller_t;
typedef struct gset_ gset_t;
typedef struct xpeak_ xpeak_t;

class rmatrix {
 public:
   rmatrix() {};
   rmatrix(const rmatrix& src);
   rmatrix& setRotateX(double theta);
   rmatrix& setRotateY(double theta);
   rmatrix& setRotateZ(double theta);
   rmatrix& setGoniometer(const gset_t& goni);
   rmatrix& setVertical(const miller_t& indi);
   gset_t getGoniometer();
   rmatrix inverse() const;
   rmatrix operator*(const rmatrix& src) const;
   rmatrix& operator=(const rmatrix& src);
   rmatrix& operator*=(const rmatrix& src);
   void dump();
   double& operator()(int i, int j) { return rmat[i-1][j-1]; }
 private:
   double rmat[3][3];
};

rmatrix::rmatrix(const rmatrix& src)
{
   rmat[0][0] = src.rmat[0][0];
   rmat[1][0] = src.rmat[1][0];
   rmat[2][0] = src.rmat[2][0];
   rmat[0][1] = src.rmat[0][1];
   rmat[1][1] = src.rmat[1][1];
   rmat[2][1] = src.rmat[2][1];
   rmat[0][2] = src.rmat[0][2];
   rmat[1][2] = src.rmat[1][2];
   rmat[2][2] = src.rmat[2][2];
}

rmatrix& rmatrix::setRotateX(double theta)   // rotate y->z, passive, radians
{
   rmat[0][0] = 1;
   rmat[1][0] = 0;
   rmat[2][0] = 0;
   rmat[0][1] = 0;
   rmat[1][1] = cos(theta);
   rmat[2][1] = -sin(theta);
   rmat[0][2] = 0;
   rmat[1][2] = sin(theta);
   rmat[2][2] = cos(theta);
   return *this;
}

rmatrix& rmatrix::setRotateY(double theta)   // rotate z->x, passive, radians
{
   rmat[0][0] = cos(theta);
   rmat[1][0] = 0;
   rmat[2][0] = sin(theta);
   rmat[0][1] = 0;
   rmat[1][1] = 1;
   rmat[2][1] = 0;
   rmat[0][2] = -sin(theta);
   rmat[1][2] = 0;
   rmat[2][2] = cos(theta);
   return *this;
}

rmatrix& rmatrix::setRotateZ(double theta)   // rotate x->y, passive, radians
{
   rmat[0][0] = cos(theta);
   rmat[1][0] = -sin(theta);
   rmat[2][0] = 0;
   rmat[0][1] = sin(theta);
   rmat[1][1] = cos(theta);
   rmat[2][1] = 0;
   rmat[0][2] = 0;
   rmat[1][2] = 0;
   rmat[2][2] = 1;
   return *this;
}

rmatrix& rmatrix::setGoniometer(const gset_t& goni)
{
   rmatrix stage;
   this->setRotateX(+goni.theta);         // theta is active z->y
   *this *= stage.setRotateZ(+goni.chi);  // chi is active y->x
   *this *= stage.setRotateX(-goni.phi);  // phi is active y->z
   return *this;
}

rmatrix& rmatrix::setVertical(const miller_t& indi)
{
   double q[3];
   q[0] = indi.h;
   q[1] = indi.k;
   q[2] = indi.l;
   double qperp=sqrt(q[0]*q[0]+q[1]*q[1]);
   if (qperp > 1e-10) {
      rmatrix stage;
      double theta = atan2(qperp,q[2]);
      double phi = atan2(q[1],q[0]);
      setRotateX(M_PI/2-theta);
      *this *= stage.setRotateZ(phi-M_PI/2);
   }
   else {
      setRotateX(M_PI/2);
   }
   return *this;
}

gset_t rmatrix::getGoniometer()
{
   gset_t res;

   double cchi = rmat[0][0];
   if (fabs(cchi) < 1) {
      res.theta = atan2(rmat[2][0],-rmat[1][0]);
      if (res.theta > M_PI/2) {
         res.theta -= M_PI;
      }
      else if (res.theta < -M_PI/2) {
         res.theta += M_PI;
      }
      double stheta = sin(res.theta);
      if (fabs(stheta) > 0.5) {
         res.chi = atan2(rmat[2][0]/stheta,rmat[0][0]);
      }
      else if (fabs(stheta) < 1e-6) {
         res.theta = 0;
         res.chi = atan2(-rmat[1][0],rmat[0][0]);
      }
      else {
         double ctheta = cos(res.theta);
         res.chi = atan2(-rmat[1][0]/ctheta,rmat[0][0]);
      }
      double schi = sin(res.chi);
      if (schi > 1e-10) {
         res.phi = atan2(-rmat[0][2],rmat[0][1]);
      }
      else if (schi < -1e-10) {
         res.phi = atan2(rmat[0][2],-rmat[0][1]);
      }
      else {
         res.phi = 0;
      }
   }
   else if (fabs(cchi) < 1) {
      double cc = (rmat[2][2]-rmat[1][1]*cchi)/(1-cchi*cchi);
      double ss = (rmat[1][1]-rmat[2][2]*cchi)/(1-cchi*cchi);
      double sc = (rmat[1][2]+rmat[2][1]*cchi)/(1-cchi*cchi);
      double cs = (rmat[2][1]+rmat[1][2]*cchi)/(1-cchi*cchi);
      if (fabs(ss) > fabs(sc)) {
         res.theta = atan2(ss,cs);
         res.phi = atan2(ss,sc);
         if (res.theta > M_PI/2) {
            res.theta -= M_PI;
            res.phi = atan2(-ss,-sc);
         }
         else if (res.theta < -M_PI/2) {
            res.theta += M_PI;
            res.phi = atan2(-ss,-sc);
         }
         res.chi = atan2(rmat[2][0]/sin(res.theta),cchi);
      }
      else if (fabs(sc) > 0) {
         res.theta = atan2(sc,cc);
         res.phi = atan2(cs,cc);
         if (res.theta > M_PI/2) {
            res.theta -= M_PI;
            res.phi = atan2(-cs,-cc);
         }
         else if (res.theta < -M_PI/2) {
            res.theta += M_PI;
            res.phi = atan2(-cs,-cc);
         }
         res.chi = atan2(-rmat[1][0]/cos(res.theta),cchi);
      }
      else {
         res.theta = 0;
         res.phi = atan2(cs,cc);
         if (res.theta > M_PI/2) {
            res.theta -= M_PI;
            res.phi = atan2(-cs,-cc);
         }
         else if (res.theta < -M_PI/2) {
            res.theta += M_PI;
            res.phi = atan2(-cs,-cc);
         }
         res.chi = atan2(-rmat[1][0]/cos(res.theta),cchi);
      }
   }
   else {
      res.phi = 0;
      res.chi = (cchi > 0)? 0 : M_PI;
      res.theta = atan2(-rmat[2][1]/cchi,rmat[1][1]/cchi);
      if (res.theta > M_PI/2) {
         res.theta -= M_PI;
         res.phi = M_PI;
      }
      else if (res.theta < -M_PI/2) {
         res.theta += M_PI;
         res.phi = M_PI;
      }
   }

   return res;
}

rmatrix rmatrix::inverse() const
{
   rmatrix res;
   res.rmat[0][0] = rmat[0][0];
   res.rmat[1][0] = rmat[0][1];
   res.rmat[2][0] = rmat[0][2];
   res.rmat[0][1] = rmat[1][0];
   res.rmat[1][1] = rmat[1][1];
   res.rmat[2][1] = rmat[1][2];
   res.rmat[0][2] = rmat[2][0];
   res.rmat[1][2] = rmat[2][1];
   res.rmat[2][2] = rmat[2][2];
   return res;
}

rmatrix rmatrix::operator*(const rmatrix& src) const
{
   rmatrix res;
   res.rmat[0][0] = rmat[0][0]*src.rmat[0][0]
                  + rmat[0][1]*src.rmat[1][0]
                  + rmat[0][2]*src.rmat[2][0];
   res.rmat[0][1] = rmat[0][0]*src.rmat[0][1]
                  + rmat[0][1]*src.rmat[1][1]
                  + rmat[0][2]*src.rmat[2][1];
   res.rmat[0][2] = rmat[0][0]*src.rmat[0][2]
                  + rmat[0][1]*src.rmat[1][2]
                  + rmat[0][2]*src.rmat[2][2];

   res.rmat[1][0] = rmat[1][0]*src.rmat[0][0]
                  + rmat[1][1]*src.rmat[1][0]
                  + rmat[1][2]*src.rmat[2][0];
   res.rmat[1][1] = rmat[1][0]*src.rmat[0][1]
                  + rmat[1][1]*src.rmat[1][1]
                  + rmat[1][2]*src.rmat[2][1];
   res.rmat[1][2] = rmat[1][0]*src.rmat[0][2]
                  + rmat[1][1]*src.rmat[1][2]
                  + rmat[1][2]*src.rmat[2][2];

   res.rmat[2][0] = rmat[2][0]*src.rmat[0][0]
                  + rmat[2][1]*src.rmat[1][0]
                  + rmat[2][2]*src.rmat[2][0];
   res.rmat[2][1] = rmat[2][0]*src.rmat[0][1]
                  + rmat[2][1]*src.rmat[1][1]
                  + rmat[2][2]*src.rmat[2][1];
   res.rmat[2][2] = rmat[2][0]*src.rmat[0][2]
                  + rmat[2][1]*src.rmat[1][2]
                  + rmat[2][2]*src.rmat[2][2];
   return res;
}

rmatrix& rmatrix::operator*=(const rmatrix& src)
{
   rmatrix res = (*this)*src;
   return *this = res;
}

rmatrix& rmatrix::operator=(const rmatrix& src)
{
   rmat[0][0] = src.rmat[0][0];
   rmat[1][0] = src.rmat[1][0];
   rmat[2][0] = src.rmat[2][0];
   rmat[0][1] = src.rmat[0][1];
   rmat[1][1] = src.rmat[1][1];
   rmat[2][1] = src.rmat[2][1];
   rmat[0][2] = src.rmat[0][2];
   rmat[1][2] = src.rmat[1][2];
   rmat[2][2] = src.rmat[2][2];
   return *this;
}

void rmatrix::dump()
{
   printf(" / %12.6f  %12.6f  %12.6f \\\n", rmat[0][0], rmat[0][1], rmat[0][2]);
   printf("|  %12.6f  %12.6f  %12.6f  |\n", rmat[1][0], rmat[1][1], rmat[1][2]);
   printf(" \\ %12.6f  %12.6f  %12.6f /\n", rmat[2][0], rmat[2][1], rmat[2][2]);
}

void finder(double k, xpeak_t peaks[2])
{
}

int main()
{
   xpeak_t inset[2];

   std::cout << " axisfinder: determine gonimeter settings that set up correct"
             << std::endl
             << " diffraction conditions for a given set of crystal planes in "
             << crystal << std::endl;
   std::cout << std::endl;
   std::cout << " 1) Input two sets of rocking curve peaks" 
             << " to determine offsets" << std::endl;
   std::cout << "    Enter Miller indices h k l for peak 1: ";
   std::cin >> inset[0].miller.h >> inset[0].miller.k >> inset[0].miller.l;
   std::cout << "    Enter gonimeter theta chi phi for peak 1: ";
   std::cin >> inset[0].goni.theta >> inset[0].goni.chi >> inset[0].goni.phi;
   std::cout << "    Enter Miller indices h k l for peak 2: "; 
   std::cin >> inset[1].miller.h >> inset[1].miller.k >> inset[1].miller.l;
   std::cout << "    Enter gonimeter theta chi phi for peak 2: ";
   std::cin >> inset[1].goni.theta >> inset[1].goni.chi >> inset[1].goni.phi;

   inset[0].goni.theta *= M_PI/180;
   inset[0].goni.chi *= M_PI/180;
   inset[0].goni.phi *= M_PI/180;
   inset[1].goni.theta *= M_PI/180;
   inset[1].goni.chi *= M_PI/180;
   inset[1].goni.phi *= M_PI/180;

 /* Solve this problem by expressing the offset matrix R that rotates
  * crystal coordinates into goniometer mount coordinates as follows.
  *
  *       R[i][j] = sum(a=1,3){ V[a][i] U[a][j] }
  *
  * where the rows U[a] are an orthogonal basis that rotates into a
  * second orthogonal basis contained in the rows V[a].  The first
  * reciprocal vector provided by the user gives a=1 after normalization.
  * The second provides a=2 after orthonormalization with the a=1 vectors.
  * The third ones are generated from the first two by completeness.
  */
   rmatrix U,V;
   rmatrix R,G;
   double q[4];
   double theta;

   q[0] = hbarc*(2*M_PI/aLattice);
   q[1] = q[0]*inset[0].miller.h;
   q[2] = q[0]*inset[0].miller.k;
   q[3] = q[0]*inset[0].miller.l;
   q[0] = sqrt(q[1]*q[1]+q[2]*q[2]+q[3]*q[3]);
   U(1,1) = q[1]/q[0];
   U(1,2) = q[2]/q[0];
   U(1,3) = q[3]/q[0];
   G.setGoniometer(inset[0].goni);
   theta = asin(q[0]/(2*kX));
   V(1,1) = G(1,1)*0 + G(2,1)*cos(theta) - G(3,1)*sin(theta);
   V(1,2) = G(1,2)*0 + G(2,2)*cos(theta) - G(3,2)*sin(theta);
   V(1,3) = G(1,3)*0 + G(2,3)*cos(theta) - G(3,3)*sin(theta);

   q[0] = hbarc*(2*M_PI/aLattice);
   q[1] = q[0]*inset[1].miller.h;
   q[2] = q[0]*inset[1].miller.k;
   q[3] = q[0]*inset[1].miller.l;
   q[0] = sqrt(q[1]*q[1]+q[2]*q[2]+q[3]*q[3]);
   U(2,1) = q[1]/q[0];
   U(2,2) = q[2]/q[0];
   U(2,3) = q[3]/q[0];
   G.setGoniometer(inset[1].goni);
   theta = asin(q[0]/(2*kX));
   V(2,1) = G(1,1)*0 + G(2,1)*cos(theta) - G(3,1)*sin(theta);
   V(2,2) = G(1,2)*0 + G(2,2)*cos(theta) - G(3,2)*sin(theta);
   V(2,3) = G(1,3)*0 + G(2,3)*cos(theta) - G(3,3)*sin(theta);

   double u1_u2 = U(1,1)*U(2,1) + U(1,2)*U(2,2) + U(1,3)*U(2,3);   
   double unorm = sqrt(1-u1_u2*u1_u2);
   U(2,1) = (U(2,1)-U(1,1)*u1_u2)/unorm;
   U(2,2) = (U(2,2)-U(1,2)*u1_u2)/unorm;
   U(2,3) = (U(2,3)-U(1,3)*u1_u2)/unorm;
   U(3,1) = U(1,2)*U(2,3)-U(1,3)*U(2,2);
   U(3,2) = U(1,3)*U(2,1)-U(1,1)*U(2,3);
   U(3,3) = U(1,1)*U(2,2)-U(1,2)*U(2,1);
   double v1_v2 = V(1,1)*V(2,1) + V(1,2)*V(2,2) + V(1,3)*V(2,3);   
   double vnorm = sqrt(1-v1_v2*v1_v2);
   V(2,1) = (V(2,1)-V(1,1)*v1_v2)/vnorm;
   V(2,2) = (V(2,2)-V(1,2)*v1_v2)/vnorm;
   V(2,3) = (V(2,3)-V(1,3)*v1_v2)/vnorm;
   V(3,1) = V(1,2)*V(2,3)-V(1,3)*V(2,2);
   V(3,2) = V(1,3)*V(2,1)-V(1,1)*V(2,3);
   V(3,3) = V(1,1)*V(2,2)-V(1,2)*V(2,1);
   
   std::cout << "    Check equality of dot products: "
             << u1_u2 << ", " << v1_v2 << std::endl;

   for (int i=1; i<4; i++) {
      for (int j=1; j<4; j++) {
         R(i,j) = V(1,i)*U(1,j)+V(2,i)*U(2,j)+V(3,i)*U(3,j);
      }
   }

   gset_t oset = (R.inverse()).getGoniometer();
   std::cout << std::endl << "  Results for goniometer offsets:"
             << std::endl << "     theta=" << oset.theta*(180/M_PI)
             << std::endl << "     chi=" << oset.chi*(180/M_PI)
             << std::endl << "     phi=" << oset.phi*(180/M_PI)
             << std::endl;

   std::cout << std::endl << " 2) Now I will predict goniometer settings"
                          << " for additional peaks" << std::endl;
   for (int step=3; step<9999; step++) {
      xpeak_t tset;
      std::cout << "    Enter Miller indices h k l for peak " << step << ": ";
      std::cin >> tset.miller.h >> tset.miller.k >> tset.miller.l;
      q[0] = hbarc*(2*M_PI/aLattice);
      q[1] = q[0]*tset.miller.h;
      q[2] = q[0]*tset.miller.k;
      q[3] = q[0]*tset.miller.l;
      q[0] = sqrt(q[1]*q[1]+q[2]*q[2]+q[3]*q[3]);
      theta = asin(q[0]/(2*kX));

      rmatrix Q,H,S;
      Q.setRotateX(theta);
      H.setVertical(tset.miller);
      double wobble=1e-30;
      while (wobble != 0) {
         S.setRotateY(wobble*M_PI/180);
         S *= H;
         oset = (Q * S * R.inverse()).getGoniometer();
         std::cout << "    The peak should appear at or near:"
                << std::endl << "     twotheta=" << 2*theta*(180/M_PI)
                << std::endl << "     theta=" << oset.theta*(180/M_PI)
                << std::endl << "     chi=" << oset.chi*(180/M_PI)
                << std::endl << "     phi=" << oset.phi*(180/M_PI)
                << std::endl;
         std::cout << "    Enter a wobble angle (or 0 to quit): ";
         std::cin >> wobble;
      }
   }
}