//
// destripe.C - utility functions for taking out the effects of differential
//              nonlinearity in the stepper motor for the Bragg angle used
//              for rocking curve measurements at CLS.
//
// author: richard.t.jones at uconn.edu
// version: june 22, 2019
//
// usage example:
//   $ root -l
//   root [1] .L Map2D_cc.so
//   root [2] .L destripe.C++O
//   root [3] destripe("JD70-101_030_results.root");
//   root [4] fourier(43, 5, 300, 794);
//   root [5] apply();
//   root [6] save();

#include <TROOT.h>
#include <TStyle.h>
#include <TCanvas.h>
#include <TFile.h>
#include <TH2D.h>
#include <TH1D.h>
#include <TF1.h>
#include <TProfile.h>
#include <Map2D.h>

#include <vector>
#include <utility>

TDirectory *gRint(0);

void roiMasker();
void roiPicker();
std::vector< std::pair<int, int> > roiVertex;

TString resultsFile("");

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 destripe(const char *results_file)
{
   // Generates the profile histogram sigpro which will contain the
   // average rocking curve peak width as a function of step number
   // for all pixels that peak at that step, for the rocking curve
   // scan results in the results_file. The resulting TProfile
   // called sigpro, and is saved empty in memory for subsequent
   // filling by fourier.

   if (gRint == 0) {
      TH1D test1("t1", "", 1, 0, 1);
      gRint = test1.GetDirectory();
   }
   TFile fresults(results_file);
   TH2D *hmu = (TH2D*)fresults.Get("hmu;2");
   TH2D *hsigma = (TH2D*)fresults.Get("hsigma");
   TH1D *rcurve = (TH1D*)fresults.Get("rcurve");
   int nsteps = rcurve->GetNbinsX();
   TProfile *sigpro = new TProfile("sigpro", "sigma profile",
                                   nsteps, 0, nsteps, 0, 1000, "");
   rootexport(sigpro);
   rootexport(hmu);
   rootexport(hsigma);

   TFile f1("destripe.root");
   Map2D *sigma = (Map2D*)f1.Get("hsigma");
   if (sigma) {
      printf("found hsigma in destripe.root, using that one!\n");
      hsigma = sigma;
   }
   else {
      sigma = new Map2D(*hsigma);
   }
   sigma->SetName("sigma");
   rootexport (sigma);

   // Display the peak sigma map
   int nxbins = sigma->GetNbinsX();
   int nybins = sigma->GetNbinsY();
   double topograph_width = sigma->GetXaxis()->GetNbins();
   double topograph_height = sigma->GetYaxis()->GetNbins();
   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);
   TCanvas *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();
   double xmax = sigma->GetXaxis()->GetBinUpEdge(nxbins);
   double ymax = sigma->GetXaxis()->GetBinUpEdge(nybins);
   if (xmax == nxbins and ymax == nybins) {
      sigma->GetXaxis()->SetTitle("u (pixels)");
      sigma->GetYaxis()->SetTitle("v (pixels)");
      sigma->GetYaxis()->SetTitleOffset(1.4);
   }
   else {
      sigma->GetXaxis()->SetTitle("u (mm)");
      sigma->GetYaxis()->SetTitle("v (mm)");
      sigma->GetYaxis()->SetTitleOffset(1.1);
   }
   if (sigma->GetContour() == 100) {
      sigma->GetZaxis()->SetTitle("#mu^{}rad");
      sigma->GetZaxis()->SetTitleOffset(1.5);
   }
   else {
      sigma->GetZaxis()->SetTitle("steps");
      sigma->GetZaxis()->SetTitleOffset(1.5);
   }
   sigma->SetStats(0);
   sigma->Draw("colz");

   // ask the user to select a region of interest on the displayed image
   roiVertex.clear();
   c1->DeleteExec("roiPicker");
   c1->AddExec("roiPicker","roiPicker();");
   c1->Update();
   std::cout << "Left-click the corners of the region of interest "
                "in the displayed image, then right-click in the image."
                << std::endl;
   resultsFile = TString(results_file);
}

double dsdtheta(double var[], double par[])
{
   // Represents the function ds/d_theta which expresses the Bragg
   // angle motor nonlinearity as a function of step number s in
   // terms of a Fourier series whose coefficients are stored as:
   //    par[0] = number of harmonics in the series
   //    par[1] = angular frequency of fundamental
   //    par[2] = constant term
   //    par[3] = unused
   //    par[4] = cos amplitude for fundamental
   //    par[5] = sin amplitude for fundamental
   //    par[6] = cos amplitude for second harmonic
   //    par[7] = sin amplitude for second harmonic
   //    ...
   //    par[2*N+4] = cos amplitude for N'th harmonic
   //    par[2*N+5] = sin amplitude for N'th harmonic
 
   int harmonics = par[0];
   double omega0 = par[1];
   double result = 0;
   for (int n=0; n<=harmonics; n++) {
      result += par[2*n+2] * cos(n * omega0 * var[0]);
      result += par[2*n+3] * sin(n * omega0 * var[0]);
   }
   return result;
}

void roiPicker()
{
   // 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 any number of times to select a polygonal
   // region of interest, then right-clicks in the image to execute the
   // cropping action.

   TCanvas *c1 = (TCanvas*)gROOT->FindObject("c1");
   TObject *select = c1->GetSelected();
   if (!select) {
      return;
   }
   if (!select->InheritsFrom("TH2")) {
      return;
   }
   Map2D *sigma = (Map2D*)gROOT->FindObject("sigma");
   int event = c1->GetEvent();
   if (event == 1) {
      double x = c1->AbsPixeltoX(c1->GetEventX());
      double y = c1->AbsPixeltoY(c1->GetEventY());
      int px = sigma->GetXaxis()->FindBin(x);
      int py = sigma->GetYaxis()->FindBin(y);
      roiVertex.push_back(std::pair<int,int>(px,py));
      printf("recorded vertex x,y=%lf,%lf\n",x,y);
   }
   else if (event == 12) {
      printf("end of mouse input\n");
      c1->DeleteExec("roiPicker");
      roiMasker();
   }
}

void roiMasker()
{
   // Clear all pixels in the sigma map that lie outside the region of interest
   // defined as the interior of a polygon with vertices selected by the user.
   // This algorithm assumes that the sides of the polygon are touching only
   // at the vertices.

   Map2D *sigma = (Map2D*)gROOT->FindObject("sigma"); 
   if (sigma == 0) {
      printf("Map2D sigma not found, cannot continue!\n");
      return;
   }
   Map2D *sigmask = new Map2D(*sigma);
   sigmask->SetName("sigmask");
   sigmask->Flood(1);
   int nv = roiVertex.size();
   int nborder = 0;
   for (int iv=0; iv<nv; ++iv) {
      int i = roiVertex[iv].first;
      int j = roiVertex[iv].second;
      sigmask->SetBinContent(i,j,1);
      int ie = roiVertex[(iv+1)%nv].first;
      int je = roiVertex[(iv+1)%nv].second;
      sigmask->SetBinContent(ie,je,0);
      nborder += 1;
      int dx = (ie > i)? +1 : -1;
      int dy = (je > j)? +1 : -1;
      double x = (i + 0.5) * dx;
      double y = (j + 0.5) * dy;
      double slope = fabs((je - j)/(ie - i + 1e-9));
      while (sigmask->GetBinContent(i,j) > 0) {
         sigmask->SetBinContent(i,j,0);
         nborder += 1;
         double xb = floor(x + 1) - x;
         double yb = floor(y + 1) - y;
         if (yb > xb * slope) {
            x += xb;
            y += xb * slope;
            i += dx;
         }
         else if (yb < xb * slope) {
            x += yb / slope;
            y += yb;
            j += dy;
         }
         else {
            x += xb;
            y += yb;
            i += dx;
            j += dy;
         }
      }
   }
   int nx = sigmask->GetNbinsX();
   int ny = sigmask->GetNbinsY();
   sigmask->Despeckle((nx*ny - nborder - 1) / 2);
   if (sigmask->GetBinContent(1,1) > 0 ||
       sigmask->GetBinContent(nx,1) > 0 ||
       sigmask->GetBinContent(1,ny) > 0 ||
       sigmask->GetBinContent(nx,ny) > 0)
   {
      Map2D map1(*sigmask);
      map1.Flood(2).Mask(sigmask,1);
      sigmask->Add(&map1);
      sigmask->Shift(0,0,-1);
   }
   rootexport(sigmask);
   sigma->Mask(sigmask);
   sigma->Draw("colz");
   ((TCanvas*)gROOT->FindObject("c1"))->Update();
}

void fourier(double period, int harmonics, int s1=0, int s2=0)
{
   // Reads the sigma map from memory, or from the scratch root file
   // destripe.root if sigma is not found in memory, and extracts a
   // periodic estimate of the motor nonlinearity function ds/d_theta
   // from a section of this histogram from step s1 to s2, using the
   // specified period and number of harmonics. The result is encoded
   // in a TF1 object fdsdtheta, but since that cannot be safely saved
   // in a root file, it is converted to a TH1D named hdsdtheta, which
   // is stored in memory and saved to the scratch destripe.root file.
   // In spite of its name, the dsdtheta function and histogram are
   // not scaled to the actual pitch of the Bragg theta motor in the
   // beamline. That match is made later when dsdtheta is inverted and
   // integrated to find theta(s). Rather, they are scaled to match
   // sigpro, so that the two may be overlaid for visual comparison.

   if (gRint == 0) {
      printf("error - destripe must be invoked before calling fourier!\n");
      return;
   }
   TProfile *sigpro = (TProfile*)gDirectory->Get("sigpro");
   if (sigpro == 0) {
      TFile f("destripe.root");
      sigpro = (TProfile*)f.Get("sigpro");
      if (sigpro == 0) {
         printf("error - sigpro not found in destripe.root, "
                "must invoke destripe first before calling fourier!\n");
         return;
      }
      else {
         printf("warning - sigpro not found in memory, "
                "fetching a copy from destripe.root.\n");
      }
      rootexport(sigpro);
   }
   TH2D *hmu = (TH2D*)gDirectory->Get("hmu");
   Map2D *sigma = (Map2D*)gDirectory->Get("sigma");
   if (hmu == 0 || sigma == 0) {
      TFile f("destripe.root");
      hmu = (TH2D*)f.Get("hmu");
      sigma = (Map2D*)f.Get("sigma");
      if (hmu == 0 || sigma == 0) {
         printf("error - hmu and/or sigma not found in destripe.root, "
                "must invoke destripe first before calling fourier!\n");
         return;
      }
      else {
         printf("warning - hmu and/or sigpma not found in memory, "
                "fetching a copy from destripe.root.\n");
      }
      rootexport(sigma);
      rootexport(hmu);
   }
   sigpro->Reset();
   int nxbins = sigma->GetNbinsX();
   int nybins = sigma->GetNbinsY();
   for (int i=1; i <= nxbins; ++i) {
      for (int j=1; j <= nybins; ++j) {
         double s = sigma->GetBinContent(i,j);
         if (s > 0) {
            double step = hmu->GetBinContent(i,j);
            sigpro->Fill(step,s);
         }
      }
   }
   sigpro->SetStats(0);
   sigpro->GetXaxis()->SetTitle("step number");
   sigpro->GetYaxis()->SetTitle("sigma (steps)");
   TFile f("destripe.root", "update");
   sigpro->Write();

   int nsteps = sigpro->GetNbinsX();
   s1 = (s1 > 0)? s1 : 1;
   s2 = (s2 > s1)? s2 : nsteps;
   int nsamples = s2 - s1 + 1;
   double sample_periods = nsamples / period;
   double ccoef[harmonics+1];
   double scoef[harmonics+1];
   double omega0 = 2 * 3.1415926 / period;
   for (int n=0; n<=harmonics; ++n) {
      ccoef[n] = scoef[n] = 0;
   }
   for (int i=s1; i<=s2; ++i) {
      double s = sigpro->GetBinContent(i);
      ccoef[0] += s;
      for (int n=1; n<=harmonics; ++n) {
         ccoef[n] += s * cos(n * omega0 * i); 
         scoef[n] += s * sin(n * omega0 * i); 
      }
   }
   TF1 *fdsdtheta;
   fdsdtheta = new TF1("fdsdtheta", dsdtheta, 0, nsteps, 2*(harmonics + 1));
   fdsdtheta->SetParameter(0, harmonics);
   fdsdtheta->SetParameter(1, omega0);
   double sfact = 1.0;
   ccoef[0] /= 2 * sfact;
   for (int n=0; n<=harmonics; ++n) {
      ccoef[n] *= 2*sfact / nsamples;
      scoef[n] *= 2*sfact / nsamples;
      fdsdtheta->SetParameter(2*n+2, ccoef[n]);
      fdsdtheta->SetParameter(2*n+3, scoef[n]);
      printf("%lf %lf\n", ccoef[n], scoef[n]);
   }
   for (int n=1; n<=harmonics; ++n) {
      fdsdtheta->SetParameter(2*n+2, 1.1*ccoef[n]);
      fdsdtheta->SetParameter(2*n+3, 1.1*scoef[n]);
   }
   printf("\n");
   TH1D *hdsdtheta;
   TString title;
   title.Form("ds/d_theta with period %lf steps", period);
   hdsdtheta = new TH1D("hdsdtheta", title, nsteps, 0, nsteps);
   hdsdtheta->Eval(fdsdtheta);
   hdsdtheta->GetXaxis()->SetTitle("step");
   hdsdtheta->GetYaxis()->SetTitle("ds/d#theta (arb. units)");
   hdsdtheta->SetLineColor(kRed);
   hdsdtheta->SetStats(0);
   hdsdtheta->Draw();
   rootexport(hdsdtheta);
   sigpro->Draw();
   hdsdtheta->Draw("same");

   // Now integrate the reciprocal to get the function theta(s).
   TH1D *hthetaofs = (TH1D*)hdsdtheta->Clone("hthetaofs");
   double meanslope = 0;
   int periods = 0;
   double sum = 0;
   for (int i=1; i<=nsteps; ++i) {
      double dsdt = hdsdtheta->GetBinContent(i);
      sum += 1/dsdt;
      hthetaofs->SetBinContent(i, sum);
      if (i - int(i/period)*period < 1) {
         meanslope = sum;
         periods += 1;
      }
   }
   meanslope /= periods * period;
   double trueslope = 0.00014; // degrees per step, see logbook
   hthetaofs->Scale(trueslope*(3.1415926/180)*1e6 / meanslope);
   hthetaofs->GetYaxis()->SetTitle("#theta(s) (#murad)");
   hthetaofs->GetYaxis()->SetTitleOffset(1.4);
   hthetaofs->SetTitle("#theta vs step number");
   rootexport(hthetaofs);

   // save copies to a scratch directory
   hdsdtheta->Write();
   hthetaofs->Write();
}

void apply(const char* results_file="")
{
   // Read in the rocking curve peak centroid (hmu) and width (hsigma)
   // topographs from resultsfile, and apply the correction that takes
   // out the effects of nonlinearity in the Bragg theta angle motor.
   // The algorithm that does this requires as input the function theta(s)
   // which it reads in the form of a 1d histogram from the destripe.root
   // root file. It should have already been computed using the other
   // tools in this suite and stored there for use with this resultsfile.
   // In general, one needs a different theta(s) for each scan. At exit,
   // corrected versions of hmu and hsigma named hmu_moco and hsigma_moco
   // are saved to resultsfile.

   if (gRint == 0) {
      printf("error - destripe must be invoked first, "
             "followed by fourier, before calling apply!\n");
      return;
   }
   TH1D *hthetaofs = (TH1D*)gDirectory->Get("hthetaofs");
   if (hthetaofs == 0) {
      TFile f1("destripe.root");
      TH1D *hthetaofs = (TH1D*)f1.Get("hthetaofs");
      if (hthetaofs == 0) {
         printf("error in apply - motor response function hthetaofs "
                "was not found in destripe.root, cannot continue!\n");
         return;
      }
      rootexport(hthetaofs);
   }
   if (strlen(results_file) > 0)
      resultsFile = TString(results_file);
   TFile fresults(resultsFile);
   if (!fresults.IsOpen()) {
      printf("error in apply - unable to open results file %s"
             "for output, cannot continue!\n", resultsFile.Data());
      return;
   }
   TH2D *hmu = (TH2D*)fresults.Get("hmu;2");
   TH2D *hsigma = (TH2D*)fresults.Get("hsigma;2");
   if (hmu == 0 or hsigma == 0) {
      printf("error in apply - rocking curve topograms hmu, hsigma "
             "were not found in the input file, cannot continue!\n");
      return;
   }
   TH2D *hmu_moco = (TH2D*)hmu->Clone("hmu_moco");
   TH2D *hsigma_moco = (TH2D*)hsigma->Clone("hsigma_moco");
   int nxbins = hmu_moco->GetNbinsX();
   int nybins = hmu_moco->GetNbinsY();
   for (int i=1; i<=nxbins; ++i) {
      for (int j=1; j<=nybins; ++j) {
         double mu = hmu->GetBinContent(i,j);
         double sig = hsigma->GetBinContent(i,j);
         double theta0 = hthetaofs->Interpolate(mu - sig);
         double theta1 = hthetaofs->Interpolate(mu + sig);
         hmu_moco->SetBinContent(i,j, (theta1 + theta0)/2);
         hsigma_moco->SetBinContent(i,j, (theta1 - theta0)/2);
      }
   }
   TH2D *hmu_resized = (TH2D*)fresults.Get("hmu;3");
   double hmu_width = hmu_resized->GetXaxis()->GetBinUpEdge(nxbins);
   double hmu_height = hmu_resized->GetYaxis()->GetBinUpEdge(nybins);
   hmu_moco->GetXaxis()->SetLimits(0,hmu_width);
   hmu_moco->GetYaxis()->SetLimits(0,hmu_height);
   hmu_moco->GetXaxis()->SetTitle("u (mm)");
   hmu_moco->GetYaxis()->SetTitle("v (mm)");
   hmu_moco->GetYaxis()->SetTitleOffset(1.1);
   hmu_moco->SetContour(100);
   hmu_moco->GetZaxis()->SetTitle("#murad");
   hmu_moco->GetZaxis()->SetTitleOffset(1.5);
   hsigma_moco->GetXaxis()->SetLimits(0,hmu_width);
   hsigma_moco->GetYaxis()->SetLimits(0,hmu_height);
   hsigma_moco->GetXaxis()->SetTitle("u (mm)");
   hsigma_moco->GetYaxis()->SetTitle("v (mm)");
   hsigma_moco->GetYaxis()->SetTitleOffset(1.1);
   hsigma_moco->SetContour(100);
   hsigma_moco->GetZaxis()->SetTitle("#murad");
   hsigma_moco->GetZaxis()->SetTitleOffset(1.5);
   hmu_moco->SetStats(0);
   hsigma_moco->SetStats(0);
   rootexport(hmu_moco);
   rootexport(hsigma_moco);
   TFile f1("destripe.root", "update");
   hmu_moco->Write();
   hsigma_moco->SetMaximum(20);
   hsigma_moco->Write();
   hsigma_moco->Draw("colz");
   ((TCanvas*)gROOT->FindObject("c1"))->Update();
}

void save(const char* results_file="")
{
   // Saves the hsigma_moco and hmu_moco motor-corrected topographs
   // in memory to the results_file.

   TH2D *hmu_moco = (TH2D*)gDirectory->Get("hmu_moco");
   TH2D *hsigma_moco = (TH2D*)gDirectory->Get("hsigma_moco");
   if (hmu_moco == 0 || hsigma_moco == 0) {
      printf("save error - motor-corrected topograph maps hmu_moco "
             "and/or hsigma_moco not found, cannot continue!\n");
      return;
   }
   if (strlen(results_file) > 0)
      resultsFile = TString(results_file);
   TFile fresults(resultsFile, "update");
   hmu_moco->Write();
   hsigma_moco->Write();
   printf("hmu_moco and hsigma_moco successfully saved to %s\n",
          resultsFile.Data());
}