/*
  *  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/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 SecondaryBlockCD: public Module {
     public:
 
         SecondaryBlockCD(PoolPtr pool, const ParameterSet& parameters): Module(pool, parameters.getModuleName()) {
                 sqrt_s = parameters.globalParameters().get<double>("energy");
 
                 s12 = get<double>(parameters.get<InputTag>("s12"));
 
                 p1 = get<LorentzVector>(parameters.get<InputTag>("p1"));
                 p2 = get<LorentzVector>(parameters.get<InputTag>("p2"));
             };
 
         virtual Status work() override {
 
             solutions->clear();
 
             const double m1 = p1->M();
             const double sq_m1 = SQ(m1);
             const double m2 = p2->M();
             const double sq_m2 = SQ(m2);
 
             // Don't spend time on unphysical part of phase-space
             if (*s12 >= SQ(sqrt_s) || sq_m1 + sq_m2 >= *s12)
                return Status::NEXT;
 
             std::vector<double> E1_solutions; // up to two solutions
             const double theta1 = p1->Theta();
             const double phi1 = p1->Phi();
 
             const double E2 = p2->E();
             const double norm2 = p2->P();
 
             const double cos_theta12 = ROOT::Math::VectorUtil::CosTheta(*p1, *p2);
 
             // Equation to be solved for E1 : s12 = (p1+p2)^2 ==> [ 4*SQ(cos_theta12)*SQ(norm2)-4*SQ(E2) ] * SQ(E1) + [ 4*(s12 - SQ(m1) - SQ(m2))*E2 ] * E1 + 4*SQ(m1)*SQ(cos_theta12)*(SQ(m2)-SQ(E2) = 0)
             const double quadraticTerm = 4 * SQ(E2) - 4 * SQ(norm2) * SQ(cos_theta12);
             const double linearTerm = 4 * E2 * (sq_m1 + sq_m2 - *s12);
             const double indepTerm = SQ(sq_m1 + sq_m2 - *s12) + 4 * sq_m1 * SQ(norm2) * SQ(cos_theta12);
 
             bool foundSolution = solveQuadratic(quadraticTerm, linearTerm, indepTerm, E1_solutions);
 
             if (!foundSolution) {
                 return Status::NEXT;
             }
 
             for (const double& E1: E1_solutions) {
                 // Skip unphysical solutions
                 if (E1 <= 0 || E1 <= m1)
                     continue;
 
                 // Avoid introducing spurious solution from the equation s12 = (p1+p2)^2 that has to be squared in the computation
                 if ((sq_m1 + sq_m2 + 2 * E1 * E2 - *s12) * cos_theta12 < 0)
                     continue;
 
                 LorentzVector gen_p1_sol;
                 double norm1 = std::sqrt(SQ(E1) - sq_m1);
                 double gen_pt1_sol = norm1 * std::sin(theta1);
                 gen_p1_sol.SetPxPyPzE(
                         gen_pt1_sol * std::cos(phi1),
                         gen_pt1_sol * std::sin(phi1),
                         norm1 * std::cos(theta1),
                         E1);
 
                 if (!ApproxComparison(gen_p1_sol.M() / gen_p1_sol.E(), m1 / gen_p1_sol.E())) {
 #ifndef NDEBUG
                     LOG(trace) << "[SecondaryBlockCD] Throwing solution because of invalid mass. " <<
                         "Expected " << m1 << ", got " << gen_p1_sol.M();
 #endif
                     continue;
                 }
 
 
                 if (!ApproxComparison((gen_p1_sol + *p2).M2(), *s12)) {
 #ifndef NDEBUG
                     LOG(trace) << "[SecondaryBlockCD] Throwing solution because of invalid invariant mass. " <<
                         "Expected " << *s12 << ", got " << (gen_p1_sol + *p2).M2();
 #endif
                     continue;
                 }
                 
                 // Compute jacobian
                 double jacobian = std::abs( std::sin(theta1) * SQ(norm1) / (32 * CB(M_PI) * (norm1 * E2 - E1 * norm2 * cos_theta12)) );
 
                 Solution s { {gen_p1_sol}, jacobian, true };
                 solutions->push_back(s);
             }
 
             return (solutions->size() > 0) ? Status::OK : Status::NEXT;
         }
 
     private:
         double sqrt_s;
 
         // Inputs
         Value<double> s12;
         Value<LorentzVector> p1;
         Value<LorentzVector> p2;
 
         // Output
         std::shared_ptr<SolutionCollection> solutions = produce<SolutionCollection>("solutions");
 };
 
 REGISTER_MODULE(SecondaryBlockCD)
         .Input("s12")
         .Input("p1")
         .Input("p2")
         .Output("solutions")
         .GlobalAttr("energy: double");