//
// beamspot.C - displays an X-ray rocking curve topograph image of the
//              peak angle vs position, and allows the user to select
//              a spot to center the electron beam. A gaussian intensity
//              profile is generated centered on the selected spot and
//              the whole-crystal rocking curve is computed, weighted
//              by the beam intensity.
//
// author: richard.t.jones at uconn.edu
// version: june 22, 2019
//
// usage example:
//   $ root -l
//   root [1] .L beamspot.C++O
//   root [2] spoton("JD70-101_030_results.root");
//   root [3] rcurve_spot.Draw();

#define CACHE_RCTREE_AND_MULTIPROCESS 1

#include <TROOT.h>
#include <TString.h>
#include <TStyle.h>
#include <TColor.h>
#include <TFile.h>
#include <TH2D.h>
#include <TH1D.h>
#include <TF1.h>
#include <TTree.h>
#include <TCanvas.h>
#include <TExec.h>

#include <vector>
#include <iostream>
#include <thread>

double angle_step_size = 0.00014 * 3.1415926e+6/180;

double beamspot_sigma_x(1.0); // mm
double beamspot_sigma_y(0.7); // mm
double beamspot_center_x(0);
double beamspot_center_y(0);
int topograph_width(0);
int topograph_height(0);
TH2D *topograph_image(0);

TString rootfile("root://nod26.phys.uconn.edu//Gluex/beamline/diamonds/%s/results/%s_%3.3d_rocking_curves.root");
TString resultsfile("root://nod26.phys.uconn.edu//Gluex/beamline/diamonds/%s/results/%s_%3.3d_results.root");

TCanvas *c1 = 0;
TDirectory *gRint(0);
TFile *frocking_curve = 0;
TFile *fresults = 0;

std::vector<TH1D*> hcurve_cache;
TString cached_rctree_id;

int rootexport(TH1 *obj)
{
   // Exports ROOT object *obj to the global ROOT namespace.

   TString name(obj->GetName());
   TString temp("X9pe8BqlcuTXlhHu");
   obj->SetName(temp);
   TH1 *old;
   int overwrites = 0;
   while ((old = (TH1*)gRint->Get(name))) {
      gRint->Delete(name.Data());
      overwrites++;
   }
   obj->SetDirectory(gRint);
   obj->SetName(name);
   return overwrites;
}

void spoton(const char* run, const char* sample, int scan)
{
   // Displays the rocking curve peak angle (hmu) topograph image
   // and invites the user to select a location where the electron
   // beam spot should be centered. Once the user left-clicks the
   // mouse at the desired position, a gaussian beam spot is placed
   // there and the summed rocking curve of the entire crystal is
   // taken, weighted by the beam intensity, and saved under the
   // name rcurve_spot. The displayed hmu topograph is updated to
   // show the position and intensity pattern of the beam spot
   // overlayed on the crystal image.

   if (gRint == 0) {
      TH1D test1("test1", "", 1, 0, 1);
      gRint = test1.GetDirectory();
   }
   if (frocking_curve)
      delete frocking_curve;
   frocking_curve = TFile::Open(TString::Format(rootfile, run, sample, scan));
   if (!frocking_curve->IsOpen()) {
      printf("Error in spoton - input rootfile does not contain the expected "
             "rocking curve tree rctree, cannot continue!\n");
      frocking_curve = 0;
      return;
   }
   if (fresults)
      delete fresults;
   fresults = TFile::Open(TString::Format(resultsfile, run, sample, scan));
   if (!fresults->IsOpen()) {
      printf("Error in spoton - failed to open input file %s\n",
             TString::Format(resultsfile, run, sample, scan).Data());
      fresults = 0;
      return;
   }
   TH2D *hmu = (TH2D*)fresults->Get("hmu_moco");
   if (hmu == 0) {
      printf("Error in spoton - input results file does not contain "
             "the expected rocking curve topograph hmu, cannot continue!\n");
      return;
   }
   gStyle->SetPalette(1,0);
   gStyle->SetNumberContours(50);
   gStyle->SetOptStat(0);
   gStyle->SetOptFit(1);
   gStyle->SetStatX(0.95);
   gStyle->SetStatY(0.93);
   gStyle->SetStatW(0.18);
   gStyle->SetStatH(0.15);
   topograph_width = hmu->GetXaxis()->GetNbins();
   topograph_height = hmu->GetYaxis()->GetNbins();
   topograph_image = hmu;
   double plotheight = 400;
   double plotwidth = plotheight*topograph_width/topograph_height;
   double leftmargin = 50;
   double bottommargin = 50;
   double rightmargin = 90;
   double topmargin = 50;
   double canvaswidth = plotwidth+leftmargin+rightmargin;
   double canvasheight = plotheight+topmargin+bottommargin;
   gStyle->SetCanvasPreferGL(true);
   c1 = new TCanvas("c1","rocking curve topograph",50,10,
                    int(canvaswidth),int(canvasheight));
   c1->SetLeftMargin(leftmargin/canvaswidth);
   c1->SetRightMargin(rightmargin/canvaswidth);
   c1->SetTopMargin(topmargin/canvasheight);
   c1->SetBottomMargin(bottommargin/canvasheight);
   c1->cd();
   hmu->GetXaxis()->SetTitle("u (mm)");
   hmu->GetYaxis()->SetTitle("v (mm)");
   hmu->GetYaxis()->SetTitleOffset(1.1);
   if (hmu->GetContour() == 100) {
      hmu->GetZaxis()->SetTitle("#mu^{}rad");
      hmu->GetZaxis()->SetTitleOffset(1.5);
   }
   hmu->Draw("colz");
   hmu->SetDirectory(gRint);
   c1->DeleteExec("rcBeamspotPicker");
   c1->AddExec("rcBeamspotPicker","beamspotPicker();");
   c1->Update();
   std::cout << "Click on the topograph image to select "
                "the center of the beam spot." << std::endl;
}

void make_beamspot()
{
   // Generates a 2d gaussian beam spot with the properties
   //   beamspot_center_x, beamspot_center_y,
   //   beamspot_sigma_x, beamspot_sigma_y
   // all in mm, and saves it as a 2d histogram named beamspot.
   // It then computes the whole-crystal rocking curve weighted
   // by the gaussian beam intensity, and saves it in the 1d
   // histogram rcurve_spot. These two output histograms are
   // saved in the output root file beamspot.root, together with
   // the topograph_image that goes with them.
 
   if (gRint == 0) {
      printf("error in make_beamspot - must call spoton first!\n");
      return;
   }
   if (!frocking_curve->IsOpen()) {
      printf("error in make_beamspot - must call spoton first!\n");
      return;
   }
   TTree *rctree = (TTree*)frocking_curve->Get("rctree");
   if (rctree == 0) {
      printf("Error in spoton - input rootfile does not contain the expected "
             "rocking curve tree rctree, cannot continue!\n");
      return;
   }
   printf("making spot at (%lf,%lf)\n",
          beamspot_center_x, beamspot_center_y);
   TH1D *hcurve(0);
   rctree->SetBranchAddress("rchist", &hcurve);
   rctree->GetEntry(0);
   if (hcurve == 0) {
      printf("GetEntry(0) on rctree returned null, cannot continue!\n");
      return;
   }

#if CACHE_RCTREE_AND_MULTIPROCESS
   if (frocking_curve->GetName() != cached_rctree_id) {
      unsigned int nentries = rctree->GetEntries();
      for (unsigned int i=0; i < nentries; i++) {
         while (i >= hcurve_cache.size())
            hcurve_cache.push_back(0);
         rctree->SetBranchAddress("rchist", &hcurve_cache[i]);
         rctree->GetEntry(i);
      }
      cached_rctree_id = frocking_curve->GetName();
      printf("set unique id %s\n", cached_rctree_id.Data());
   }
#endif

   TH2D *beamspot = (TH2D*)topograph_image->Clone("beamspot");
   TH1D *rcurve_spot = (TH1D*)hcurve->Clone("rcurve_spot");
   int nbins = rcurve_spot->GetNbinsX();
   rcurve_spot->Reset();
   beamspot->Reset();

   // parallelize the sum
   const int nthreads = 8;
   std::thread thr[nthreads];
   TH1D *rc[nthreads];
   TH2D *bs[nthreads];
   for (int it=0; it<nthreads; it++) {
      int j1 = topograph_height/nthreads * it + 1;
      int j2 = topograph_height/nthreads * (it + 1);
      j2 = (j2 < topograph_height)? j2 : topograph_height;
#if CACHE_RCTREE_AND_MULTIPROCESS
      rc[it] = (TH1D*)rcurve_spot->Clone();
      bs[it] = (TH2D*)beamspot->Clone();
      thr[it] = std::thread([rc, bs](int it, int j1, int j2) {
#else
      rc[it] = rcurve_spot;
      bs[it] = beamspot;
#endif
      for (int j=j1; j<=j2; j++) {
         double y = topograph_image->GetYaxis()->GetBinCenter(j);
         double y2 = pow((y - beamspot_center_y) / beamspot_sigma_y, 2);
         for (int i=1; i<=topograph_width; i++) {
            double x = topograph_image->GetXaxis()->GetBinCenter(i);
            double x2 = pow((x - beamspot_center_x) / beamspot_sigma_x, 2);
            double w = exp(-0.5 * (x2 + y2));
            if (w > 1e-9) {
#if CACHE_RCTREE_AND_MULTIPROCESS
               TH1D *hcurve = hcurve_cache[i + (j-1)*topograph_width];
#else
               rctree->GetEntry(i + (j-1)*topograph_width);
#endif
               rc[it]->Add(hcurve, w);
               bs[it]->Fill(x, y, w);
            }
         }
      }
#if CACHE_RCTREE_AND_MULTIPROCESS
                                                  }, it, j1, j2);
#endif
   }
#if CACHE_RCTREE_AND_MULTIPROCESS
printf("spawned %d threads, reaping now...\n", nthreads);
   for (int it=0; it<nthreads; it++) {
      thr[it].join();
      rcurve_spot->Add(rc[it]);
      beamspot->Add(bs[it]);
      rc[it]->Delete();
      bs[it]->Delete();
   }
#endif

   rcurve_spot->GetXaxis()->SetLimits(0,nbins*angle_step_size);
   rcurve_spot->SetTitle("whole-crystal rocking curve with beam spot");
   rcurve_spot->GetXaxis()->SetTitle("Bragg angle (#murad)");

   // fit the spot rocking curve to a gaussian over constant bg
   double ymin = rcurve_spot->GetMinimum();
   double ymax = rcurve_spot->GetMaximum();
   double ymean = rcurve_spot->Integral() / rcurve_spot->GetNbinsX();
   double xpeak = rcurve_spot->GetXaxis()
                           ->GetBinCenter(rcurve_spot->GetMaximumBin());
   double xrms = 50;
   double xrms_min = 10;
   double xrange[2] = {rcurve_spot->GetXaxis()->GetBinLowEdge(1),
                       rcurve_spot->GetXaxis()->GetBinUpEdge(nbins)};
   TF1 fitfunc("rcshape","[0]+[1]*exp(-0.5*(x-[2])*(x-[2])/([3]*[3]))",
               xrange[0],xrange[1]);
   fitfunc.SetParName(0,"ybase");
   fitfunc.SetParName(1,"yamp");
   fitfunc.SetParName(2,"xmu");
   fitfunc.SetParName(3,"xsigma");
   fitfunc.SetParameter(0,ymean);
   fitfunc.SetParameter(1,ymax-ymean);
   fitfunc.SetParameter(2,xpeak);
   fitfunc.SetParameter(3,xrms);
   fitfunc.SetParLimits(0,ymin*0.999,ymax*1.001);
   fitfunc.SetParLimits(1,0.,2*(ymax-ymin)+1e-3);
   fitfunc.SetParLimits(2,xrange[0],xrange[1]);
   fitfunc.SetParLimits(3,xrms_min,xrange[1]-xrange[0]);
   TCanvas *c2 = new TCanvas("c2","rocking curve", 600, 10, 600, 500);
   TFitResultPtr res = rcurve_spot->Fit(&fitfunc,"SQR");
   c2->Update();
   rootexport(rcurve_spot);

   // save resulting images to output file
   TFile f1("beamspot.root", "update");
   rcurve_spot->Write();
   beamspot->Write();
   topograph_image->Write();

   // superimpose the beamspot image on the topograph 
   c1->cd();
   TExec *ex1 = new TExec("ex1", "palette1();");
   TExec *ex2 = new TExec("ex2", "palette2();");
   ex1->Draw();
   topograph_image->Draw("col same");
   ex2->Draw();
   beamspot->SetMinimum(0.1);
   beamspot->Draw("col same");
   ex1->Draw();
}

void palette1()
{
   gStyle->SetPalette(1);
}

void palette2()
{
   gStyle->SetPalette(kInvertedDarkBodyRadiator,0,0.6);
}

void beamspotPicker()
{
   // This callback is called whenever a mouse event (move, click, etc.)
   // happens inside one of the open ROOT graphics windows. This one
   // activates when the user left-clicks inside the active region of a
   // displayed TH2D histogram and sets the beamspot_center coordinates
   // to the location clicked on, in units of the TH2D axes.

   TObject *select = c1->GetSelected();
   if (!select) {
      return;
   }
   if (!select->InheritsFrom("TH2")) {
      return;
   }
   if (c1->GetEvent() == 1) {
      int px = c1->GetEventX();
      int py = c1->GetEventY();
      beamspot_center_x = c1->AbsPixeltoX(px);
      beamspot_center_y = c1->AbsPixeltoY(py);
      c1->DeleteExec("rcBeamspotPicker");
      make_beamspot();
   }
}

void set_beamspot(double x0, double y0, double xsig=0, double ysig=0)
{
   beamspot_center_x = x0;
   beamspot_center_y = y0;
   if (xsig > 0)
      beamspot_sigma_x = xsig;
   if (ysig > 0)
      beamspot_sigma_y = ysig;
}

void get_beamspot(double *x0, double *y0, double *xsig=0, double *ysig=0)
{
   *x0 = beamspot_center_x;
   *y0 = beamspot_center_y;
   if (xsig != 0)
      *xsig = beamspot_sigma_x;
   if (ysig != 0)
      *ysig = beamspot_sigma_y;
}