//
// phonoCbrems.C
//
// Generates single phonons weighted by their coupling,
// and uses this generator to compute the spectrum and
// polarization of coherent bremsstrahlung from the
// single-phonon coherent bremsstrahlung process.
//
// author: richard.t.jones at uconn.edu
// version: february 13, 2017

#include "Complex.h"
#include "TPhoton.h"
#include "TLepton.h"
#include "TCrossSection.h"
#include "TThreeVectorReal.h"
#include "TFourVectorReal.h"
#include "TLorentzBoost.h"
#include "constants.h"

#include <TROOT.h>
#include <TCanvas.h>
#include <TF1.h>
#include <TF2.h>
#include <TH1D.h>
#include <TTree.h>
#include <TAxis.h>
#include <TFile.h>
#include <TRandom3.h>
#include <math.h>

TCanvas *c1 = 0;

#ifndef DEFINE_SQR_ON_STANDARD_TYPES
#define DEFINE_SQR_ON_STANDARD_TYPES
inline unsigned int sqr(unsigned int x) { return x*x; }
inline Int_t sqr(Int_t x) { return x*x; }
inline Float_t sqr(Float_t x) { return x*x; }
inline Double_t sqr(Double_t x) { return x*x; }
inline LDouble_t sqr(LDouble_t x) { return x*x; }
inline Complex_t sqr(Complex_t x) { return x*x; }
#endif

#ifndef STANDARD_VECTOR_CONSTANTS
#define STANDARD_VECTOR_CONSTANTS
const TThreeVectorReal zeroVector(0,0,0);
const TThreeVectorReal posXhat(1,0,0);
const TThreeVectorReal negXhat(-1,0,0);
const TThreeVectorReal posYhat(0,1,0);
const TThreeVectorReal negYhat(0,-1,0);
const TThreeVectorReal posZhat(0,0,1);
const TThreeVectorReal negZhat(0,0,-1);
#endif

#define IMPORTANCE_SAMPLING_HIST_FILE "phonon_weights.root"

// parameters for the beamline
double Ebeam = 12.;     // GeV
double radcoldist = 76.; // m
double pbeam = sqrt(sqr(Ebeam) - sqr(mElectron));

// physical parameters for the target crystal
const double Eedge = 9.;      // GeV
const double qtotal_220 = 9.8e-6; // GeV
const double qlong_220 = Eedge / (Ebeam - Eedge) * 
                         sqr(mElectron) / (2 * Ebeam);
const double aLattice = 3.567e-10;   // m
const double hbarc = 1.97327e-16;    // GeV.m
const double kBoltzmann = 8.617e-14; // GeV/K
const double kT = 8.617e-14 * 300;   // GeV
const double TDebye = 2200;          // K
const double Ncell = 8;
const double Vcell = 45.5e-30;       // m^3
const double Matom = 12 * 0.932;     // GeV/c^2
const double ThetaD = 8.617e-14 * 2200; // GeV
const double Qmax = 2 * M_PI * hbarc * 
                    pow(3 / (4 * M_PI) * Ncell / Vcell, 1/3.); // GeV/c
const double Bphonons = 0.561;      // dimensionless

TH1D *fImportSample[2] = {0, 0};
TRandom3 *fRandom = new TRandom3(0);

void normalize(TH1D *hist)
{
   double sum = hist->Integral();
   double dx = hist->GetBinWidth(1);
   double normfactor = 1 / (sum * dx);
   int nbins = hist->GetNbinsX();
   for (int i=1; i <= nbins; ++i) {
      hist->SetBinContent(i, hist->GetBinContent(i) * normfactor);
   }
}

void prepare_sampling()
{
   TFile samples(IMPORTANCE_SAMPLING_HIST_FILE);
   fImportSample[0] = (TH1D*)samples.Get("x_weight");
   bool found = true;
   for (int i=0; i < 1; ++i) {
      if (fImportSample[i]) {
         fImportSample[i]->SetDirectory(0);
         normalize(fImportSample[i]);
      }
      else {
         found = false;
      }
   }
   if (found) {
      std::cout << "importance sampling histograms read from "
                << IMPORTANCE_SAMPLING_HIST_FILE
                << std::endl;
   }
}

double sample_hkl(int hkl[3])
{
   // Sample the reciprocal lattice site and return the weight.

   static double Niterations = 0;
   static int Nentries = 0;

   // bootstrap the normalization
   if (Nentries == 0) {
      Nentries = 1;
      for (int n=0; n < 1000000; ++n) {
         sample_hkl(hkl);
      }
   }

   ++Nentries;
   while (++Niterations) {
      const double hexp = 20.;
      const double kexp = 20.;
      const double lexp = 1.5;
      const double hexp_factor = 1 - exp(-1/hexp);
      const double kexp_factor = 1 - exp(-1/kexp);
      const double lexp_factor = 1 - exp(-1/lexp);
      hkl[0] = floor(-hexp * log(fRandom->Uniform()));
      hkl[1] = floor(-kexp * log(fRandom->Uniform()));
      hkl[2] = floor(-lexp * log(fRandom->Uniform()));
      double wgt = 8;
      wgt *= exp(hkl[0] / hexp) / hexp_factor;
      wgt *= exp(hkl[1] / kexp) / kexp_factor;
      wgt *= exp(hkl[2] / lexp) / lexp_factor;
      int sector = floor(8 * fRandom->Uniform());
      if (sector & 1) {
         hkl[0] = -hkl[0] - 1;
      }
      if (sector & 2) {
         hkl[1] = -hkl[1] - 1;
      }
      if (sector & 4) {
         hkl[2] = -hkl[2] - 1;
      }
      if (hkl[0] / 2 * 2 == hkl[0]) {
         int hklsum = hkl[0] + hkl[1] + hkl[2];
         if (hkl[1] / 2 * 2 != hkl[1] ||
             hkl[2] / 2 * 2 != hkl[2] ||
             hklsum / 4 * 4 != hklsum)
         {
            continue;
         }
      }
      else if (hkl[1] / 2 * 2 == hkl[1] || 
               hkl[2] / 2 * 2 == hkl[2])
      {
        continue;
      }
      return wgt * Nentries / Niterations;
   }
   return 0;
}

double sample_x(TThreeVectorReal &x)
{
   // Sample the phonon momentum parameter x from the
   // phonon spectrum, and return the weight.

   double xmag;
   double wgt;
   if (fImportSample[0]) {
      xmag = fImportSample[0]->GetRandom();
      int bin = fImportSample[0]->FindBin(xmag);
      wgt = 1 / fImportSample[0]->GetBinContent(bin);
   }
   else {
      xmag = sqrt(fRandom->Uniform());
      wgt = 1 / (2 * xmag);
   }
   double phi = fRandom->Uniform() * 2*M_PI;
   double theta = acos(2 * fRandom->Uniform() - 1);
   x.SetPolar(xmag, theta, phi);
   double nx = 0;
   if (xmag > 1e-6 && kT > 0) {
      nx = 1 / (exp(ThetaD * xmag / kT) - 1);
   }
   else {
      nx = kT / (ThetaD * xmag);
   }
   wgt *= xmag * (2 * nx + 1);
   return wgt;
}

double sample_Nphonons(int &Nphonons)
{
   // Sample the number of phonons and return the weight

   int Nmax = 10;
   Nphonons = floor(fRandom->Uniform() * Nmax);
   return Nmax;
}

int gentree(int nsamples, TString rootfile)
{
   // Generate coherent bremsstrahlung photons from a single 
   // reciprocal lattice vector, where the gamma radiation is
   // accompanied by a phonons radiated or absorbed in the
   // radiator crystal assumed to be in equilibrium at a fixed
   // temperature T. The results are saved in a root tree.

   if (fImportSample[0] == 0) {
      prepare_sampling();
   }

   // Define the vectors that orient the crystal as follows:
   //  1. (2,2,0) lies in the x-z plane with a small negative z component
   //  2. (-2,2,0) points close to the y axis with a larger -z component
   //  3. (0,0,1) points close the the z axis (beam direction)

   double tiltV = asin(qlong_220 / qtotal_220);
   double tiltH = 0.050; // radians
   double qUnit = 2*M_PI * hbarc / aLattice;
   TThreeVectorReal qAxis[3];
   qAxis[0].SetPolar(qUnit, M_PI/2, 0);
   qAxis[1].SetPolar(qUnit, M_PI/2, M_PI/2);
   qAxis[2].SetPolar(qUnit, 0, 0);
   for (int i=0; i < 3; ++i) {
      qAxis[i].Rotate(posZhat, +M_PI/4);
#define DO_PERP 1
#if DO_PERP
      qAxis[i].Rotate(posYhat, -tiltH);
      qAxis[i].Rotate(posXhat, +tiltV);
#else
      qAxis[i].Rotate(posXhat, +tiltH);
      qAxis[i].Rotate(posYhat, -tiltV);
#endif
   }
   TThreeVectorReal q220(2 * qAxis[0] + 2 * qAxis[1]);
   TThreeVectorReal q_220(-2 * qAxis[0] + 2 * qAxis[1]);
   std::cout << "direction (2,2,0) is "; (q220 / q220.Length()).Print();
   std::cout << "direction (-2,2,0) is "; (q_220 / q_220.Length()).Print();

   TFile treefile(rootfile, "recreate");

   struct treerow {
      double qphoton[3];
      double Qphonons[3];
      double diffXS;
      double polar;
      double k[3];
      double coupling;
      double wgt;
      int hkl[3];
      int Nphonons;
   } row;

   TObject *o = gROOT->FindObject("phono");
   if (o)
      o->Delete();
   TTree tree("phono", "coherent bremsstrahlung one-phonon process");
   tree.Branch("qphoton", &row.qphoton[0], "qphoton[3]/D");
   tree.Branch("Qphonons", &row.Qphonons[0], "Qphonons[3]/D");
   tree.Branch("diffXS", &row.diffXS, "diffXS/D");
   tree.Branch("polar", &row.polar, "polar/D");
   tree.Branch("k", &row.k[0], "k[3]/D");
   tree.Branch("coupling", &row.coupling, "coupling/D");
   tree.Branch("wgt", &row.wgt, "wgt/D");
   tree.Branch("hkl", &row.hkl[0], "hkl[3]/I");
   tree.Branch("Nphonons", &row.Nphonons, "Nphonons/I");

   TLepton eIn(mElectron), eOut(mElectron);
   TPhoton gOut;

   TThreeVectorReal uhat;
   for (int sample=0; sample < nsamples; ++sample) {
      row.wgt = sample_hkl(row.hkl);
      TThreeVectorReal qLattice;
      qLattice = row.hkl[0] * qAxis[0] + 
                 row.hkl[1] * qAxis[1] +
                 row.hkl[2] * qAxis[2];
      int Nphonons;
      row.wgt *= sample_Nphonons(Nphonons);
      TThreeVectorReal Qphonons(0,0,0);
      for (int phonon=0; phonon < Nphonons; ++phonon) {
         TThreeVectorReal Qphonon(0,0,0);
         row.wgt *= sample_x(Qphonon);
         Qphonons += Qmax * Qphonon;
         row.wgt /= phonon + 1;
      }
      row.Nphonons = Nphonons;
      row.Qphonons[0] = Qphonons[1];
      row.Qphonons[1] = Qphonons[2];
      row.Qphonons[2] = Qphonons[3];
      TFourVectorReal Qphoton(0, qLattice - Qphonons);
      row.qphoton[0] = Qphoton[1];
      row.qphoton[1] = Qphoton[2];
      row.qphoton[2] = Qphoton[3];
      row.wgt *= pow(3 * Qphoton.LengthSqr() / 
                    (2 * Matom * ThetaD), Nphonons);
      row.wgt *= exp(-3 * qLattice.LengthSqr() / 
                    (2 * Matom * ThetaD) * Bphonons);

      TFourVectorReal pin(Ebeam, 0, 0, pbeam);
      TFourVectorReal ptot(pin + Qphoton);
      double M2 = ptot.InvariantSqr();
      if (M2 < sqr(mElectron)) {
         continue;
         row.k[0] = 0;
         row.k[1] = 0;
         row.k[2] = 0;
         row.diffXS = 0;
         row.polar = 0;
         tree.Fill();
      }
      double kstar = (M2 - sqr(mElectron)) / (2 * sqrt(M2));
      double phistar = fRandom->Uniform() * 2*M_PI;
      double thetastar = acos(fRandom->Uniform() * 2 - 1);
      uhat.SetPolar(1, thetastar, phistar);
      TFourVectorReal kout(kstar, kstar * uhat);
      TThreeVectorReal beta(-(TThreeVector)ptot / ptot[0]);
      kout.Boost(beta);
      if (kout[0] < 1.0) {
         continue;
      }
      double dthetastar = 1e-3;
      double dphistar = 1e-3;
      uhat.SetPolar(1, thetastar + dthetastar, phistar);
      TFourVectorReal kout1(kstar, kstar * uhat);
      uhat.SetPolar(1, thetastar, phistar + dphistar);
      TFourVectorReal kout2(kstar, kstar * uhat);
      kout1.Boost(beta);
      kout2.Boost(beta);
      double dcosthetastar = -sin(thetastar) * dthetastar;
      double dk_dcosthetastar = (kout1[0] - kout[0]) / dcosthetastar;
      double dk_dphistar = (kout2[0] - kout[0]) / dphistar;
      double phi = atan2(kout[2], kout[1]);
      double phi1 = atan2(kout1[2], kout1[1]);
      double phi2 = atan2(kout2[2], kout2[1]);
      double dphi_dcosthetastar = (phi1 - phi) / dcosthetastar;
      double dphi_dphistar = (phi2 - phi) / dphistar;
      double jacob = dk_dcosthetastar * dphi_dphistar -
                     dphi_dcosthetastar * dk_dphistar;
      row.wgt *= jacob * 4 * M_PI;
      if (row.wgt < 0)
         row.wgt = 0;
      row.k[0] = kout[1];
      row.k[1] = kout[2];
      row.k[2] = kout[3];
      TFourVectorReal pout(ptot - kout);
      eIn.SetMom(pin);
      eOut.SetMom(pout);
      gOut.SetMom(kout);
      eIn.SetPol(zeroVector);
      eOut.AllPol();
      gOut.AllPol();
      row.diffXS = TCrossSection::Bremsstrahlung(eIn, eOut, gOut);
      gOut.SetPol(uhat.SetPolar(1, M_PI/2., -phi));
      row.polar = TCrossSection::Bremsstrahlung(eIn, eOut, gOut);
      row.polar /= row.diffXS;

      // include the atomic and crystal form factors
      const double Z = 6;
      const double Sff = 8 * Z;
      const double betacut = 111 * pow(Z, -1/3.) / mElectron;
      const double Fff = 1 / (1 + sqr(betacut) * Qphoton.LengthSqr());
      const double XffSqr_d3q = pow(2 * PI_ * hbarc, 3) / Vcell;
      row.diffXS *= sqr(Sff) * sqr(1-Fff) * XffSqr_d3q;

      tree.Fill();
   }
   tree.Write();
   return nsamples;
}

double phonon_integrand(double *var, double *par)
{
   double x = var[0];
   double nx = 0;
   if (x > 1e-6 && kT > 0) {
      nx = 1 / (exp(ThetaD * x / kT) - 1);
   }
   else {
      nx = kT / (ThetaD * x);
   }
   return x * (2 * nx + 1);
}