//
// Cbrems.C
//
// Generates maps of coherent bremsstrahlung from a single reciprocal
// lattice vector in a diamond crystal.
//
// author: richard.t.jones at uconn.edu
// version: february 11, 2017

#include "Complex.h"
#include "TPhoton.h"
#include "TLepton.h"
#include "TCrossSection.h"
#include "constants.h"

#include <TROOT.h>
#include <TCanvas.h>
#include <TF1.h>
#include <TF2.h>
#include <TH1D.h>
#include <TAxis.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

double Ebeam = 12.; // GeV
double Eedge = 9.; // GeV
double qtotal = 9.8e-6;
double qlong = Eedge / (Ebeam - Eedge) * sqr(mElectron) / (2 * Ebeam);
TThreeVectorReal qRecoil(qtotal, 0, -qlong);
double radcoldist = 76.; // m
double pbeam = sqrt(sqr(Ebeam) - sqr(mElectron));

Double_t energy(Double_t *var, Double_t *par)
{
   TThreeVectorReal khat(var[0], var[1], radcoldist * 1e3);
   khat.Normalize(1);
   double kmag = -(pbeam * qRecoil[3] + qRecoil.LengthSqr() / 2) /
                  (Ebeam - pbeam * khat[3] - qRecoil.Dot(khat));
   return kmag;
}

Double_t intensity(Double_t *var, Double_t *par)
{
   TThreeVectorReal khat(var[0], var[1], radcoldist * 1e3);
   khat.Normalize(1);
   double kmag = -(pbeam * qRecoil[3] + qRecoil.LengthSqr() / 2) /
                  (Ebeam - pbeam * khat[3] - qRecoil.Dot(khat));
   TFourVectorReal kout(kmag, kmag * khat);
   TFourVectorReal pin(Ebeam, 0, 0, pbeam);
   TFourVectorReal qin(0, qRecoil);
   TFourVectorReal pout(pin + qin - kout);

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

   // Solve for the rest of the kinematics
   eIn.SetMom(pin);
   gOut.SetMom(kout);
   eOut.SetMom(pout);

   // Set the initial,final polarizations
   eIn.SetPol(zeroVector);
   gOut.AllPol();
   eOut.AllPol();

   // Multiply the basic cross section by the form factors
   double result = TCrossSection::Bremsstrahlung(eIn, eOut, gOut);
   const double Z = 6;
   const double Sff = 8 * Z;
   const double Aphonon = 0.5e9;			                // in /GeV**2
   const double Gff = exp(-Aphonon * qRecoil.LengthSqr() / 2);
   const double beta = 111 * pow(Z, -1/3.) / mElectron;	    // ff cutoff in /GeV
   const double Fff = 1 / (1 + sqr(beta) * qRecoil.LengthSqr());
   const double hbarc = 1.97327e-6;		                  	// in GeV.Angstroms
   const double Vcell = 45.5;	                            // in Angstroms**3
   const double XffSqr_d3q = pow(2 * PI_ * hbarc, 3) / Vcell;
   result *= sqr(Sff) * sqr(Gff) * sqr(1-Fff) * XffSqr_d3q;
   return result;                                           // ub/GeV/radian
}

Double_t polarization(Double_t *var, Double_t *par)
{
   TThreeVectorReal khat(var[0], var[1], radcoldist * 1e3);
   khat.Normalize(1);
   double kmag = -(pbeam * qRecoil[3] + qRecoil.LengthSqr() / 2) /
                  (Ebeam - pbeam * khat[3] - qRecoil.Dot(khat));
   TFourVectorReal kout(kmag, kmag * khat);
   TFourVectorReal pin(Ebeam, 0, 0, pbeam);
   TFourVectorReal qin(0, qRecoil);
   TFourVectorReal pout(pin + qin - kout);

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

   // Solve for the rest of the kinematics
   eIn.SetMom(pin);
   gOut.SetMom(kout);
   eOut.SetMom(pout);

   // Measure the transverse polarization in the plane of qRecoil
   double phi = atan2(var[1], var[0]);
   eIn.SetPol(zeroVector);
   eOut.AllPol();
   TThreeVectorReal xhat,yhat;
   xhat.SetPolar(1, PI_/2, -phi);
   yhat.SetPolar(1, PI_/2, -phi + PI_/2);
   gOut.SetPol(xhat);
   double Xrate = TCrossSection::Bremsstrahlung(eIn, eOut, gOut);
   gOut.SetPol(yhat);
   double Yrate = TCrossSection::Bremsstrahlung(eIn, eOut, gOut);
   return (Xrate - Yrate) / (Xrate + Yrate);
}

void plot_colz(TF2 *f2, TString plotfile)
{
   f2->SetNpx(500);
   f2->SetNpy(500);
   f2->GetXaxis()->SetTitle("x (mm)");
   f2->GetYaxis()->SetTitle("y (mm)");
   f2->SetContour(20);
   f2->Draw("colz");
   c1->Update();
   c1->Print(plotfile);
}

void plot_slix(TF2 *f2, TString plotfile)
{
   TObject *o = (TH1D*)gROOT->FindObject("h");
   if (o)
      o->Delete();
   double xmin, xmax, ymin, ymax;
   f2->GetRange(xmin, ymin, xmax, ymax);
   TH1D *h = new TH1D("h", f2->GetTitle(), 500, xmin, xmax);
   for (int i = 1; i <= 500; ++i) { 
      double x = h->GetXaxis()->GetBinCenter(i);
      h->SetBinContent(i, (*f2)(x, 0));
   }
   c1->SetLeftMargin(0.15);
   c1->SetRightMargin(0.10);
   h->GetXaxis()->SetTitle("x (mm), y=0");
   h->SetStats(0);
   h->Draw();
   c1->Update();
   c1->Print(plotfile);
}

void plot_sliy(TF2 *f2, TString plotfile)
{
   TObject *o = (TH1D*)gROOT->FindObject("h");
   if (o)
      o->Delete();
   double xmin, xmax, ymin, ymax;
   f2->GetRange(xmin, ymin, xmax, ymax);
   TH1D *h = new TH1D("h", f2->GetTitle(), 500, ymin, ymax);
   for (int i = 1; i <= 500; ++i) { 
      double y = h->GetXaxis()->GetBinCenter(i);
      h->SetBinContent(i, (*f2)(0, y));
   }
   c1->SetLeftMargin(0.15);
   c1->SetRightMargin(0.10);
   h->GetXaxis()->SetTitle("y (mm), x=0");
   h->SetStats(0);
   h->Draw();
   c1->Update();
   c1->Print(plotfile);
}

void make_plots()
{
   static int initialized = 0;
   if (!initialized) {
      c1 = new TCanvas("c1", "", 550, 500);
      initialized = 1;
   }
   TF2 *fen = new TF2("fen", energy, -10., 10., -10., 10);
   fen->SetTitle("coherent energy vs position, 9 GeV edge");
   TF2 *fin = new TF2("fin", intensity, -10., 10., -10., 10);
   fin->SetTitle("coherent intensity vs position, 9 GeV edge");
   TF2 *fpo = new TF2("fpo", polarization, -10., 10., -10., 10);
   fpo->SetTitle("coherent polarization vs position, 9 GeV edge");
   c1->SetLeftMargin(0.10);
   c1->SetRightMargin(0.15);
   plot_colz(fen, "cohEnergy.png");
   plot_colz(fin, "cohIntensity.png");
   plot_colz(fpo, "cohPolarization.png");
   plot_slix(fen, "cohEnergy_x.png");
   plot_slix(fin, "cohIntensity_x.png");
   plot_slix(fpo, "cohPolarization_x.png");
   plot_sliy(fen, "cohEnergy_y.png");
   plot_sliy(fin, "cohIntensity_y.png");
   plot_sliy(fpo, "cohPolarization_y.png");
}

void make_hires_plots()
{
   static int initialized = 0;
   if (!initialized) {
      c1 = new TCanvas("c1", "", 550, 500);
      initialized = 1;
   }
   TF2 *fenh = new TF2("fenh", energy, -2.5, 2.5, -2.5, 2.5);
   fenh->SetTitle("coherent energy vs position, 9 GeV edge");
   TF2 *finh = new TF2("finh", intensity, -2.5, 2.5, -2.5, 2.5);
   finh->SetTitle("coherent intensity vs position, 9 GeV edge");
   TF2 *fpoh = new TF2("fpoh", polarization, -2.5, 2.5, -2.5, 2.5);
   fpoh->SetTitle("coherent polarization vs position, 9 GeV edge");
   c1->SetLeftMargin(0.10);
   c1->SetRightMargin(0.15);
   plot_colz(fenh, "cohEnergy_h.png");
   plot_colz(finh, "cohIntensity_h.png");
   plot_colz(fpoh, "cohPolarization_h.png");
   plot_slix(fenh, "cohEnergy_hx.png");
   plot_slix(finh, "cohIntensity_hx.png");
   plot_slix(fpoh, "cohPolarization_hx.png");
   plot_sliy(fenh, "cohEnergy_hy.png");
   plot_sliy(finh, "cohIntensity_hy.png");
   plot_sliy(fpoh, "cohPolarization_hy.png");
}