//
// rockmean.C - root macro to extract the rocking curve peak (mean value)
//              for each cell in a rocking curve series and display the
//              result in a high-resolution color-coded map.
//
// author: richard.t.jones at uconn.edu
// version: april 29, 2013
//
// usage: 
//	$ root -l
//	root [0] .L Map2D.cc+O
//	root [1] .L rcfitter.C+O
//	root [2] .L rcpicker.C+O
//	root [3] .L rockmean.C
//	root [4] set_run_period("chess-5-2014");
//	root [5] Map2D *u40_6_2 = rockmean("U40study06",2);
//	root [6] rcpicker *rcp = picker("U40study06",2);
//	root [7] rcp->PickOn(hmu);

#include <TROOT.h>
#include <TFile.h>
#include <TString.h>
#include <iostream>
#include <rcpicker.h>
#include <Map2D.h>
#include <unistd.h>
#include <TF1.h>
#include <TAxis.h>

double plot_size_mm = 8.0;
int reverse_images = 0;

TString xrootd_url("root://nod26.phys.uconn.edu");
TString data_dir("/Gluex/beamline/diamonds/");
TString run_period("cls-8-2016");

#define THETA_BRAGG_DEG 19.2
//#define THETA_BRAGG_DEG 14.2
#define RESULTS_ON_PNFS 1
#define USE_RECTANGULAR_RIO 1

Map2D *rockmean(const char* file, int scan, double riothresh=160,
                                            double finthresh=50, 
                                            double pctthresh=20, 
                                            double xmin=0,
                                            double xmax=1e30,
                                            double ymin=0,
                                            double ymax=1e30,
                                            double max_mu=0,
                                            double max_sigma=0,
                                            int specsize=10000)
{
   // Generate a rocking curve centroid color topograph
   // from the fitted rocking curve results found in
   // <file>_<scan>_results.root, normally stored on pnfs.
   // The images found in the root file are rescaled into
   // length units of mm and cropped to the box {(xmin,ymin),
   // (xmax,ymax)}. The parameters max_mu and max_sigma
   // set the maximum of the color scale for the mu and
   // sigma color topographs. The images are then despeckled
   // with the requested specsize. The meaning of the
   // threshold arguments is as follows.
   //   1) riothresh = threshold baseline value used to
   //                  define the region of interest
   //                  for the image; pixels outside the
   //                  region of interest are set to the
   //                  matte value, normally 0.
   //  2) finthresh = threshold amplitude value used to
   //                 define what constitutes a fitted
   //                 peak; fits with amplitudes less than
   //                 this are reset to matte, normally 0.
   //  3) pctthresh = threshold percent in terms of the
   //                 regional average of amplitude over
   //                 baseline values; pixels with fits
   //                 giving amplitudes less than this
   //                 are reset to matte, normally 0.

   TFile *f;
   TString rootfile;
   TString rootURL(xrootd_url + data_dir + run_period + "/results");
#if RESULTS_ON_PNFS
   rootfile.Form("%s/%s_%3.3d_results.root",rootURL.Data(),file,scan);
   std::cout << "file is " << rootfile << std::endl;
#else
   rootfile.Form("%s/%s_%3.3d_results.root",".",file,scan);
#endif
   f = TFile::Open(rootfile.Data());
   if (! f->IsOpen()) {
      std::cerr << "Error - unable to open input root file \"" << rootfile
                << "\"" << std::endl;
      return 0;
   }
   TString rcfile;
   rcfile.Form("%s/%s_%3.3d_rocking_curves.root",rootURL.Data(),file,scan);
   TString logfile;
   logfile.Form("%s/%s/%s",run_period.Data(),file,file);
   rcpicker *rcp = new rcpicker(file, scan, rcfile.Data(), logfile.Data());
   f->cd();
   TH2D* hmu;
   f->GetObject("hmu",hmu);
   TH2D* hamp;
   f->GetObject("hamp",hamp);
   TH2D* hbase;
   f->GetObject("hbase",hbase);
   TH2D* hsigma;
   f->GetObject("hsigma",hsigma);
   if (hmu == 0 || hamp == 0 || hbase == 0 || hsigma == 0) {
      std::cerr << "Error - unable to read hmu, hsigma, hbase, hamp"
                << " from input file \"" << rootfile << "\"" << std::endl;
   }
   Double_t distanceScale = rcp->GetMicronsPerPixel()/1000;
  
   Int_t nbinx = hbase->GetNbinsX();
   Int_t nbiny = hbase->GetNbinsY();
   Int_t rescale_needed = (hbase->GetXaxis()->GetXmax() == nbinx);
   if (rescale_needed) {
      hbase->GetXaxis()->SetLimits(0,nbinx*distanceScale);
      hbase->GetYaxis()->SetLimits(0,nbiny*distanceScale);
      hamp->GetXaxis()->SetLimits(0,nbinx*distanceScale);
      hamp->GetYaxis()->SetLimits(0,nbiny*distanceScale);
      hmu->GetXaxis()->SetLimits(0,nbinx*distanceScale);
      hmu->GetYaxis()->SetLimits(0,nbiny*distanceScale);
      hsigma->GetXaxis()->SetLimits(0,nbinx*distanceScale);
      hsigma->GetYaxis()->SetLimits(0,nbiny*distanceScale);
   }
   Map2D* mbase = new Map2D(*hbase);
   mbase->SetName("mbase");
   Map2D* mamp = new Map2D(*hamp);
   mamp->SetName("mamp");
   Map2D* mmu = new Map2D(*hmu);
   mmu->SetName("mmu");
   Map2D* msigma = new Map2D(*hsigma);
   msigma->SetName("msigma");
   Map2D* mcut = new Map2D(*mbase);
   mcut->SetName("mcut");
   mcut->Add(mamp);
   mcut->ZeroSuppress(riothresh);
   mcut->Despeckle(specsize);
#if USE_RECTANGULAR_RIO
   mcut->ZeroSuppress(1e30,1);
   TH1D *mcut_px = mcut->ProjectionX();
   TH1D *mcut_py = mcut->ProjectionY();
   int ixmin=0;
   int ixmax=0;
   int iymin=0;
   int iymax=0;
   for (int ix=1; ix <= nbinx; ++ix) {
      if (mcut_px->GetBinContent(ix) != 0)
         ixmax = ix;
      else if (ixmax == 0)
         ixmin = ix + 1;
   }
   for (int iy=1; iy <= nbiny; ++iy) {
      if (mcut_py->GetBinContent(iy) != 0)
         iymax = iy;
      else if (iymax == 0)
         iymin = iy + 1;
   }
   ixmin = (xmin == 0)? ixmin : mcut->GetXaxis()->FindBin(xmin);
   ixmax = (xmax == 0)? ixmax : mcut->GetXaxis()->FindBin(xmax);
   iymin = (ymin == 0)? iymin : mcut->GetYaxis()->FindBin(ymin);
   iymax = (ymax == 0)? iymax : mcut->GetYaxis()->FindBin(ymax);
   printf("ixmin=%d,ixmax=%d,iymin=%d,iymax=%d\n",ixmin,ixmax,iymin,iymax);
   mcut->Reset();
   TH1D* mamp_py = mamp->ProjectionY();
   for (int iy = iymin; iy <= iymax; ++iy) {
      double binsum = mcut_py->Integral(iy-150,iy+150);
      double ampsum = mamp_py->Integral(iy-150,iy+150);
      double minthresh = (ampsum / binsum) * pctthresh/100;
      for (int ix = ixmin; ix <= ixmax; ++ix) {
         double a = mamp->GetBinContent(ix,iy);
         if (a > finthresh && a > minthresh)
            mcut->SetBinContent(ix, iy, 1);
         else
            mcut->SetBinContent(ix, iy, 0);
      }
   }
   mcut->Despeckle(specsize);
#endif
   mmu->Mask(mcut);
   mmu->Rescale(1,1,rcp->GetMicroradPerStep());
   mmu->GetXaxis()->SetTitle("u (mm)");
   mmu->GetYaxis()->SetTitle("v (mm)");
   mmu->GetZaxis()->SetTitle("#mu^{}rad");
   mmu->GetZaxis()->SetTitleOffset(1.6);
   msigma->Mask(mcut);
   msigma->Rescale(1,1,rcp->GetMicroradPerStep());
   msigma->GetXaxis()->SetTitle("u (mm)");
   msigma->GetYaxis()->SetTitle("v (mm)");
   msigma->GetZaxis()->SetTitle("#mu^{}rad");
   msigma->GetZaxis()->SetTitleOffset(1.6);

   if (reverse_images) {
      mbase->Rotate(0,0,180,1);
      mamp->Rotate(0,0,180,1);
      mmu->Rotate(0,0,180,1);
      msigma->Rotate(0,0,180,1);
      mcut->Rotate(0,0,180,1);
   }
   double xmean = mcut->GetMean(1);
   double ymean = mcut->GetMean(2);
   double xshift = plot_size_mm/2 - xmean;
   double yshift = plot_size_mm/2 - ymean;
   if (fabs(xshift) > 0.2 || fabs(yshift) > 0.2) {
      mbase->Shift(xshift,yshift,0);
      mamp->Shift(xshift,yshift,0);
      mmu->Shift(xshift,yshift,0);
      msigma->Shift(xshift,yshift,0);
      mcut->Shift(xshift,yshift,0);
   }

   TString name, title;
   name.Form("hscan%d", scan);
   title.Form("%s scan %d",file,scan);
   TH2D* hcol = new TH2D(name,title,2000,0.,plot_size_mm,2000,0.,plot_size_mm);
   Map2D* mcol = new Map2D(*hcol);
   name.Form("mscan%d", scan);
   mcol->SetName(name);
   mcol->Fill(mmu);
   mcol->Center();
   mcol->Rescale(1,1/cos(THETA_BRAGG_DEG*3.1416/180));

   Map2D* mcol2 = new Map2D(*hcol);
   name.Form("m2scan%d", scan);
   mcol2->SetName(name);
   mcol2->Fill(msigma);
   mcol2->Center();
   mcol2->Rescale(1,1/cos(THETA_BRAGG_DEG*3.1416/180));

   TCanvas *c1 = (TCanvas*)gROOT->FindObject("c1");
   if (c1 == 0) {
      c1 = new TCanvas("c1","c1",5,10,550,500);
      c1->SetRightMargin(0.20);
   }
   c1->cd();
   mcol->GetXaxis()->SetTitle("u (mm)");
   mcol2->GetXaxis()->SetTitle("u (mm)");
   mcol->GetYaxis()->SetTitle("v (mm)");
   mcol2->GetYaxis()->SetTitle("v (mm)");
   mcol->GetZaxis()->SetTitle("#mu^{}rad");
   mcol2->GetZaxis()->SetTitle("#mu^{}rad");
   mcol->GetZaxis()->SetTitleOffset(1.6);
   mcol2->GetZaxis()->SetTitleOffset(1.6);
   mcol->SetContour(100);
   mcol2->SetContour(100);
   if (max_mu > 0) {
      mcol->SetMinimum(0);
      mcol->SetMaximum(max_mu);
   }
   else {
      mcol->Normalize();
      if (max_mu < 0)
         mcol->SetMinimum(0);
   }
   if (max_sigma > 0) {
      mcol2->SetMinimum(0);
      mcol2->SetMaximum(max_sigma);
   }
   else {
      mcol2->Normalize();
      if (max_sigma < 0)
         mcol2->SetMinimum(0);
   }
   if (mcut->GetMaximum() > 0) {
      mcol->Draw("colz");
   }
   return mcol;
}

Map2D *rocksigma(const char* file, int scan)
{
   TString name;
   name.Form("m2scan%d", scan);
   return (Map2D*)gROOT->FindObject(name);
}

TCanvas *rockpixel()
{
   return (TCanvas*)gROOT->FindObject("c2");
}

Map2D *downsample(Map2D *src, Int_t nbins)
{
   TString name;
   name.Form("%s_%d",src->GetName(),nbins);
   Double_t xmin = src->GetXaxis()->GetXmin();
   Double_t xmax = src->GetXaxis()->GetXmax();
   Double_t ymin = src->GetYaxis()->GetXmin();
   Double_t ymax = src->GetYaxis()->GetXmax();
   TH2D *h = new TH2D(name,src->GetTitle(),nbins,xmin,xmax,nbins,ymin,ymax);
   Map2D *map = new Map2D(*h);
   map->SetDirectory(0);
   map->SetName(name);
   map->Fill(src);
   return map;
}

rcpicker *picker(const char *file, int scan)
{
   TFile *f;
   TString resultsfile;
   TString rootURL(xrootd_url + data_dir + run_period + "/results");
#if RESULTS_ON_PNFS
   resultsfile.Form("%s/%s_%3.3d_results.root",rootURL.Data(),file,scan);
#else
   resultsfile.Form("%s/%s_%3.3d_results.root",".",file,scan);
#endif
   f = TFile::Open(resultsfile.Data());
   if (! f->IsOpen()) {
      std::cerr << "Error - unable to open input root file \"" << resultsfile
                << "\"" << std::endl;
      return 0;
   }
   TString rcfile;
   rcfile.Form("%s/%s_%3.3d_rocking_curves.root",rootURL.Data(),file,scan);
   TString logfile;
   logfile.Form("%s/%s/%s",run_period.Data(),file,file);
   rcpicker *rcp = new rcpicker(file, scan, rcfile.Data(), logfile.Data());
   f->cd();
   return rcp;
}

void set_run_period(const char *pname)
{
   run_period = TString(pname);
}

TString get_run_period()
{
   return run_period;
}

void set_plot_size(double size_mm)
{
   plot_size_mm = size_mm;
}

double get_plot_size()
{
   return plot_size_mm;
}

void set_reverse_images(int rev)
{
   reverse_images = rev;
}

int get_reverse_images()
{
   return reverse_images;
}