// rcfitter_rdf.C
//
// RDataFrame-based replacement for the old PROOF-based rcfitter TSelector.
// PROOF (TProofOutputFile, TProof, etc.) has been fully removed from
// current ROOT installations, so this rewrites the same per-pixel
// rocking-curve fit analysis using RDataFrame::ForeachSlot, which gives
// the same "one set of accumulator histograms per worker" structure that
// PROOF's SlaveBegin/Process/SlaveTerminate/Terminate lifecycle provided,
// just collapsed into a single process using ROOT's implicit multithreading.
//
// Usage:
//   root [0] .x rcfitter_rdf.C
// or, once you've confirmed correctness, enable parallelism by
// uncommenting the ROOT::EnableImplicitMT() line below.

#include <vector>
#include <string>
#include <iostream>

#include <ROOT/RDataFrame.hxx>
#include <TH1D.h>
#include <TH2D.h>
#include <TF1.h>
#include <TFitResultPtr.h>
#include <TFitResult.h>
#include <TFile.h>
#include <TString.h>
#include <TObjArray.h>
#include <TObjString.h>

#include "rcfitter_rdf.h"

// ---- image dimensions, same as the original fWidth/fHeight ----
#if ANDOR_CAMERA_PIXEL_DIMENSIONS
static const int gWidth = 2560;
static const int gHeight = 2160;
#else
static const int gWidth = 4000;
static const int gHeight = 2672;
#endif

// ---- the same fit function as rcfitter::FitRC(TH1D*), verbatim logic ----
TFitResultPtr FitRC(TH1D *rc)
{
   // Fit a single peak on top of a flat background.

   Int_t nbins = rc->GetNbinsX();
   Double_t binwidth = rc->GetXaxis()->GetBinWidth(1);
   Double_t ymax = 0;
   Double_t ymin = 1e30;
   Double_t sumy = 0;
   Double_t sumy2 = 1./12;
   for (Int_t ibin=1; ibin <= nbins; ++ibin) {
      Double_t y = rc->GetBinContent(ibin);
      sumy += y;
      sumy2 += y*y;
      ymax = (y > ymax)? y : ymax;
      ymin = (y < ymin)? y : ymin;
   }
   Double_t ymean = sumy / nbins;
   Double_t yrms = sqrt(sumy2/nbins - ymean*ymean);
   yrms = (yrms > sqrt(1/12.))? yrms : sqrt(1/12.);

   Double_t xpeak, xpeak_bound[2];
   Double_t xsigma, xsigma_bound[2];

   if (ymax > ymean + 6 * yrms) {
      Int_t imax = rc->GetMaximumBin();
      Int_t ihigh = imax;
      while (++ihigh <= nbins)
         if (rc->GetBinContent(ihigh) < ymean + 3 * yrms)
            break;
      Int_t ilow = imax;
      while (--ilow >= 1)
         if (rc->GetBinContent(ilow) < ymean + 3 * yrms)
            break;
      xpeak = rc->GetXaxis()->GetBinCenter(imax);
      xpeak_bound[0] = xpeak + 5 * (ilow-imax) * binwidth;
      xpeak_bound[1] = xpeak + 5 * (ihigh-imax) * binwidth;
      xsigma = binwidth * (ihigh-ilow) / 2.3;
      xsigma_bound[0] = binwidth;
      xsigma_bound[1] = xsigma * 5;
   }
   else {
      Int_t scan_window = 12;
      Double_t sum_window = 0;
      Double_t max_window = 0;
      Int_t imax = 0;
      for (Int_t ibin=1; ibin <= nbins; ++ibin) {
         sum_window += rc->GetBinContent(ibin);
         if (ibin - scan_window > 0)
            sum_window -= rc->GetBinContent(ibin - scan_window);
         if (sum_window > max_window) {
            max_window = sum_window;
            imax = ibin - scan_window/2;
         }
      }
      xpeak = rc->GetXaxis()->GetBinCenter(imax);
      xpeak_bound[0] = xpeak - scan_window * binwidth;
      xpeak_bound[1] = xpeak + scan_window * binwidth;
      xsigma = binwidth * scan_window / 2.3;
      xsigma_bound[0] = binwidth;
      xsigma_bound[1] = xsigma * 5;
   }

   // NOTE: this TF1 is a fresh local object on every call -- no shared
   // global fit-function state -- which is exactly what makes this safe
   // to call concurrently from multiple threads under EnableImplicitMT().
   TF1 fitfunc("rcshape","[0]+[1]*exp(-0.5*(x-[2])*(x-[2])/([3]*[3]))",
               xpeak_bound[0],xpeak_bound[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,xsigma);
   fitfunc.SetParLimits(0,ymin*0.999,ymax*1.001);
   fitfunc.SetParLimits(1,0.,2*(ymax-ymin)+1e-3);
   fitfunc.SetParLimits(2,xpeak_bound[0],xpeak_bound[1]);
   fitfunc.SetParLimits(3,xsigma_bound[0],xsigma_bound[1]);

   return rc->Fit(&fitfunc,"SQ0R");
}

// Runs the full rocking-curve fit analysis over the tree "rctree" found
// in treeUrl (a local path or an xrootd URL, exactly as TChain::Add()
// would accept). Returns the 8 final result histograms in the order
// base, amp, mu, sigma, mean, rms, max, peak -- the caller can Write()
// them directly, same as the old Terminate() output.
//
// If start/count restrict to a sub-range of entries, RDataFrame::Range()
// is used, which requires falling back to single-threaded processing
// (RDataFrame does not support Range() together with implicit MT) --
// this matches how the old code was normally invoked with the full
// range anyway (count defaults to "effectively all entries").
// The actual per-entry logic, shared between the ranged and non-ranged
// paths below via a template so the same code works whether df's type
// is RDataFrame or the RInterface returned by .Range(...).
template <typename DFType>
static std::vector<TH2D> processFrame(DFType &&df)
{
   unsigned int nSlots = df.GetNSlots();
   std::cout << "Running with " << nSlots << " slot(s)." << std::endl;

   // one set of accumulator histograms per slot -- direct analogue of
   // the old per-PROOF-worker fHbase, fHamp, ... created in SlaveBegin()
   std::vector<TH2D> hbase, hamp, hmu, hsigma, hmean, hrms, hmax, hpeak;
   for (unsigned int s = 0; s < nSlots; ++s) {
      hbase.emplace_back(Form("hbase_%u",s),"constant term",
                          gWidth,0,gWidth,gHeight,0,gHeight);
      hamp.emplace_back(Form("hamp_%u",s),"fit peak height",
                         gWidth,0,gWidth,gHeight,0,gHeight);
      hmu.emplace_back(Form("hmu_%u",s),"fit peak centroid",
                        gWidth,0,gWidth,gHeight,0,gHeight);
      hsigma.emplace_back(Form("hsigma_%u",s),"fit peak sigma",
                           gWidth,0,gWidth,gHeight,0,gHeight);
      hmean.emplace_back(Form("hmean_%u",s),"rocking curve mean",
                          gWidth,0,gWidth,gHeight,0,gHeight);
      hrms.emplace_back(Form("hrms_%u",s),"rocking curve rms",
                         gWidth,0,gWidth,gHeight,0,gHeight);
      hmax.emplace_back(Form("hmax_%u",s),"rocking curve max intensity",
                         gWidth,0,gWidth,gHeight,0,gHeight);
      hpeak.emplace_back(Form("hpeak_%u",s),"rocking curve max angle",
                          gWidth,0,gWidth,gHeight,0,gHeight);
   }

   df.ForeachSlot(
      [&](unsigned int slot, TH1D &rchist_in) {
         // Work on a private copy since we may rebin it, and different
         // threads must never touch the same TH1D concurrently.
         TH1D rchist(rchist_in);

#ifdef SCAN_COMBINE_STEPS
         rchist.Rebin(SCAN_COMBINE_STEPS);
#endif

         TString name(rchist.GetName());
         TObjArray *tokens = name.Tokenize("_,");
         TObjString *arg1 = (TObjString *)(*tokens)[1];
         TObjString *arg2 = (TObjString *)(*tokens)[2];
         Int_t iu = arg1->String().Atoi();
         Int_t iv = arg2->String().Atoi();
         delete tokens;

         Int_t nbins = rchist.GetNbinsX();
         Double_t sum[3] = {1e-30,1e-30,1e-30};
         Double_t xmax=0, ymax=0;
         for (int ibin = 1; ibin <= nbins; ++ibin) {
            double y = rchist.GetBinContent(ibin);
            if (y > ymax) {
               xmax = ibin-1;
               ymax = y;
            }
            sum[0] += 1;
            sum[1] += y;
            sum[2] += y*y;
         }
         Double_t ymean = sum[1]/sum[0];
         Double_t yrms = sqrt(sum[2]/sum[0]-ymean*ymean);
         hmean[slot].Fill(iu,iv,ymean);
         hrms[slot].Fill(iu,iv,yrms);
         hmax[slot].Fill(iu,iv,ymax);
         hpeak[slot].Fill(iu,iv,xmax);
         if (yrms > 0.1) {
            TFitResultPtr result = FitRC(&rchist);
            if (result->Status() == 0) {
               hbase[slot].Fill(iu,iv,result->Value(0));
               hamp[slot].Fill(iu,iv,result->Value(1));
               hmu[slot].Fill(iu,iv,result->Value(2));
               hsigma[slot].Fill(iu,iv,result->Value(3));
            }
         }
      },
      {"rchist"}
   );

   // merge all per-slot histograms -- direct analogue of the old
   // Terminate() method that summed results across PROOF workers
   std::vector<TH2D> out;
   out.emplace_back(hbase[0]);   out.back().SetName("hbase");
   out.emplace_back(hamp[0]);    out.back().SetName("hamp");
   out.emplace_back(hmu[0]);     out.back().SetName("hmu");
   out.emplace_back(hsigma[0]);  out.back().SetName("hsigma");
   out.emplace_back(hmean[0]);   out.back().SetName("hmean");
   out.emplace_back(hrms[0]);    out.back().SetName("hrms");
   out.emplace_back(hmax[0]);    out.back().SetName("hmax");
   out.emplace_back(hpeak[0]);   out.back().SetName("hpeak");
   for (unsigned int s = 1; s < nSlots; ++s) {
      out[0].Add(&hbase[s]);
      out[1].Add(&hamp[s]);
      out[2].Add(&hmu[s]);
      out[3].Add(&hsigma[s]);
      out[4].Add(&hmean[s]);
      out[5].Add(&hrms[s]);
      out[6].Add(&hmax[s]);
      out[7].Add(&hpeak[s]);
   }
   return out;
}

// Runs the full rocking-curve fit analysis over the tree "rctree" found
// in treeUrl (a local path or an xrootd URL, exactly as TChain::Add()
// would accept). Returns the 8 final result histograms in the order
// base, amp, mu, sigma, mean, rms, max, peak -- the caller can Write()
// them directly, same as the old Terminate() output.
//
// If start/count restrict to a sub-range of entries, RDataFrame::Range()
// is used, which requires falling back to single-threaded processing --
// RDataFrame does not support Range() together with implicit MT. The old
// code was normally invoked over the full range anyway (count defaults
// to "effectively all entries"), so this only matters for partial runs.
std::vector<TH2D> RunRcFitter(const char *treeUrl,
                              Long64_t start,
                              Long64_t count)
{
   bool doRange = (start != 0 || count >= 0);
   if (doRange) {
      // Range() is incompatible with implicit MT -- if a prior call in
      // this session already turned MT on, turn it back off for this
      // one, since Range() must run serially regardless.
      if (ROOT::IsImplicitMTEnabled()) {
         ROOT::DisableImplicitMT();
      }
   }
   else {
      // Must happen *before* constructing the RDataFrame below --
      // RDataFrame fixes its slot count at construction time, so
      // enabling MT afterward has no effect.
      ROOT::EnableImplicitMT();
   }
   ROOT::RDataFrame df("rctree", treeUrl);
   if (doRange) {
      // end=0 is RDataFrame's actual "to the end of the dataset" value.
      // Only compute a real end bound if count looks like a genuine,
      // reasonably-sized limit -- a huge sentinel count (as the old
      // PROOF-era code used to mean "no limit") must map to 0, not to
      // a literal astronomical entry number, which Range() does not
      // handle the same way and was causing it to silently stop a
      // little short of the actual end of the tree.
      const Long64_t kSentinelThreshold = 1000000000LL; // 1 billion
      Long64_t end = (count < 0 || count >= kSentinelThreshold) ? 0 : start + count;
      return processFrame(df.Range(start, end));
   }
   else {
      return processFrame(df);
   }
}