////////////////////////////////////////////////////////////////////
/// file ApplyEnergySystematics.cc
///
/// brief: Examples of applying energy systematics
///
/// author Eric Marzece <marzece@hep.upenn.edu
///
/// detail:
/// These functions will look at RAT-ROOT DS files and apply some
/// energy systematic values and evaluate the effect it has on
/// the energy distribution of the given file.
/// REVISION HISTORY:\n
///             June 10 2018 - New file
///
////////////////////////////////////////////////////////////////////

#include <RAT/DU/DSReader.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/EV.hh>
#include <RAT/DS/FitResult.hh>
#include <RAT/DB.hh>
#include <RAT/DU/EnergySystematics.hh>
#include <RAT/DU/ReconCalibrator.hh>

#include <TH1D.h>
#include <TCanvas.h>
#include <TRandom2.h>

#include <string>
#include <algorithm>
using std::string;

// Values take from DocDB 5083 v10 slide 24
// Make sure these values are appropriate for your purposes before using
const static double escale_uncertainty = 0.020; // 2.0%
const static double resolution_uncertainty = 0.018; // 1.8%

// Convenience function for getting the energy then applying calibrations to it
// This is only valid for rat version 6.5.x
//
// Also, should apply RSP energy correction to the energy if using rat version
// below 6.15.0.
// See examples/root/PlotFit::PlotFitEnergy for how to do that.
double GetCalibratedEnergy(const RAT::DS::FitVertex& fitVertex, const bool& isData,
                           const RAT::DU::ReconCalibrator* calibrator) {
      // Get energy and apply calibrations to it
      double energy = fitVertex.GetEnergy();
      double z = fitVertex.GetPosition().Z();
      double x = fitVertex.GetPosition().X();
      double y = fitVertex.GetPosition().Y();
      double rho = sqrt(x*x + y*y);
      return calibrator->CalibrateEnergyRSP(isData, energy, rho, z);
}

// Provide a input filename for a RAT-ROOT data file, a name for the for the
// energy fitter that should be used (e.g. waterFitter, scintFitter), and
// a name for a filename for output histograms to be written to
// This function shows "before" and "after" distributions for applying
// energy scale/resolution systematic uncertainties.
void ApplyEnergySystematics(const string& fileName, const string& output_filename,
                           const string& fitName) {
  // If this is being done on data that does not require remote database connection
  // eg.: a simple simulation with default run number (0)
  // We can disable the remote connections:
  //
  // NOTE: Don't do this if you are using real data!!!
  RAT::DB::Get()->SetAirplaneModeStatus(true);

  RAT::DU::DSReader dsReader(fileName);
  const bool isData = !dsReader.GetRun().GetMCFlag();
  RAT::DU::ReconCalibrator* calibrator = RAT::DU::ReconCalibrator::Get();

  // Create histograms for the best fit energy and for the energy scaled
  // up (plus) and down (minus) and smeared by the energy resolution uncertainty.
  TH1D* energy_hist = new TH1D("nominal_energy", "nominal energy", 500, 0, 10.0);
  TH1D* scaled_hist_p = new TH1D("scaled_energy_p", "scaled energy p", 500, 0, 10.0);
  TH1D* scaled_hist_m = new TH1D("scaled_energy_m", "scaled energy m", 500, 0, 10.0);
  TH1D* smeared_hist = new TH1D("smeared_energy", "smeared energy", 500, 0, 10.0);

  RAT::DU::EnergySystematics sys(escale_uncertainty, resolution_uncertainty);

  // Event Loop
  for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ ){
    const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );

    for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ ) {
      const RAT::DS::EV& rEV = rDS.GetEV( iEV );
      if(!rEV.FitResultExists(fitName)) {
          continue;
      }
      RAT::DS::FitResult fitResult = rEV.GetFitResult(fitName);
      RAT::DS::FitVertex fitVertex = fitResult.GetVertex(0);
      if(!fitVertex.ContainsEnergy() || !fitVertex.ValidEnergy() ||
         !fitVertex.ContainsPosition() || !fitVertex.ValidPosition()) {
          continue;
      }

      // Get energy and apply calibrations to it
      double energy = GetCalibratedEnergy(fitVertex, isData, calibrator);

      double scaled_energy_p = sys.ApplyEnergyScale(energy, false);
      double scaled_energy_m = sys.ApplyEnergyScale(energy, true);
      double smeared_energy = sys.ApplyEnergyResolution(energy);

      energy_hist->Fill(energy);
      scaled_hist_p->Fill(scaled_energy_p);
      scaled_hist_m->Fill(scaled_energy_m);
      smeared_hist->Fill(smeared_energy);
    }
  }
  TCanvas* c = new TCanvas();
  c->Divide(2,2);
  c->cd(1);
  energy_hist->Draw();
  c->cd(2);
  smeared_hist->Draw();
  c->cd(3);
  scaled_hist_p->Draw();
  c->cd(4);
  scaled_hist_m->Draw();
  c->Print(output_filename.c_str());

  delete energy_hist;
  delete scaled_hist_p;
  delete scaled_hist_m;
  delete smeared_hist;
}

// This shows a simple method of evaluating the effect of uncertainty on
// energy scale. It shifts the energy distribution up and down and gives the
// size of the effect of the systematic as the difference of those two with the
// nominal distribution.
void EvaluateEnergyScale_Simple(const string& fileName, const string& fitName,
        const double &elow, const double& ehigh) {
    RAT::DB::Get()->SetAirplaneModeStatus(true);

    RAT::DU::DSReader dsReader(fileName);
    const bool isData = !dsReader.GetRun().GetMCFlag();
    RAT::DU::ReconCalibrator* calibrator = RAT::DU::ReconCalibrator::Get();

    // Counter to get the total number of valid fits
    unsigned int valid_fit_count = 0;
    // Fraction of events that pass energy cut
    double nominal_efficiency = 0.0;
    // Fraction of events that pass energy cut when the plus energy scale systematic is applied
    double efficiency_p = 0.0;
    // Fraction of events that pass energy cut when the minus energy scale systematic is applied
    double efficiency_m = 0.0;

    // Object for applying energy scale systematic
    RAT::DU::EnergySystematics sys(escale_uncertainty, 0);

    // Start of event loop
    for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ ) {
        const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );

        for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ ) {
            const RAT::DS::EV& rEV = rDS.GetEV( iEV );
            if(!rEV.FitResultExists(fitName)) {
                continue;
            }
            RAT::DS::FitResult fitResult = rEV.GetFitResult(fitName);
            RAT::DS::FitVertex fitVertex = fitResult.GetVertex(0);
            if(!fitVertex.ContainsEnergy() || !fitVertex.ValidEnergy() ||
               !fitVertex.ContainsPosition() || !fitVertex.ValidPosition()) {
                continue;
            }
            valid_fit_count += 1;

            double energy = GetCalibratedEnergy(fitVertex, isData, calibrator);
            if(energy > elow && energy < ehigh) {
                nominal_efficiency+=1;
            }
            double new_energy_p = sys.ApplyEnergyScale(energy);
            if(new_energy_p > elow && new_energy_p < ehigh) {
                efficiency_p += 1;
            }
            double new_energy_m = sys.ApplyEnergyScale(energy, true);
            if(new_energy_m > elow && new_energy_m < ehigh) {
                efficiency_m += 1;
            }
        }
    } // End of event loop
    // Change counts to efficiencies by dividing by total number of valid events
    nominal_efficiency /= (double) valid_fit_count;
    efficiency_p = efficiency_p/(double)valid_fit_count - nominal_efficiency;
    efficiency_m = efficiency_m/(double) valid_fit_count - nominal_efficiency;

    if(efficiency_p < efficiency_m) {
        std::swap(efficiency_p, efficiency_m);
    }

    // Output results
    printf("Energy cut efficiency =%0.2f +- (%0.2f, %0.2f)\n", nominal_efficiency,
                                                                efficiency_p,
                                                                efficiency_m);
}

// This code evaluates the 1-sigma systematic bounds on the efficiency
// of the energy cut defined by the thresholds elow and ehigh. It does so
// by sampling energy scales from a guassian (whose width is the energy scale
// uncertainty) and gives the result (and error) as the inner 1-sigma
// of sampled efficiencies
void EvaluateEnergyScale_Fancy(const string& fileName, const string& fitName,
                         const double& elow, const double &ehigh) {

    RAT::DB::Get()->SetAirplaneModeStatus(true);

    RAT::DU::DSReader dsReader(fileName);
    const bool isData = !dsReader.GetRun().GetMCFlag();
    RAT::DU::ReconCalibrator* calibrator = RAT::DU::ReconCalibrator::Get();

    // Random number generator for gaussian sampling
    TRandom2 root_rng;
    const int N_SAMPLES = 1000;

    // Counter to get the total number of valid fits
    unsigned int valid_fit_count = 0;
    // Fraction of events that pass energy cut
    double nominal_efficiency = 0.0;

    // Array of objects for applying many different energy scales
    std::vector<RAT::DU::EnergySystematics> sys;
    // Array of efficiencies to keep track of what efficiency each scale gives you
    std::vector<double> cut_efficiencies;
    cut_efficiencies.resize(N_SAMPLES);

    // Initialize arrays and create a bunch of objects for applying different
    // energy scales. Where the energy scales applied are sampled from a gaussian
    // centered at 0 with the width of the energy scale uncertainty;
    for(int i=0; i < N_SAMPLES; i++) {
        double escale = root_rng.Gaus(0.0, escale_uncertainty);
        // Set the resolution values to zero b/c we only care about energy scale here.
        sys.push_back(RAT::DU::EnergySystematics(escale, 0));
        cut_efficiencies[i] = 0.0;
    }

    // Start of event loop
    for( size_t iEntry = 0; iEntry < dsReader.GetEntryCount(); iEntry++ ) {
        const RAT::DS::Entry& rDS = dsReader.GetEntry( iEntry );

        for( size_t iEV = 0; iEV < rDS.GetEVCount(); iEV++ ) {
            const RAT::DS::EV& rEV = rDS.GetEV( iEV );

            if(!rEV.FitResultExists(fitName)) {
                continue;
            }
            RAT::DS::FitResult fitResult = rEV.GetFitResult(fitName);
            RAT::DS::FitVertex fitVertex = fitResult.GetVertex(0);
            if(!fitVertex.ContainsEnergy() || !fitVertex.ValidEnergy() ||
               !fitVertex.ContainsPosition() || !fitVertex.ValidPosition()) {
                continue;
            }
            double energy = GetCalibratedEnergy(fitVertex, isData, calibrator);
            valid_fit_count += 1;
            if(energy > elow && energy < ehigh) {
                nominal_efficiency+=1;
            }
            // Loop over all the energy scales and apply them to this event
            // And if it falls within the energy bounds, count it towards the
            // efficiency for this energy scale.
            for(int i=0; i < N_SAMPLES; i++) {
                double new_energy = sys[i].ApplyEnergyScale(energy);
                if(new_energy > elow && new_energy < ehigh) {
                    cut_efficiencies[i] += 1;
                }
            }
        }
    } // End of event loop
    // Change counts to efficiencies by dividing by total number of valid events
    nominal_efficiency /= (double) valid_fit_count;
    for(int i=0; i < N_SAMPLES; i++) {
        cut_efficiencies[i] /= (double) valid_fit_count;
    }


    // Calculate the lower and upper (1-sigma) bound on the energy efficiency.
    // I do so here by simply finding edges the middle 68.2% of samples.
    // This is not necessarily a good way to find the bounds but for an example
    // it's good enough.
    double one_sigma_fraction = 0.682;
    double lower_bound = (1 - one_sigma_fraction)*N_SAMPLES/2.0;
    double upper_bound = N_SAMPLES - lower_bound;

    // Sort the samples by lowest efficiency to highest efficiency
    // That way the middle 68% of the array should be the 1 sigma range
    // Assuming the distribution is gaussian-ish (which isn't a good assumption)
    std::sort(cut_efficiencies.begin(), cut_efficiencies.end());

    // linear interpolate the edge values
    double lower_val = cut_efficiencies[floor(lower_bound)];
    lower_val += (lower_bound - floor(lower_bound))*(cut_efficiencies[ceil(lower_bound)] - lower_val);

    double upper_val = cut_efficiencies[floor(upper_bound)];
    upper_val += (upper_bound - floor(upper_bound))*(cut_efficiencies[ceil(upper_bound)] - upper_val);

    upper_val -= nominal_efficiency;
    lower_val -= nominal_efficiency;

    // Output results
    printf("Energy cut efficiency = %0.2f +- (%0.2f, %0.2f)\n", nominal_efficiency,
                                                                upper_val,
                                                                lower_val);
}