#define rcshifter_cxx
// This is a utility class for correcting for target walk across a single
// rocking curve scan, which occurs as a result of the target being not
// exactly centered on the rotation axis of the theta stage. This class
// assumes that rcmaker (see rcmaker.C) has already been run on the raw
// image data to produce the rctree for this scan in a root file saved
// on local disk. The output from rcshifter is a modified instance of
// rctree written to the same root file and directory where the input
// rctree was found. The output object has the same name rctree, so that
// any consumers of the rctree will find and use the shifted instance
// unless explicitly directed to get the unshifted instance by a lookup
// with a versioned key.
//
// To use this macro, first you should set the walk constant defined as
// WALK_PIXELS_PER_STEP below, then execute rcshifter from within the
// directory containing your SAMPLE_nnn_rocking_curves.root root file
// as illustrated in the following example.
//
//  Root > .L rcshifter.C+O
//  Root > rcshifter rcs("JD70-104_010_rocking_curves.root");
//  Root > rcs.Loop();
//  Root > rcs.save();

#include "rcshifter.h"
#include <iostream>
#include <TKey.h>

// value adapted for JD70-104 scans, cls-6-2019
#define WALK_PIXELS_PER_STEP 0.0813

rcshifter::rcshifter(const char *rootfile) : fChain(0) 
{
   rcfile = new TFile(rootfile,"update");
   rctree = (TTree*)rcfile->Get("rctree");
   if (rctree == 0) {
      std::cerr << "Error in rcshifter constructor - no rctree found"
                << " in input file " << rootfile << std::endl;
      return;
   }
   rcrows = rccolumns = rcsteps = 0;
   TIter nextkey(rcfile->GetListOfKeys());
   TKey *key;
   while ((key = (TKey*)nextkey())) {
      TString name(key->GetName());
      if (name(0,6) == TString("pixel_")) {
         rchist_corr = (TH1D*)key->ReadObj();
         rcsteps = rchist_corr->GetNbinsX();
         sscanf(name.Data(), "pixel_%d_%d", &rccolumns, &rcrows);
         ++rccolumns;
         ++rcrows;
         break;
      }
   }
   if (rcsteps == 0 || rctree->GetEntries() != rcrows*rccolumns) {
      std::cerr << "Error in rcshifter constructor - unable to find"
                << " a consistent set of dimensions for the rctree in "
                << rootfile << std::endl;
      return;
   }

   TString title(rctree->GetTitle());
   rctree_corr = new TTree("rctree_corr",title + ", walk-corrected");
   rctree_corr->Branch("rchist",rchist_corr,128000,0);

   // for efficiency we try to cache the entire tree in memory
   // Init(rctree);
   rctree->SetBranchAddress("rchist",&rchist,&b_rchist);
   int nentries = rctree->GetEntries();
   int entry = 0;
   for (int j=0; j<rcrows; ++j) {
      for (int i=0; i<rccolumns; ++i) {
         rchist = 0;
         rctree->GetEntry(entry++);
         srchist.push_back(rchist);
      }
   }
}

void rcshifter::Loop()
{
   Long64_t entry = -1;
   int j0 = WALK_PIXELS_PER_STEP * rcsteps/2;
   for (int j=0; j<rcrows; ++j) {
      for (int i=0; i<rccolumns; ++i) {
         rchist_corr->Reset();
         int entry_ij = i+j*rccolumns;
         rchist_corr->SetName(srchist[entry_ij]->GetName());
         rchist_corr->SetTitle(srchist[entry_ij]->GetTitle());
         for (int k=0; k<rcsteps; ++k) {
            int jk = j + k * WALK_PIXELS_PER_STEP - j0;
            jk = (jk < 0)? 0 : (jk >= rcrows)? rcrows-1 : jk;
            int entry_ijk = i + jk * rccolumns;
            double h0, h1;
            if (k < rccolumns/2) {
               h0 = (srchist[entry_ij]->GetBinContent(1) +
                     srchist[entry_ij]->GetBinContent(2) +
                     srchist[entry_ij]->GetBinContent(3))/3;
               h1 = (srchist[entry_ijk]->GetBinContent(1) +
                     srchist[entry_ijk]->GetBinContent(2) +
                     srchist[entry_ijk]->GetBinContent(3))/3;
            }
            else {
               h0 = (srchist[entry_ij]->GetBinContent(rcsteps-2) +
                     srchist[entry_ij]->GetBinContent(rcsteps-1) +
                     srchist[entry_ij]->GetBinContent(rcsteps))/3;
               h1 = (srchist[entry_ijk]->GetBinContent(rcsteps-2) +
                     srchist[entry_ijk]->GetBinContent(rcsteps-1) +
                     srchist[entry_ijk]->GetBinContent(rcsteps))/3;
            }
            double h = srchist[entry_ijk]->GetBinContent(k+1);
            rchist_corr->SetBinContent(k+1,h-h1+h0);
         }
         rctree_corr->Fill();
      }
   }
}

void rcshifter::save()
{
   rctree_corr->Write();
   rctree_corr->Write("rctree");
}

void rcshifter::drawold(int i, int j)
{
   int entry = (i-1) + (j-1)*rccolumns;
   if (entry >= 0 && entry < rctree->GetEntries())
      srchist[entry]->Draw("colz");
   else
      printf("pixel is out of range!\n");
}

void rcshifter::drawnew(int i, int j)
{
   int entry = (i-1) + (j-1)*rccolumns;
   if (entry >= 0 && entry < rctree_corr->GetEntries()) {
      rctree_corr->GetEntry(entry);
      rchist_corr->Draw("colz");
   }
   else
      printf("pixel is out of range!\n");
}