/*
  *  MoMEMta: a modular implementation of the Matrix Element Method
  *  Copyright (C) 2017  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/ParameterSet.h>
 #include <momemta/Module.h>
 #include <momemta/Solution.h>
 #include <momemta/Math.h>
 #include <momemta/InputTag.h>
 #include <momemta/Types.h>
 
 #include <Math/GenVector/VectorUtil.h>
 
 class SecondaryBlockE: public Module {
     public:
 
         SecondaryBlockE(PoolPtr pool, const ParameterSet& parameters): Module(pool, parameters.getModuleName()),
             sqrt_s(parameters.globalParameters().get<double>("energy")) {
                 s12 = get<double>(parameters.get<InputTag>("s12"));
                 s123 = get<double>(parameters.get<InputTag>("s123"));
 
                 m_p1 = get<LorentzVector>(parameters.get<InputTag>("p1"));
                 m_p2 = get<LorentzVector>(parameters.get<InputTag>("p2"));
                 m_p3 = get<LorentzVector>(parameters.get<InputTag>("p3"));
             };
 
         virtual Status work() override {
 
             solutions->clear();
 
             const double m1 = m_p1->M();
             const double sq_m1 = SQ(m1);
             const double m3 = m_p3->M();
             const double sq_m3 = SQ(m3);
 
             // Don't spend time on unphysical part of phase-space
             if (*s123 >= SQ(sqrt_s) || *s12 + sq_m3 >= *s123 || sq_m1 >= *s12)
                return Status::NEXT;
 
             const double p3 = m_p3->P();
             const double E3 = m_p3->E();
             const double sq_E3 = SQ(E3);
 
             const double c12 = ROOT::Math::VectorUtil::CosTheta(*m_p1, *m_p2);
             const double c13 = ROOT::Math::VectorUtil::CosTheta(*m_p1, *m_p3);
             const double c23 = ROOT::Math::VectorUtil::CosTheta(*m_p2, *m_p3);
 
             double X = p3 * c23 - E3;
             double Y = *s123 - *s12 - sq_m3;
 
             std::vector<double> abs_p1, abs_p2;
             solve2Quads(SQ(X), SQ(p3 * c13) - sq_E3, 2 * p3 * c13 * X,  X * Y, p3 * c13 * Y, 0.25 * SQ(Y) - sq_E3 * sq_m1,
                         2 * X / E3, 0, 2 * (p3 * c13 / E3 - c12), Y / E3, 0, sq_m1 - *s12,
                         abs_p2, abs_p1);
 
             // Use now the obtained |p1| and |p2| solutions to build p1 and p2 (m2=0!)
             for (std::size_t i = 0; i < abs_p1.size(); i++) {
                 // Skip unphysical solutions
                 if (abs_p1[i] <= 0 || abs_p2[i] <= 0)
                     continue;
 
                 const double sin_theta_1 = std::sin(m_p1->Theta());
                 const double sin_theta_2 = std::sin(m_p2->Theta());
                 const double E1 = std::sqrt(SQ(abs_p1[i]) + sq_m1);
 
                 LorentzVector p1_sol, p2_sol;
                 p1_sol.SetPxPyPzE(
                         abs_p1[i] * std::cos(m_p1->Phi()) * sin_theta_1,
                         abs_p1[i] * std::sin(m_p1->Phi()) * sin_theta_1,
                         abs_p1[i] * std::cos(m_p1->Theta()),
                         E1);
                 p2_sol.SetPxPyPzE(
                         abs_p2[i] * std::cos(m_p2->Phi()) * sin_theta_2,
                         abs_p2[i] * std::sin(m_p2->Phi()) * sin_theta_2,
                         abs_p2[i] * std::cos(m_p2->Theta()),
                         abs_p2[i]);
 
                 if (!ApproxComparison(p1_sol.M() / p1_sol.E(), m1 / p1_sol.E())) {
 #ifndef NDEBUG
                     LOG(trace) << "[SecondaryBlockE] Throwing solution because of invalid mass. " <<
                         "Expected " << m1 << ", got " << p1_sol.M();
 #endif
                     continue;
                 }
 
                 if (!ApproxComparison(p2_sol.M() / p2_sol.E(), 0.)) {
 #ifndef NDEBUG
                     LOG(trace) << "[SecondaryBlockE] Throwing solution because of invalid mass. " <<
                         "Expected 0., got " << p2_sol.M();
 #endif
                     continue;
                 }
 
                 if (!ApproxComparison((p1_sol + p2_sol + *m_p3).M2(), *s123)) {
 #ifndef NDEBUE
                     LOG(trace) << "[SecondaryBlockE] Throwing solution because of invalid invariant mass. " <<
                         "Expected " << *s123 << ", got " << (p1_sol + p2_sol + *m_p3).M2();
 #endif
                     continue;
                 }
 
                 if (!ApproxComparison((p1_sol + p2_sol).M2(), *s12)) {
 #ifndef NDEBUG
                     LOG(trace) << "[SecondaryBlockE] Throwing solution because of invalid invariant mass. " <<
                         "Expected " << *s12 << ", got " << (p1_sol + p2_sol).M2();
 #endif
                     continue;
                 }
                 
                 // Compute jacobian
                 const double jacobian = abs_p2[i] * SQ(abs_p1[i]) * sin_theta_1 * sin_theta_2 /
                                             (1024 * std::pow(M_PI, 6) * std::abs(
                                                     abs_p2[i] * (abs_p1[i] - E1 * c12) * X + (E3 * abs_p1[i] - E1 * p3 * c13) * (E1 - abs_p1[i] * c12) )
                                             );
 
                 Solution solution { { p1_sol, p2_sol }, jacobian, true };
                 solutions->push_back(solution);
             }
 
             return (solutions->size() > 0) ? Status::OK : Status::NEXT;
 
         }
 
     private:
         double sqrt_s;
 
         // Inputs
         Value<double> s12;
         Value<double> s123;
         Value<LorentzVector> m_p1;
         Value<LorentzVector> m_p2;
         Value<LorentzVector> m_p3;
 
         // Output
         std::shared_ptr<SolutionCollection> solutions = produce<SolutionCollection>("solutions");
 };
 
 REGISTER_MODULE(SecondaryBlockE)
         .Input("s12")
         .Input("s123")
         .Input("p1")
         .Input("p2")
         .Input("p3")
         .Output("solutions")
         .GlobalAttr("energy: double");