/*
  *  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/Types.h>
 #include <momemta/Math.h>
 
 
 class BlockB: public Module {
     public:
 
         BlockB(PoolPtr pool, const ParameterSet& parameters): Module(pool, parameters.getModuleName()) {
 
             sqrt_s = parameters.globalParameters().get<double>("energy");
             pT_is_met = parameters.get<bool>("pT_is_met", false);
 
             m1 = parameters.get<double>("m1", 0.);
 
             s12 = get<double>(parameters.get<InputTag>("s12"));
 
             p2 = get<LorentzVector>(parameters.get<InputTag>("p2"));
 
             if (parameters.exists("branches")) {
                 auto branches_tags = parameters.get<std::vector<InputTag>>("branches");
                 for (auto& t: branches_tags)
                     m_branches.push_back(get<LorentzVector>(t));
             }
 
             // If the met input is specified, get it, otherwise retrieve default
             // one ("met::p4")
             InputTag met_tag;
             if (parameters.exists("met")) {
                 met_tag = parameters.get<InputTag>("met");
             } else {
                 met_tag = InputTag({"met", "p4"});
             }
 
             m_met = get<LorentzVector>(met_tag);
         };
 
         virtual Status work() override {
 
             solutions->clear();
 
             // Equations to solve:
             //(1) (p1 + p2)^2 = s12 = M1^2 + M2^2 + 2 E1 E2 - 2 p1x p2x - 2 p1y p2y - 2 p1z p2z
             //(2)  p1x = - pTx  #Coming from pT neutrino = -pT visible = - (p2 + ISR)
             //(3)  p1y = - pTy  #Coming from pT neutrino = -pT visible = - (p2 + ISR)
             //(4)  E1^2 - p1x^2 - p1y^2 - p1z^2 = M1^2
 
             const double p11 = SQ(m1);
             const double p22 = p2->M2();
 
             // Don't spend time on unphysical part of phase-space
             if (*s12 >= SQ(sqrt_s) || *s12 <= p11 + p22)
                 return Status::NEXT;
 
             LorentzVector pT;
             if (pT_is_met) {
                 pT = - *m_met;
             } else {
                 pT = *p2;
                 for (size_t i = 0; i < m_branches.size(); i++) {
                     pT += *m_branches[i];
                 }
             }
 
             // From eq.(1) p1z = B*E1 + A
             // From eq.(4) + eq.(1) (1 - B^2) E1^2 - 2 A B E1 + C - A^2 - M1^2 = 0
 
             const double A = - (*s12 - p11 - p22 - 2 * (pT.Px() * p2->Px() + pT.Py() * p2->Py())) / (2 * p2->Pz());
             const double B = p2->E() / p2->Pz();
             const double C = - SQ(pT.Px()) - SQ(pT.Py());
 
             // Solve quadratic a*E1^2 + b*E1 + c = 0
             const double a = 1 - SQ(B);
             const double b = - 2 * A * B;
             const double c = C - SQ(A) - p11;
 
             std::vector<double> E1;
 
             solveQuadratic(a, b, c, E1, false);
 
             if (E1.size() == 0)
                 return Status::NEXT;
 
             for (unsigned int i = 0; i < E1.size(); i++){
                 const double e1 = E1.at(i);
 
                 if (e1 <= 0) continue;
 
                 LorentzVector p1(-pT.Px(), -pT.Py(), A + B*e1, e1);
 
                 // Check if solutions are physical
                 LorentzVector tot = p1 + *p2;
                 for (size_t i = 0; i < m_branches.size(); i++)
                     tot += *m_branches[i];
                 double q1Pz = std::abs(tot.Pz() + tot.E()) / 2.;
                 double q2Pz = std::abs(tot.Pz() - tot.E()) / 2.;
                 if (q1Pz > sqrt_s / 2 || q2Pz > sqrt_s / 2)
                     continue;
 
                 if (!ApproxComparison(p1.M() / p1.E(), m1 / p1.E())) {
 #ifndef NDEBUG
                     LOG(trace) << "[BlockB] Throwing solution because of invalid mass. " <<
                                "Expected " << m1 << ", got " << p1.M();
 #endif
                     continue;
                 }
 
                 if (!ApproxComparison((p1 + *p2).M2(), *s12)) {
 #ifndef NDEBUG
                     LOG(trace) << "[BlockB] Throwing solution because of invalid invariant mass. " <<
                                "Expected " << *s12 << ", got " << (p1 + *p2).M2();
 #endif
                     continue;
                 }
 
                 const double inv_jacobian = SQ(sqrt_s) * std::abs(p2->Pz() * e1 - p2->E() * p1.Pz());
 
                 Solution s { {p1}, M_PI / inv_jacobian, true };
                 solutions->push_back(s);
             }
 
             return solutions->size() > 0 ? Status::OK : Status::NEXT;
         }
 
     private:
         double sqrt_s;
         bool pT_is_met;
         double m1;
 
         // Inputs
         Value<double> s12;
         Value<LorentzVector> p2;
         std::vector<Value<LorentzVector>> m_branches;
         Value<LorentzVector> m_met;
 
         // Outputs
         std::shared_ptr<SolutionCollection> solutions = produce<SolutionCollection>("solutions");
 };
 
 REGISTER_MODULE(BlockB)
         .Input("s12")
         .Input("p2")
         .OptionalInputs("branches")
         .Input("met=met::p4")
         .Output("solutions")
         .GlobalAttr("energy:double")
         .Attr("pT_is_met:bool=false")
         .Attr("m1:double=0.");