/*
  *  MoMEMta: a modular implementation of the Matrix Element Method
  *  Copyright (C) 2016  Universite catholique de Louvain (UCL), Belgium
  *
  *  This program is free software: you can redistribute it and/or modify
  *  it under the terms of the GNU General Public License as published by
  *  the Free Software Foundation, either version 3 of the License, or
  *  (at your option) any later version.
  *
  *  This program is distributed in the hope that it will be useful,
  *  but WITHOUT ANY WARRANTY; without even the implied warranty of
  *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  *  GNU General Public License for more details.
  *
  *  You should have received a copy of the GNU General Public License
  *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
 #include <momemta/Logging.h>
 #include <momemta/Module.h>
 #include <momemta/ParameterSet.h>
 #include <momemta/Types.h>
 #include <momemta/Math.h>
 
 
 #include <TFile.h>
 #include <TH2.h>
 #include <TAxis.h>
 
 class BinnedTransferFunctionOnEnergyBase: public Module {
     public:
 
         BinnedTransferFunctionOnEnergyBase(PoolPtr pool, const ParameterSet& parameters): Module(pool, parameters.getModuleName()) {
             m_reco_input = get<LorentzVector>(parameters.get<InputTag>("reco_particle"));
 
             std::string file_path = parameters.get<std::string>("file");
             std::string th2_name = parameters.get<std::string>("th2_name");
 
             std::unique_ptr<TFile> file(TFile::Open(file_path.c_str()));
             if(!file->IsOpen() || file->IsZombie())
                 throw file_not_found_error("Could not open file " + file_path);
 
             m_th2 = std::unique_ptr<TH2>(static_cast<TH2*>(file->Get(th2_name.c_str())));
             if(!m_th2->InheritsFrom("TH2") || !m_th2.get())
                 throw th2_not_found_error("Could not retrieve object " + th2_name + " deriving from class TH2 in file " + file_path + ".");
             m_th2->SetDirectory(0);
 
             TAxis* yAxis = m_th2->GetYaxis();
             m_deltaMin = yAxis->GetXmin();
             m_deltaMax = yAxis->GetXmax();
             m_deltaRange = m_deltaMax - m_deltaMin;
 
             TAxis* xAxis = m_th2->GetXaxis();
             double E_cut = parameters.get<double>("min_E", 0.);
             m_EgenMin = std::max(xAxis->GetXmin(), E_cut);
             m_EgenMax = xAxis->GetXmax();
 
             // Since we assume the TF continues as a constant for E->infty,
             // we need to be able to retrieve the X axis' last bin, to avoid
             // fetching the TH2's overflow bin.
             m_fallBackEgenMax = xAxis->GetBinCenter(xAxis->GetNbins());
 
             LOG(debug) << "Loaded TH2 " << th2_name << " from file " << file_path << ".";
             LOG(debug) << "\tDelta range is " << m_deltaMin << " to " << m_deltaMax << ".";
             LOG(debug) << "\tEnergy range is " << m_EgenMin << " to " << m_EgenMax << ".";
             LOG(debug) << "\tWill use values at Egen = " << m_fallBackEgenMax << " for out-of-range values.";
 
             file->Close();
             file.reset();
         };
 
     protected:
         std::unique_ptr<TH2> m_th2;
 
         double m_deltaMin, m_deltaMax, m_deltaRange;
         double m_EgenMin, m_EgenMax;
         double m_fallBackEgenMax;
 
         // Input
         Value<LorentzVector> m_reco_input;
 
     private:
         class file_not_found_error: public std::runtime_error{
             using std::runtime_error::runtime_error;
         };
 
         class th2_not_found_error: public std::runtime_error{
             using std::runtime_error::runtime_error;
         };
 };
 
 class BinnedTransferFunctionOnEnergy: public BinnedTransferFunctionOnEnergyBase {
     public:
 
         BinnedTransferFunctionOnEnergy(PoolPtr pool, const ParameterSet& parameters): BinnedTransferFunctionOnEnergyBase(pool, parameters) {
             m_ps_point = get<double>(parameters.get<InputTag>("ps_point"));
         }
 
         virtual Status work() override {
             const double rec_E = m_reco_input->E();
             const double rec_M = m_reco_input->M();
             const double range = GetDeltaRange(rec_E, rec_M);
             const double gen_E = rec_E - GetDeltaMax(rec_E, rec_M) + range * (*m_ps_point);
             const double delta = rec_E - gen_E;
 
             // To change the particle's energy without changing its direction and mass
             // Forcing positive value of (gen_E - rec_M) due to numeric precision issue
             double gen_pt = std::sqrt(std::max(gen_E - rec_M, 0.) * (gen_E + rec_M)) / std::cosh(m_reco_input->Eta());
             output->SetCoordinates(
                     gen_pt * std::cos(m_reco_input->Phi()),
                     gen_pt * std::sin(m_reco_input->Phi()),
                     gen_pt * std::sinh(m_reco_input->Eta()),
                     gen_E);
 
             // The bin number is a ROOT "global bin number" using a 1D representation of the TH2
             const int bin = m_th2->FindFixBin(std::min(gen_E, m_fallBackEgenMax), delta);
 
             // Compute TF*jacobian, where the jacobian includes the transformation of [0,1]->[range_min,range_max] and d|P|/dE
             *TF_times_jacobian = m_th2->GetBinContent(bin) * range * dP_over_dE(*output);
 
             return Status::OK;
         }
 
     private:
 
         // Input
         Value<double> m_ps_point;
 
         // Outputs
         std::shared_ptr<LorentzVector> output = produce<LorentzVector>("output");
         std::shared_ptr<double> TF_times_jacobian = produce<double>("TF_times_jacobian");
 
         // We assume TF=0 for Egen < lower limit of the TH2, and
         //           TF=constant for Egen > upper limit of the TH2
         // In the former case, the integrated range is adapted (shortened) to the range where TF!=0
         // The range also takes into account the particle's given mass: Egen has to be larger than M
         inline double GetDeltaRange(const double Erec, const double Mrec) const {
             return GetDeltaMax(Erec, Mrec) - m_deltaMin;
         }
 
         inline double GetDeltaMax(const double Erec, const double Mrec) const {
             if (Erec <= Mrec || m_deltaMax <= Mrec)
                 return 0;
             return std::min(m_deltaMax, Erec - m_EgenMin) - Mrec;
         }
 };
 
 class BinnedTransferFunctionOnEnergyEvaluator: public BinnedTransferFunctionOnEnergyBase {
     public:
 
         BinnedTransferFunctionOnEnergyEvaluator(PoolPtr pool, const ParameterSet& parameters): BinnedTransferFunctionOnEnergyBase(pool, parameters) {
             m_gen_input = get<LorentzVector>(parameters.get<InputTag>("gen_particle"));
         }
 
         virtual Status work() override {
             const double rec_E = m_reco_input->E();
             const double gen_E = m_gen_input->E();
             double delta = rec_E - gen_E;
             if (delta <= m_deltaMin || delta >= m_deltaMax) {
                 *TF_value = 0;
                 return Status::OK;
             }
 
             // The bin number is a ROOT "global bin number" using a 1D representation of the TH2
             const int bin = m_th2->FindFixBin(std::min(gen_E, m_fallBackEgenMax), delta);
 
             // Retrieve TF value
             *TF_value = m_th2->GetBinContent(bin);
 
             return Status::OK;
         }
 
     private:
 
         // Input
         Value<LorentzVector> m_gen_input;
 
         // Outputs
         std::shared_ptr<double> TF_value = produce<double>("TF");
 };
 
 REGISTER_MODULE(BinnedTransferFunctionOnEnergy)
         .Input("reco_particle")
         .Input("ps_point")
         .Output("output")
         .Output("TF_times_jacobian")
         .Attr("file:string")
         .Attr("th2_name:string")
         .Attr("min_E:double=0");
 
 REGISTER_MODULE(BinnedTransferFunctionOnEnergyEvaluator)
         .Input("reco_particle")
         .Input("gen_particle")
         .Output("TF")
         .Attr("file:string")
         .Attr("th2_name:string")
         .Attr("min_E:double=0");