#define rcfitter_cxx
#define RETURN_OUTPUT_HISTOGRAMS_IN_FILE 1
#define SUPERIMPOSE_FIT_ON_PLOTS 1
//#define SCAN_COMBINE_STEPS 4

//
// Normally this class is invoked using the helper script run_rcfitter.C
//
// The class definition in rcfitter.h has been generated automatically
// by the ROOT utility TTree::MakeSelector(). This class is derived
// from the ROOT class TSelector. For more information on the TSelector
// framework see $ROOTSYS/README/README.SELECTOR or the ROOT User Manual.

// The following methods are defined in this file:
//    Begin():        called every time a loop on the tree starts,
//                    a convenient place to create your histograms.
//    SlaveBegin():   called after Begin(), when on PROOF called only on the
//                    slave servers.
//    Process():      called for each event, in this function you decide what
//                    to read and fill your histograms.
//    SlaveTerminate: called at the end of the loop on the tree, when on PROOF
//                    called only on the slave servers.
//    Terminate():    called at the end of the loop on the tree,
//                    a convenient place to draw/fit your histograms.
//
// To use this file, try the following session on your Tree T:
//
// Root > T->Process("rcfitter.C")
// Root > T->Process("rcfitter.C","some options")
// Root > T->Process("rcfitter.C+")
//

#include <sstream>
#include <fstream>
#include <unistd.h>

#include "rcfitter.h"

#include <TFitResultPtr.h>
#include <TFitResult.h>
#include <TObjString.h>
#include <TCanvas.h>
#include <TF1.h>
#include <TH2D.h>
#include <TStyle.h>
#include <TString.h>
#include <TObject.h>
#include <TProof.h>
#include <TEnv.h>

//int gethostname(char *name, int len) {return 0;}

rcfitter::rcfitter(TTree * /*tree*/)
{
   fHbase = 0;
   fHamp = 0;
   fHmu = 0;
   fHsigma = 0;
   fHmean = 0;
   fHrms = 0;
   fHmax = 0;
   fHpeak = 0;
   fHload = 0;
#if ANDOR_CAMERA_PIXEL_DIMENSIONS
   fWidth = 2560;
   fHeight = 2160;
#else
   fWidth = 4000;
   fHeight = 2672;
#endif
   fSlaveId = 0;
   fFile = 0;
   fProofFile = 0;
}

rcfitter::~rcfitter()
{
}

void rcfitter::Begin(TTree * /*tree*/)
{
   // The Begin() function is called at the start of the query.
   // When running with PROOF Begin() is only called on the client.
   // The tree argument is deprecated (on PROOF 0 is passed).

   TString option = GetOption();
}

void rcfitter::SlaveBegin(TTree * /*tree*/)
{
   // The SlaveBegin() function is called after the Begin() function.
   // When running with PROOF SlaveBegin() is called on each slave server.
   // The tree argument is deprecated (on PROOF 0 is passed).

   TString option = GetOption();

   TProofOutputFile worker("rcfitter_workout");
   TString worker_host_port(worker.GetLocalHost());
   TString worker_dir(worker.GetDir(kTRUE));
   TString worker_ordinal(worker.GetWorkerOrdinal());
   fSlaveId = worker_ordinal(2,99).String().Atoi();

   if (fHload == 0) {
      fHload = new TH1D("workerId","load vs worker id",31,0,31);
      fHbase = new TH2D("hbase","constant term",fWidth,0,fWidth,fHeight,0,fHeight);
      fHamp = new TH2D("hamp","fit peak height",fWidth,0,fWidth,fHeight,0,fHeight);
      fHmu = new TH2D("hmu","fit peak centroid",fWidth,0,fWidth,fHeight,0,fHeight);
      fHsigma = new TH2D("hsigma","fit peak sigma",fWidth,0,fWidth,fHeight,0,fHeight);
      fHmean = new TH2D("hmean","rocking curve mean",fWidth,0,fWidth,fHeight,0,fHeight);
      fHrms = new TH2D("hrms","rocking curve rms",fWidth,0,fWidth,fHeight,0,fHeight);
      fHmax = new TH2D("hmax","rocking curve max intensity",fWidth,0,fWidth,fHeight,0,fHeight);
      fHpeak = new TH2D("hpeak","rocking curve max angle",fWidth,0,fWidth,fHeight,0,fHeight);
   }
   else {
      fHload->Reset();
      fHbase->Reset();
      fHamp->Reset();
      fHmu->Reset();
      fHsigma->Reset();
      fHmean->Reset();
      fHrms->Reset();
      fHmax->Reset();
      fHpeak->Reset();
   }

#if RETURN_OUTPUT_HISTOGRAMS_IN_FILE
   if (fProofFile == 0) {
      fProofFile = new TProofOutputFile("rcfitter_output.root");
      fOutput->Add(fProofFile);
   }
   fFile = fProofFile->OpenFile("RECREATE");
   if (fFile && fFile->IsZombie()) {
      SafeDelete(fFile);
      Info("SlaveBegin", "recreate '%s': file is a zombie!",
           fProofFile->GetName());
   }
   else if (!fFile) {
      Info("SlaveBegin", "could not create '%s': instance is invalid!",
           fProofFile->GetName());
   }
#else
   fOutput->Add(fHbase);
   fOutput->Add(fHamp);
   fOutput->Add(fHmu);
   fOutput->Add(fHsigma);
   fOutput->Add(fHmean);
   fOutput->Add(fHrms);
   fOutput->Add(fHmax);
   fOutput->Add(fHpeak);
   fOutput->Add(fHload);
#endif

}

Bool_t rcfitter::Process(Long64_t entry)
{
   // The Process() function is called for each entry in the tree (or possibly
   // keyed object in the case of PROOF) to be processed. The entry argument
   // specifies which entry in the currently loaded tree is to be processed.
   // It can be passed to either rcfitter::GetEntry() or TBranch::GetEntry()
   // to read either all or the required parts of the data. When processing
   // keyed objects with PROOF, the object is already loaded and is available
   // via the fObject pointer.
   //
   // This function should contain the "body" of the analysis. It can contain
   // simple or elaborate selection criteria, run algorithms on the data
   // of the event and typically fill histograms.
   //
   // The processing can be stopped by calling Abort().
   //
   // Use fStatus to set the return value of TTree::Process().
   //
   // The return value is currently not used.

   b_rchist->GetEntry(entry);

#if SCAN_COMBINE_STEPS
   rchist->Rebin(SCAN_COMBINE_STEPS);
#endif

   TString option = GetOption();
   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();
   fHload->Fill(fSlaveId,1);
   if (name.Contains(option)) {
      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);
      fHmean->Fill(iu,iv,ymean);
      fHrms->Fill(iu,iv,yrms);
      fHmax->Fill(iu,iv,ymax);
      fHpeak->Fill(iu,iv,xmax);
      if (yrms > 0.1) {
         TFitResultPtr result = FitRC(rchist);
         if (result->Status() == 0) {
            fHbase->Fill(iu,iv,result->Value(0));
            fHamp->Fill(iu,iv,result->Value(1));
            fHmu->Fill(iu,iv,result->Value(2));
            fHsigma->Fill(iu,iv,result->Value(3));
         }
      }
   }

   return kTRUE;
}

TFitResultPtr rcfitter::FitRC(Long64_t entry)
{
   b_rchist->GetEntry(entry);
   return FitRC(rchist);
}

#if FITRC_ALGORITHM_0

TFitResultPtr rcfitter::FitRC(TH1D *rc)
{
   // Fit a single peak on top of a flat background.

   Int_t nbins = rc->GetNbinsX();
   Double_t ymin = rc->GetMinimum();
   Double_t ymax = rc->GetMaximum();
   Double_t ymean = rc->GetSum()/nbins;

   // Pick out the region around the peak for fitting.

   Double_t sum[3] = {1e-30,1e-30,1e-30};
   Int_t range[2] = {0,1};
   Double_t sumy = 0;
   Double_t sumy2 = 0;
   for (Int_t ibin=1; ibin <= nbins; ++ibin) {
      Double_t y = rc->GetBinContent(ibin);
      if (y > 0.95*ymean + 0.05*ymax) {
         Double_t w = y-ymean;
         sum[0] += w;
         sum[1] += w*ibin;
         sum[2] += w*ibin*ibin;
         range[0] = (range[0] == 0)? ibin : range[0];
         range[1] = ibin;
      }
      sumy += y;
      sumy2 += y*y;
   }
   Double_t mean = sum[1]/sum[0];
   Double_t rms = sum[2]/sum[0]-mean*mean;
   rms = (rms > 0)? sqrt(rms) : 0;
   rms = (rms < 5)? 5 : (rms > 100)? 100 : rms;
   Int_t win = int(5*rms);
   range[0] = (range[0] > win)? range[0]-win : 1;
   range[1] = (range[1] < nbins-win)? range[1]+win : nbins;
   Double_t xrange[2];
   xrange[0] = rc->GetXaxis()->GetBinCenter(range[0]);
   xrange[1] = rc->GetXaxis()->GetBinCenter(range[1]);
   Int_t peak = rc->GetMaximumBin();
   Double_t xpeak = rc->GetXaxis()->GetBinCenter(peak);
   Double_t xrms = rms * rc->GetXaxis()->GetBinWidth(1);
   Double_t xrms0 = 9 * rc->GetXaxis()->GetBinWidth(1);
   Double_t xrms_min = 1.5 * rc->GetXaxis()->GetBinWidth(1);

   // Model the peak as a Gaussian plus a flat background

   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,xrms0+xrms*0.01);
   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]);

   // Do the fit and return the fit results.  Option flags are:
   //   S - return results of the fit as a TFitResult object
   //   Q - quiet mode (minimum printing)
   //   0 - do not plot the result of the fit
   //   R - limit the fit to the range specified for the fit function

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

#else

TFitResultPtr rcfitter::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.);

   // Find starting values, bounds for the peak fit
 
   Double_t xpeak, xpeak_bound[2];
   Double_t xsigma, xsigma_bound[2];

   // case 1) There is a clear peak visible over the noise,
   //         form a fit window around the peak maximum.
 
   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;
   }
 
   // case 2) There is no clear peak visible over the noise,
   //         use a scanning window to search for the best fit region.
 
   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;
   }

   // Model the peak as a Gaussian plus a flat background

   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]);

   // Do the fit and return the fit results.  Option flags are:
   //   S - return results of the fit as a TFitResult object
   //   Q - quiet mode (minimum printing)
   //   0 - do not plot the result of the fit
   //   R - limit the fit to the range specified for the fit function

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

#endif

void rcfitter::PlotRC(Long64_t entry)
{
   b_rchist->GetEntry(entry);
   return PlotRC(rchist);
}

void rcfitter::PlotRC(TH1D *rc)
{
#if SUPERIMPOSE_FIT_ON_PLOTS
   rc->ResetStats();
   if (rc->GetEntries() > 0) {
      rcfitter f;
      TFitResultPtr result = f.FitRC(rc);
      rc->GetFunction("rcshape")->ResetBit(1<<9);
   }
#endif
   TCanvas *c1 = (TCanvas*)gROOT->FindObject("c1");
   if (c1 == 0) {
      c1 = new TCanvas("c1","c1",5,5,550,500);
   }
   rc->Draw();
}

void rcfitter::SlaveTerminate()
{
   // The SlaveTerminate() function is called after all entries or objects
   // have been processed. When running with PROOF SlaveTerminate() is called
   // on each slave server.

#if RETURN_OUTPUT_HISTOGRAMS_IN_FILE
   if (fFile) {
      TDirectory *savedir = gDirectory;
      fFile->cd();
      fHbase->Write();
      fHamp->Write();
      fHmu->Write();
      fHsigma->Write();
      fHmean->Write();
      fHrms->Write();
      fHmax->Write();
      fHpeak->Write();
      fHload->Write();
      savedir->cd();
      fFile->Close();
   }
#endif
}

void rcfitter::Terminate()
{
   // The Terminate() function is the last function to be called during
   // a query. It always runs on the client, it can be used to present
   // the results graphically or save the results to file.

   TFile *fout(0);
   TIter next(fOutput);
   TObject *obj;
   while ((obj = next())) {
#if RETURN_OUTPUT_HISTOGRAMS_IN_FILE
      TString object_type(obj->ClassName());
      if (object_type.EqualTo("TProofOutputFile")) {
         fProofFile = (TProofOutputFile*)(obj);
         TString filePath(fProofFile->GetDir());
         TString fileName(fProofFile->GetFileName());
         filePath += fileName;
         fFile = TFile::Open(filePath);
         TH2D *hbase = (TH2D*)fFile->Get("hbase");
         TH2D *hamp = (TH2D*)fFile->Get("hamp");
         TH2D *hmu = (TH2D*)fFile->Get("hmu");
         TH2D *hsigma = (TH2D*)fFile->Get("hsigma");
         TH2D *hmean = (TH2D*)fFile->Get("hmean");
         TH2D *hrms = (TH2D*)fFile->Get("hrms");
         TH2D *hmax = (TH2D*)fFile->Get("hmax");
         TH2D *hpeak = (TH2D*)fFile->Get("hpeak");
         fout = new TFile(fileName,"RECREATE");
         hbase->Write();
         hamp->Write();
         hmu->Write();
         hsigma->Write();
         hmean->Write();
         hrms->Write();
         hmax->Write();
         hpeak->Write();
         fFile->Close();
      }
#else
      if (fout == 0)
         fout = new TFile("rcfitter_output.root","RECREATE");
      TString obj_type(obj->ClassName());
      std::cout << "saving object of type " << obj_type << std::endl;
      obj->Write();
#endif
   }
   if (fout != 0)
      fout->Close();
}