/*
  *  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/Module.h>
 
 #include <vector>
 
 #include <momemta/config.h>
 #include <momemta/ParameterSet.h>
 #include <momemta/Solution.h>
 
 #include <Path.h>
 
 #ifdef DEBUG_TIMING
 #include <chrono>
 using namespace std::chrono;
 #endif
 
 #define CALL(X) { for (auto& m: path.modules()) \
         m->X(); \
     }
 
 class Looper: public Module {
     public:
 
         Looper(PoolPtr pool, const ParameterSet& parameters): Module(pool, parameters.getModuleName()) {
             solutions = pool->get<SolutionCollection>(parameters.get<InputTag>("solutions"));
 
             path = parameters.get<Path>("path");
         };
 
         virtual void configure() override {
             CALL(configure);
         }
 
         virtual void beginIntegration() override {
             CALL(beginIntegration);
         }
 
         virtual void endIntegration() override {
             CALL(endIntegration);
 
 #ifdef DEBUG_TIMING
             LOG(info) << "Time spent evaluating modules of looper " << name() << ":";
             for (auto it: m_timings) {
                 LOG(info) << "    " << it.first->name() << ": " << duration_cast<duration<double>>(it.second).count() << "s";
             }
 #endif
         }
 
         virtual void beginPoint() override {
             CALL(beginPoint);
         }
 
         virtual void endPoint() override {
             CALL(endPoint);
         }
 
         virtual Status work() override {
             particles->clear();
 
             CALL(beginLoop);
 
             auto status = Status::OK;
 
             // For each solution, loop over all the modules
             for (const auto& s: *solutions) {
                 if (!s.valid)
                     continue;
 
                 *particles = s.values;
                 *jacobian = s.jacobian;
 
                 for (auto& m: path.modules()) {
 #ifdef DEBUG_TIMING
                     auto start = high_resolution_clock::now();
 #endif
                     auto module_status = m->work();
 #ifdef DEBUG_TIMING
                     m_timings[m.get()] += high_resolution_clock::now() - start;
 #endif
 
                     if (module_status == Status::OK)
                         continue;
                     else if (module_status == Status::NEXT)
                         break;
                     else {
                         status = module_status;
                         break;
                     }
                 }
 
                 if (status != Status::OK)
                     break;
             }
 
             CALL(endLoop);
 
             return status;
         }
 
     private:
         Path path;
 
         // Inputs
         Value<SolutionCollection> solutions;
 
         // Outputs
         std::shared_ptr<std::vector<LorentzVector>> particles = produce<std::vector<LorentzVector>>("particles");
         std::shared_ptr<double> jacobian = produce<double>("jacobian");
 
 #ifdef DEBUG_TIMING
         std::unordered_map<Module*, std::chrono::high_resolution_clock::duration> m_timings;
 #endif
 
 };
 
 REGISTER_MODULE(Looper)
     .Input("solutions")
     .Output("particles")
     .Output("jacobian")
     .Attr("path:path");