//
// genPXR class implementation
//
// author: richard.t.jones at uconn.edu
// version: september 2, 2021
//

#include "genpxr.h"

#include <math.h>
#include <exception>
#include <iostream>
#include <cmath>

#include <boost/python.hpp>
#include <boost/python/suite/indexing/map_indexing_suite.hpp>
#include <boost/python/suite/indexing/vector_indexing_suite.hpp>

const double genPXR::mElectron_keV(510.99891);
const double genPXR::alphaQED(7.2973525698e-3);
const double genPXR::hbarc_keV_m(0.1973269718e-9);
const double genPXR::cLight_m__s(2.99792e8);
const double genPXR::BohrRadius_m(hbarc_keV_m / (alphaQED * mElectron_keV));

genPXR::genPXR(double Ebeam_GeV) {

   // assign constants for the diamond lattice
   latticeConst_m.push_back(3.5668e-10);
   latticeConst_m.push_back(3.5668e-10);
   latticeConst_m.push_back(3.5668e-10);
   atomsPerUnitCell.push_back(8);
   atomicZ.push_back(6);
   delta_10keV = 4.6e-6;  // parameter for computing refractive index, see refIndex

   // initialize the crystal orientation
   crystalPhiAlphaXY[0] = 0;
   crystalPhiAlphaXY[1] = 0;
   crystalPhiAlphaXY[2] = 0;
   aLattice[0][0] = aLattice[1][1] = aLattice[2][2] = 1;
   aLattice[0][1] = aLattice[0][2] = aLattice[1][2] = 0;
   aLattice[1][0] = aLattice[2][0] = aLattice[2][1] = 0;

   // beam properties
   Ebeam_keV = Ebeam_GeV * 1e6;
   betabeam[0] = betabeam[1] = 0;
   betabeam[2] = sqrt(pow(Ebeam_keV, 2) - pow(mElectron_keV, 2)) / Ebeam_keV;

   // default detector frequency filter parameters
   frequency_filter_limits_keV[0] = 10;
   frequency_filter_limits_keV[1] = 100;
}

genPXR::~genPXR() {
}

double genPXR::FFatomic(double q_keV) {
   // Simple parameterization for the charge form factor
   // of the electron cloud around an atom, normalized to
   // unity at q=0.
   double FF(0);
   double Zsum(0);
   for (auto Z : atomicZ) {
      double betaq = 111 / pow(Z, 1/3.) * q_keV / mElectron_keV;
      FF += Z / (1 + betaq*betaq);
      Zsum += Z;
   }
   return FF / Zsum;
}

double genPXR::FFatomic_Lobato(double q_keV) {
   // Lobato parameterization for the charge form factor
   // of the electron cloud around a carbon atom, 
   // normalized to unity at q=0.
   double Z(6);
   double q__Ang = q_keV / (2 * M_PI * hbarc_keV_m * 1e10);
   double apar_Ang[5] = { 1.244660886213433e+02,
                         -2.203528570789638e+02,
                          1.952353522804791e+02,
                         -9.810793612697997e+01,
                          1.420230412136232e-02};
   double bpar_Ang2[5] = {2.421208492560056,
                          2.305379437524258,
                          2.048519321065642,
                          1.933525529175474,
                          7.689768184783397e-02};
   double FF_Ang(0);
   for (int i=0; i < 5; ++i) {
      FF_Ang += apar_Ang[i] * ((2 + bpar_Ang2[i] * q__Ang*q__Ang) /
                            pow(1 + bpar_Ang2[i] * q__Ang*q__Ang, 2));
   }
   return 1 - 2 * pow(M_PI * q__Ang, 2) * (BohrRadius_m * 1e10) * FF_Ang / Z;
}

std::vector<double> genPXR::crystalWavenumber(const int hkl[3]) {
   // Return the crystal wavenumber (m^-1) for the Miller indices hkl

   std::vector<double> h_hkl;
   h_hkl.resize(3);
   for (int i=0; i < 3; ++i)
      h_hkl[i] = 2 * M_PI *
                 (hkl[0] * aLattice[0][i] / latticeConst_m[0] +
                  hkl[1] * aLattice[1][i] / latticeConst_m[1] +
                  hkl[2] * aLattice[2][i] / latticeConst_m[2]);
   return h_hkl;
}

double genPXR::omega_keV(const int hkl[3], const std::vector<double> &khat) {
   // Solve for the energy (keV) of pxr radiation in direction khat
   // from lattice vector corresponding to Miller indices hkl. This
   // requires an iterative algorithm, can take some time.

   double min_omega_keV(1.0);
   std::vector<double> h_hkl(crystalWavenumber(hkl));
   double ome_keV(10);
   double ome_old(0);
   int count(0);
   while (fabs(ome_keV - ome_old) > fabs(ome_keV) * 1e-9) {
      ome_old = ome_keV;
      double h_dot_beta = h_hkl[0] * betabeam[0] + 
                          h_hkl[1] * betabeam[1] + 
                          h_hkl[2] * betabeam[2];
      double n_dot_beta = khat[0] * betabeam[0] + 
                          khat[1] * betabeam[1] + 
                          khat[2] * betabeam[2];
      ome_keV = h_dot_beta * hbarc_keV_m /
               (1 - n_dot_beta * refIndex(ome_keV));
      if (++count > 1000) {
         std::cerr << "error in omega_keV - "
                   << "iteration loop failed to converge" << std::endl;
         ome_keV = 0;
         break;
      }
   }
   if (ome_keV < min_omega_keV) {
      ome_keV = 0;
   }
   else {
      double kstar = ome_keV * refIndex(ome_keV);
      double k_hkl[3] = {h_hkl[0] * hbarc_keV_m + khat[0] * kstar,
                         h_hkl[1] * hbarc_keV_m + khat[1] * kstar,
                         h_hkl[2] * hbarc_keV_m + khat[2] * kstar};
      double ome_check = k_hkl[0] * betabeam[0] +
                         k_hkl[1] * betabeam[1] +
                         k_hkl[2] * betabeam[2];
      if (fabs(ome_keV - ome_check) > abs(ome_keV) * 1e-9) {
         std::cout << "bad ome_check: " << ome_keV << " != " << ome_check
                   << std::endl;
         throw std::runtime_error("error in omega_keV - bad omega check");
      }
   }
   return ome_keV;
}

double genPXR::Esuscept(const int hkl[3], double ome_keV) {
   // Return the dimensionless susceptibility Fourier coefficient
   // at frequency ome_keV corresponding to Miller indices hkl.

#if USE_FREE_ELECTRON_SUSCEPTIBILITY
   double rho0(0);
   for (unsigned i=0; i < atomicZ.size(); ++i) {
      rho0 += atomsPerUnitCell[i] * atomicZ[i];
   }
   rho0 /= latticeConst_m[0] * latticeConst_m[1] * latticeConst_m[2];
   double chi0 = -4 * M_PI * alphaQED * rho0 * pow(hbarc_keV_m, 3);
   chi0 /= mElectron_keV * ome_keV*ome_keV;
#else
   double chi0 = pow(refIndex(ome_keV), 2) - 1;
#endif
   std::vector<double> h_hkl(crystalWavenumber(hkl));
   double q_keV = sqrt(h_hkl[0] * h_hkl[0] + 
                       h_hkl[1] * h_hkl[1] +
                       h_hkl[2] * h_hkl[2]) * hbarc_keV_m;
   if (q_keV == 0)
      return chi0;
   else if ((hkl[0] % 2) != 0 and (hkl[1] % 2) != 0 and (hkl[2] % 2) != 0)
      return chi0 * FFatomic(q_keV) / sqrt(2);
   else if ((hkl[0] % 2) == 0 and (hkl[1] % 2) == 0 and (hkl[2] % 2) == 0) {
      if ((hkl[0] + hkl[1] + hkl[2]) % 4 == 0)
         return chi0 * FFatomic(q_keV);
   }
   return 0;
}

double genPXR::refIndex(double ome_keV) {
   // Refractive index of diamond, "X-ray Scattering at Grazing Incidence",
   // https://www.classe.cornell.edu/~dms79/refl/OldStyle/X-rayScattering.html
   // by Detlef Smilgies, CLASSE facility at CHESS.

   return 1 - delta_10keV * (100 / (ome_keV*ome_keV + 1e-99));
}

void genPXR::setCrystal(double phideg, double alphax, double alphay, int verbose)
{
   // Set the orientation of the radiator crystal, starting with the
   // horizontal beam pointing along the crystal z axis, the crystal
   // y axis pointing straight up, and the crystal x axis in the
   // horizontal plane pointing beam-left, and then rotated in a
   // sequence of three operations in the following order.
   //  1) rotate the crystal around the z axis by phideg degrees,
   //     with positive rotation defined as x moving toward y.
   //  2) rotate the crystal around the x' axis by alpha_x radians,
   //     with positive rotation defined as y' moving toward z'==z.
   //  3) rotate the cyrstal around the y'' axis by alpha_y radians,
   //     with positive rotation defined as z'' moving toward x''.

   if (phideg != 999)
      crystalPhiAlphaXY[0] = phideg * M_PI/180;
   if (alphax != 999)
      crystalPhiAlphaXY[1] = alphax;
   if (alphay != 999)
      crystalPhiAlphaXY[2] = alphay;
   double cosphi = cos(crystalPhiAlphaXY[0]);
   double sinphi = sin(crystalPhiAlphaXY[0]);
   double cosalx = cos(crystalPhiAlphaXY[1]);
   double sinalx = sin(crystalPhiAlphaXY[1]);
   double cosaly = cos(crystalPhiAlphaXY[2]);
   double sinaly = sin(crystalPhiAlphaXY[2]);
   aLattice[0][0] = cosaly * cosphi + sinaly * sinalx * sinphi;
   aLattice[0][1] = cosalx * sinphi;
   aLattice[0][2] = -sinaly * cosphi + cosaly * sinalx * sinphi;
   aLattice[1][0] = -cosaly * sinphi + sinalx * sinaly * cosphi;
   aLattice[1][1] = cosalx * cosphi;
   aLattice[1][2] = sinaly * sinphi + cosaly * sinalx * cosphi;
   aLattice[2][0] = cosalx * sinaly;
   aLattice[2][1] = -sinalx;
   aLattice[2][2] = cosaly * cosalx;
   double a_aTranspose[3][3];
   for (int i=0; i < 3; ++i) {
      for (int j=i; j < 3; ++j) {
         a_aTranspose[i][j] = aLattice[i][0] * aLattice[j][0] +
                              aLattice[i][1] * aLattice[j][1] +
                              aLattice[i][2] * aLattice[j][2];
      }
   }
   if (fabs(a_aTranspose[0][0] - 1) > 1e-15 ||
       fabs(a_aTranspose[1][1] - 1) > 1e-15 ||
       fabs(a_aTranspose[2][2] - 1) > 1e-15 ||
       fabs(a_aTranspose[0][1]) > 1e-15 ||
       fabs(a_aTranspose[0][2]) > 1e-15 ||
       fabs(a_aTranspose[1][2]) > 1e-15)
   {
      throw std::runtime_error("setCrystal error - "
                 "crystal rotation matrix is not orthogonal!");
   }
   else if (verbose > 0) {
      std::cout << "setCrystal - crystal orientation is phi="
                << crystalPhiAlphaXY[0] * 180/M_PI << "deg, alphax="
                << crystalPhiAlphaXY[1] << "rad, alphay="
                << crystalPhiAlphaXY[2] << "rad" << std::endl;
   }
}

double genPXR::frequency_filter(double ome_keV) {
  // Return the detector efficiency for PXR radiation of a given energy.

   if (ome_keV > frequency_filter_limits_keV[0] &&
       ome_keV < frequency_filter_limits_keV[1])
   {
      return 1;
   }
   return 0;
}

void genPXR::setFilterLimits(double ome_low_limit_keV, double ome_high_limit_keV) {
   frequency_filter_limits_keV[0] = ome_low_limit_keV;
   frequency_filter_limits_keV[1] = ome_high_limit_keV;
}

double genPXR::dNdOmegadz(double thetadeg, double phideg, int verbosity) {
   // Return the pxr radiation rate per high-energy electron in units
   // of per unit solid angle per unit length of radiator (m^-1),
   // calculated using Eq. 17 of Nitta, Physics Letters A 158 (1991) 270.
   // The formula written at Eq. 17 in the paper has errors that are
   // corrected by going back and substituting Eq. 12 into Eq. 16.
 
   double costheta = cos(thetadeg * M_PI/180);
   double sintheta = sin(thetadeg * M_PI/180);
   double cosphi = cos(phideg * M_PI/180);
   double sinphi = sin(phideg * M_PI/180);
   std::vector<double> khat = {sintheta * cosphi, sintheta * sinphi, costheta};
   std::vector<double> epol = {costheta * cosphi, costheta * sinphi, -sintheta};
   std::vector<double> epol2 = {sinphi, -cosphi, 0};
   int hkl000[3] = {0,0,0};
   double dNsum = 0;
   for (int h = -10; h <= 10; ++h) {
    for (int k = -10; k <= 10; ++k) {
     for (int l = -10; l <= 10; ++l) {
        if (h == 0 && k == 0 && l == 0)
           continue;
        int hkl[3] = {h,k,l};
        double ome_keV = omega_keV(hkl, khat);
        if (std::isnan(ome_keV)) {
           std::cout << "[" << h << "," << k << "," << l 
                     << "] gives ome_keV = " << ome_keV 
                     << std::endl;
           return 0;
        }
        if (ome_keV <= 0)
           continue;
        double eps0 = 1 + Esuscept(hkl000, ome_keV);
        double chi_hkl = Esuscept(hkl, ome_keV);
        if (chi_hkl == 0)
           continue;
        double nrefr = refIndex(ome_keV);
        double kstar = ome_keV * nrefr / hbarc_keV_m;
        std::vector<double> h_hkl(crystalWavenumber(hkl));
        double k_hkl[3];
        double pprime[3];
        double pprime_norm(0);
        double k_hkl_norm(0);
        for (int i=0; i < 3; ++i) {
           k_hkl[i] = kstar * khat[i] + h_hkl[i];
           k_hkl_norm += k_hkl[i] * k_hkl[i];
           pprime[i] = Ebeam_keV * betabeam[i] - k_hkl[i] * hbarc_keV_m;
           pprime_norm += pprime[i] * pprime[i];
        }
        double Eprime_keV = sqrt(pprime_norm + mElectron_keV*mElectron_keV);
        double betaprime[3] = {pprime[0] / Eprime_keV,
                               pprime[1] / Eprime_keV,
                               pprime[2] / Eprime_keV};
        double n_dot_beta = khat[0] * betabeam[0] +
                            khat[1] * betabeam[1] +
                            khat[2] * betabeam[2];
        double dN = alphaQED * (Eprime_keV / Ebeam_keV) *
                    pow(kstar, 3) * pow(chi_hkl, 2) /
                    (2*M_PI * pow(eps0, 3) * betabeam[2] *
                     (1 - n_dot_beta * nrefr -
                      ome_keV / Ebeam_keV * (1 - eps0)) *
                     pow(k_hkl_norm - kstar*kstar, 2));
        double dot1a = (kstar * nrefr * betabeam[0] - h_hkl[0]) * epol[0] +
                       (kstar * nrefr * betabeam[1] - h_hkl[1]) * epol[1] +
                       (kstar * nrefr * betabeam[2] - h_hkl[2]) * epol[2];
        double dot1b = (kstar * nrefr * betabeam[0] - h_hkl[0]) * epol2[0] +
                       (kstar * nrefr * betabeam[1] - h_hkl[1]) * epol2[1] +
                       (kstar * nrefr * betabeam[2] - h_hkl[2]) * epol2[2];
        double kdotdb = k_hkl[0] * (betabeam[0] - betaprime[0]) +
                        k_hkl[1] * (betabeam[1] - betaprime[1]) +
                        k_hkl[2] * (betabeam[2] - betaprime[2]);
        double dot2a = (kstar * nrefr * (betabeam[0] - betaprime[0])) * epol[0] +
                       (kstar * nrefr * (betabeam[1] - betaprime[1])) * epol[1] +
                       (kstar * nrefr * (betabeam[2] - betaprime[2])) * epol[2] -
                       kdotdb * hbarc_keV_m / ome_keV * h_hkl[0] * epol[0] -
                       kdotdb * hbarc_keV_m / ome_keV * h_hkl[1] * epol[1] -
                       kdotdb * hbarc_keV_m / ome_keV * h_hkl[2] * epol[2];
        double dot2b = (kstar * nrefr * (betabeam[0] - betaprime[0])) * epol2[0] +
                       (kstar * nrefr * (betabeam[1] - betaprime[1])) * epol2[1] +
                       (kstar * nrefr * (betabeam[2] - betaprime[2])) * epol2[2] -
                       kdotdb * hbarc_keV_m / ome_keV * h_hkl[0] * epol2[0] -
                       kdotdb * hbarc_keV_m / ome_keV * h_hkl[1] * epol2[1] -
                       kdotdb * hbarc_keV_m / ome_keV * h_hkl[2] * epol2[2];
        double dot3 = (betabeam[0] - betaprime[0]) * betabeam[0] +
                      (betabeam[1] - betaprime[1]) * betabeam[1] +
                      (betabeam[2] - betaprime[2]) * betabeam[2];
        double dot4a = h_hkl[0] * epol[0] +
                       h_hkl[1] * epol[1] +
                       h_hkl[2] * epol[2];
        double dot4b = h_hkl[0] * epol2[0] +
                       h_hkl[1] * epol2[1] +
                       h_hkl[2] * epol2[2];
        double dot5a = pow(dot4a * k_hkl[0] - kstar*kstar * epol[0], 2) +
                       pow(dot4a * k_hkl[1] - kstar*kstar * epol[1], 2) +
                       pow(dot4a * k_hkl[2] - kstar*kstar * epol[2], 2);
        double dot5b = pow(dot4b * k_hkl[0] - kstar*kstar * epol2[0], 2) +
                       pow(dot4b * k_hkl[1] - kstar*kstar * epol2[1], 2) +
                       pow(dot4b * k_hkl[2] - kstar*kstar * epol2[2], 2);
        dN *= dot1a * dot1a - dot1a * dot2a + 
              0.5 * (pow(mElectron_keV, 2) / pow(Ebeam_keV, 2) +
                     pow(mElectron_keV, 2) / (Ebeam_keV * Eprime_keV) - dot3) *
                     dot5a / pow(ome_keV / hbarc_keV_m, 2)
            + dot1b * dot1b - dot1b * dot2b + 
              0.5 * (pow(mElectron_keV, 2) / pow(Ebeam_keV, 2) +
                     pow(mElectron_keV, 2) / (Ebeam_keV * Eprime_keV) - dot3) *
                     dot5b / pow(ome_keV / hbarc_keV_m, 2);
        if (dN <= 0)
           continue;
        dN *= 1e-3; // normalize to a 1mm thick diamond target
        dNsum += dN * frequency_filter(ome_keV);
        // dN is now in radiated photons per electron / sr of radiated solid angle
        if (verbosity > 0 && frequency_filter(ome_keV) > 0 && dN > 1) {
           std::cout << "vector [" << h << "," << k << "," << l << "] radiates "
                     << ome_keV << "keV with emission rate " << dN << std::endl;
        }
     }
    }
   }
   return dNsum;
}

double genPXR::dNdOmegadz2(double thetadeg, double phideg, int verbosity) {
   // Return the pxr radiation rate per high-energy electron in units
   // of per unit solid angle per unit length of radiator (m^-1),
   // calculated using Eq. 19 of Nitta, Physics Letters A 158 (1991) 270.
   // This is provided as a consistency check with dNdOmegadz.

   double costheta = cos(thetadeg * M_PI/180);
   double sintheta = sin(thetadeg * M_PI/180);
   double cosphi = cos(phideg * M_PI/180);
   double sinphi = sin(phideg * M_PI/180);
   std::vector<double> khat = {sintheta * cosphi, sintheta * sinphi, costheta};
   std::vector<double> epol = {costheta * cosphi, costheta * sinphi, -sintheta};
   std::vector<double> epol2 = {sinphi, -cosphi, 0};
   int hkl000[3] = {0,0,0};
   double dNsum = 0;
   for (int h = -10; h <= 10; ++h) {
    for (int k = -10; k <= 10; ++k) {
     for (int l = -10; l <= 10; ++l) {
        if (h == 0 && k == 0 && l == 0)
           continue;
        int hkl[3] = {h,k,l};
        double ome_keV = omega_keV(hkl, khat);
        if (ome_keV <= 0)
           continue;
        double eps0 = 1 + Esuscept(hkl000, ome_keV);
        double chi_hkl = Esuscept(hkl, ome_keV);
        if (chi_hkl == 0)
           continue;
        double nrefr = refIndex(ome_keV);
        double kstar = ome_keV * nrefr / hbarc_keV_m;
        std::vector<double> h_hkl(crystalWavenumber(hkl));
        double n_dot_beta = khat[0] * betabeam[0] +
                            khat[1] * betabeam[1] +
                            khat[2] * betabeam[2];
        double dN = alphaQED * pow(kstar, 3) * pow(chi_hkl, 2) /
                   (2*M_PI * pow(eps0, 3) * betabeam[2] *
                    (1 - n_dot_beta * nrefr) *
                    pow(pow(kstar * khat[0] + h_hkl[0], 2) +
                        pow(kstar * khat[1] + h_hkl[1], 2) +
                        pow(ome_keV / (betabeam[2] * hbarc_keV_m), 2) *
                        (pow(mElectron_keV / Ebeam_keV, 2) + 
                         pow(betabeam[2], 2) * (1 - eps0)), 2));
        double dot1a = (kstar * betabeam[0] - h_hkl[0]) * epol[0] +
                       (kstar * betabeam[1] - h_hkl[1]) * epol[1] +
                       (kstar * betabeam[2] - h_hkl[2]) * epol[2];
        double dot1b = (kstar * betabeam[0] - h_hkl[0]) * epol2[0] +
                       (kstar * betabeam[1] - h_hkl[1]) * epol2[1] +
                       (kstar * betabeam[2] - h_hkl[2]) * epol2[2];
        dN *= dot1a * dot1a + dot1b * dot1b;
        if (dN <= 0)
           continue;
        dN *= 1e-3; // normalize to a 1mm thick diamond target
        dNsum += dN * frequency_filter(ome_keV);
        // dN is now in radiated photons per electron / sr of radiated solid angle
        if (verbosity > 0 && frequency_filter(ome_keV) > 0 && dN > 1) {
           std::cout << "vector [" << h << "," << k << "," << l << "] radiates "
                     << ome_keV << "keV with emission rate " << dN << std::endl;
        }
     }
    }
   }
   return dNsum;
}

// Create a python module containing the genPXR class methods.

BOOST_PYTHON_MODULE(libgenpxr)
{
   using boost::python::class_;
   using boost::python::enum_;
   using boost::python::def;

   class_<genPXR, genPXR*>
         ("genPXR",
          "class for computing rates for parametric Xray radiation"
          " by high-energy electrons")
   ;
}