#include "TopQuarkAnalysis/TopJetCombination/plugins/TtSemiLepJetCombWMassDeltaTopMass.h"

#include "FWCore/ParameterSet/interface/ParameterSet.h"

#include "TopQuarkAnalysis/TopTools/interface/MEzCalculator.h"

TtSemiLepJetCombWMassDeltaTopMass::TtSemiLepJetCombWMassDeltaTopMass(const edm::ParameterSet& cfg):
  jetsToken_             (consumes< std::vector<pat::Jet> >(cfg.getParameter<edm::InputTag>("jets"             ))),
  lepsToken_             (consumes< edm::View<reco::RecoCandidate> >(cfg.getParameter<edm::InputTag>("leps"             ))),
  metsToken_             (consumes< std::vector<pat::MET> >(cfg.getParameter<edm::InputTag>("mets"             ))),
  maxNJets_         (cfg.getParameter<int>          ("maxNJets"         )),
  wMass_            (cfg.getParameter<double>       ("wMass"            )),
  useBTagging_      (cfg.getParameter<bool>         ("useBTagging"      )),
  bTagAlgorithm_    (cfg.getParameter<std::string>  ("bTagAlgorithm"    )),
  minBDiscBJets_    (cfg.getParameter<double>       ("minBDiscBJets"    )),
  maxBDiscLightJets_(cfg.getParameter<double>       ("maxBDiscLightJets")),
  neutrinoSolutionType_(cfg.getParameter<int>       ("neutrinoSolutionType")),
  maxNComb_         (cfg.getParameter<int>          ("maxNComb")),
  scale2Wmass_      (cfg.getParameter<bool>         ("scale2Wmass")) //,
  // jetCorrectionLevel_ (cfg.getParameter<std::string> ("jetCorrectionLevel"))
{
  if(maxNJets_<4 && maxNJets_!=-1)
    throw cms::Exception("WrongConfig")
      << "Parameter maxNJets can not be set to " << maxNJets_ << ". \n"
      << "It has to be larger than 4 or can be set to -1 to take all jets.";

  produces<std::vector<std::vector<int> > >();
  produces<int>("NumberOfConsideredJets");
  produces<std::vector<double> >("Chi2");
  produces< std::vector<pat::Particle> >("PartonsHadP");
  produces< std::vector<pat::Particle> >("PartonsHadQ");
  produces< std::vector<pat::Particle> >("PartonsHadB");
  produces< std::vector<pat::Particle> >("PartonsLepB");
  produces< std::vector<pat::Particle> >("Leptons");
  produces<std::vector<pat::Particle> >("Neutrinos");
}

TtSemiLepJetCombWMassDeltaTopMass::~TtSemiLepJetCombWMassDeltaTopMass()
{
}

void
TtSemiLepJetCombWMassDeltaTopMass::produce(edm::Event& evt, const edm::EventSetup& setup)
{
  std::unique_ptr<std::vector<std::vector<int> > > pOut(new std::vector<std::vector<int> >);
  std::unique_ptr<int> pJetsConsidered(new int);

  // std::unique_ptr< std::vector<int>               > pStatus( new std::vector<int> );
  std::unique_ptr< std::vector<pat::Particle> > pPartonsHadP( new std::vector<pat::Particle> );
  std::unique_ptr< std::vector<pat::Particle> > pPartonsHadQ( new std::vector<pat::Particle> );
  std::unique_ptr< std::vector<pat::Particle> > pPartonsHadB( new std::vector<pat::Particle> );
  std::unique_ptr< std::vector<pat::Particle> > pPartonsLepB( new std::vector<pat::Particle> );
  std::unique_ptr< std::vector<pat::Particle> > pLeptons    ( new std::vector<pat::Particle> );
  std::unique_ptr<std::vector<pat::Particle> > pNeutrinos(new std::vector<pat::Particle>);

  std::vector<int> match;
  std::vector<int> match2;  //second solution
  for (unsigned int i = 0; i < 4; ++i){
    match.push_back(-1);
    match2.push_back(-1);
  }

  // get jets
  edm::Handle< std::vector<pat::Jet> > jets;
  evt.getByToken(jetsToken_, jets);

  // get leptons
  edm::Handle< edm::View<reco::RecoCandidate> > leps;
  evt.getByToken(lepsToken_, leps);

  // get MET
  edm::Handle< std::vector<pat::MET> > mets;
  evt.getByToken(metsToken_, mets);

  // skip events without lepton candidate or less than 4 jets or no MET
  if(leps->empty() || jets->size() < 4 || mets->empty()){
    pOut->push_back( match );
    evt.put(std::move(pOut));
    *pJetsConsidered = jets->size();
    evt.put(std::move(pJetsConsidered), "NumberOfConsideredJets");
    
    pPartonsHadP->push_back( pat::Particle() );
    pPartonsHadQ->push_back( pat::Particle() );
    pPartonsHadB->push_back( pat::Particle() );
    pPartonsLepB->push_back( pat::Particle() );
    pLeptons    ->push_back( pat::Particle() );
    pNeutrinos  ->push_back( pat::Particle() );

    evt.put(std::move(pPartonsHadP), "PartonsHadP");
    evt.put(std::move(pPartonsHadQ), "PartonsHadQ");
    evt.put(std::move(pPartonsHadB), "PartonsHadB");
    evt.put(std::move(pPartonsLepB), "PartonsLepB");
    evt.put(std::move(pLeptons    ), "Leptons"    );
    evt.put(std::move(pNeutrinos  ), "Neutrinos"  );;
    return;
  }

  unsigned maxNJets = maxNJets_;
  if(maxNJets_ == -1 || (int)jets->size() < maxNJets_) maxNJets = jets->size();
  *pJetsConsidered = maxNJets;
  evt.put(std::move(pJetsConsidered), "NumberOfConsideredJets");

  std::vector<bool> isBJet;
  std::vector<bool> isLJet;
  int cntBJets = 0;
  if(useBTagging_) {
    for(unsigned int idx=0; idx<maxNJets; ++idx) {
      isBJet.push_back( ((*jets)[idx].bDiscriminator(bTagAlgorithm_) > minBDiscBJets_    ) );
      isLJet.push_back( ((*jets)[idx].bDiscriminator(bTagAlgorithm_) < maxBDiscLightJets_) );
      if((*jets)[idx].bDiscriminator(bTagAlgorithm_) > minBDiscBJets_    )cntBJets++;
    }
  }

  // -----------------------------------------------------
  // if saving two solutions, the second one only makes sence
  // if there are exactly two b jets, and 4 jets.
  // -----------------------------------------------------
  if(maxNComb_ == 2 && ( maxNJets !=4 || cntBJets != 2 ) ){
    pOut->push_back( match );
    evt.put(std::move(pOut));
    // the getters return empty objects here
    pPartonsHadP->push_back( pat::Particle() );
    pPartonsHadQ->push_back( pat::Particle() );
    pPartonsHadB->push_back( pat::Particle() );
    pPartonsLepB->push_back( pat::Particle() );
    pLeptons    ->push_back( pat::Particle() );
    pNeutrinos  ->push_back( pat::Particle() );

    evt.put(std::move(pPartonsHadP), "PartonsHadP");
    evt.put(std::move(pPartonsHadQ), "PartonsHadQ");
    evt.put(std::move(pPartonsHadB), "PartonsHadB");
    evt.put(std::move(pPartonsLepB), "PartonsLepB");
    evt.put(std::move(pLeptons    ), "Leptons"    );
    evt.put(std::move(pNeutrinos  ), "Neutrinos"  );
    return;
  }


  // -----------------------------------------------------
  // associate those jets that get closest to the W mass
  // with their invariant mass to the hadronic W boson
  // -----------------------------------------------------
  double wDist =-1.;
  std::vector<int> closestToWMassIndices;
  closestToWMassIndices.push_back(-1);
  closestToWMassIndices.push_back(-1);
  for(unsigned idx=0; idx<maxNJets; ++idx){
    if(useBTagging_ && (!isLJet[idx] || (cntBJets<=2 && isBJet[idx]))) continue;
    for(unsigned jdx=(idx+1); jdx<maxNJets; ++jdx){
      if(useBTagging_ && (!isLJet[jdx] || (cntBJets<=2 && isBJet[jdx]) || (cntBJets==3 && isBJet[idx] && isBJet[jdx]))) continue;
      reco::Particle::LorentzVector sum =
	(*jets)[idx].p4()+
	(*jets)[jdx].p4();
      if( wDist<0. || wDist>fabs(sum.mass()-wMass_) ){
	wDist=fabs(sum.mass()-wMass_);
	closestToWMassIndices.clear();
	closestToWMassIndices.push_back(idx);
	closestToWMassIndices.push_back(jdx);
      }
    }
  }

  // -----------------------------------------------------
  // build a leptonic W boson from the lepton and the MET
  // -----------------------------------------------------
  double neutrino_pz = 0.;
  if(neutrinoSolutionType_!=-1) {
    MEzCalculator mez;
    mez.SetMET( *mets->begin() );
    if( dynamic_cast<const reco::Muon*>(&(leps->front())) )
      mez.SetLepton( (*leps)[0], true );
    else if( dynamic_cast<const reco::GsfElectron*>(&(leps->front())) )
      mez.SetLepton( (*leps)[0], false );
    else
      throw cms::Exception("UnimplementedFeature") << "Type of lepton given together with MET for solving neutrino kinematics is neither muon nor electron.\n";
    neutrino_pz = mez.Calculate(neutrinoSolutionType_);
  }
  const math::XYZTLorentzVector neutrino( mets->begin()->px(), mets->begin()->py(), neutrino_pz, sqrt(mets->begin()->px()*mets->begin()->px()
												      + mets->begin()->py()*mets->begin()->py()
												      + neutrino_pz*neutrino_pz) );
  const reco::Particle::LorentzVector lepW = neutrino + leps->front().p4();

  // -----------------------------------------------------
  // associate those jets to the hadronic and the leptonic
  // b quark that minimize the difference between
  // hadronic and leptonic top mass
  // -----------------------------------------------------
  double deltaTop=-1.;
  int hadB=-1;
  int lepB=-1;
  //?? Do we need to set these to empty particles if isValid(closestToWMassIndices[0], jets)=false ?
  reco::Particle::LorentzVector light_q = (*jets)[closestToWMassIndices[0]].p4();
  reco::Particle::LorentzVector light_qbar = (*jets)[closestToWMassIndices[1]].p4();
  const reco::Particle::LorentzVector hadW = light_q + light_qbar;

  if( isValid(closestToWMassIndices[0], jets) && isValid(closestToWMassIndices[1], jets)) {
    // find hadronic b candidate
    for(unsigned idx=0; idx<maxNJets; ++idx){
      if(useBTagging_ && !isBJet[idx]) continue;
      // make sure it's not used up already from the hadronic W
      if( (int)idx!=closestToWMassIndices[0] && (int)idx!=closestToWMassIndices[1] ){
	reco::Particle::LorentzVector hadTop = hadW + (*jets)[idx].p4();
	// find leptonic b candidate
	for(unsigned jdx=0; jdx<maxNJets; ++jdx){
	  if(useBTagging_ && !isBJet[jdx]) continue;
	  // make sure it's not used up already from the hadronic branch
	  if( (int)jdx!=closestToWMassIndices[0] && (int)jdx!=closestToWMassIndices[1] && jdx!=idx ){
	    reco::Particle::LorentzVector lepTop = lepW + (*jets)[jdx].p4();
	    if( deltaTop<0. || deltaTop>fabs(hadTop.mass()-lepTop.mass()) ){
	      deltaTop=fabs(hadTop.mass()-lepTop.mass());
	      hadB=idx;
	      lepB=jdx;
	    }
	  }
	}
      }
    }
  }

  if( scale2Wmass_) {
    if(( hadW.M() < 70 ) || ( 100 < hadW.M() )){
      hadB=-1;
      lepB=-1;
    }
      float W_mismatch = (float)wMass_/(float)hadW.M();
    // std::cout << "had W = " << hadW.M() << ", W_mismatch = " << W_mismatch << std::endl;
    light_q = light_q*W_mismatch;
    light_qbar = light_qbar*W_mismatch;
  }


  // -----------------------------------------------------
  // save the solution with the smallest top mass difference
  // -----------------------------------------------------
  match[TtSemiLepEvtPartons::LightQ   ] = closestToWMassIndices[0];
  match[TtSemiLepEvtPartons::LightQBar] = closestToWMassIndices[1];
  match[TtSemiLepEvtPartons::HadB     ] = hadB;
  match[TtSemiLepEvtPartons::LepB     ] = lepB;
  pOut->push_back( match );


  pPartonsHadP->push_back(pat::Particle(reco::LeafCandidate(0, light_q, math::XYZPoint())));
  pPartonsHadQ->push_back(pat::Particle(reco::LeafCandidate(0, light_qbar, math::XYZPoint())));
  pPartonsHadB->push_back(pat::Particle(reco::LeafCandidate(0, (*jets)[hadB].p4(), math::XYZPoint())));
  pPartonsLepB->push_back(pat::Particle(reco::LeafCandidate(0, (*jets)[lepB].p4(), math::XYZPoint())));
  pLeptons->push_back(pat::Particle(reco::LeafCandidate(0, leps->front().p4(), math::XYZPoint())));
  pNeutrinos->push_back(pat::Particle(reco::LeafCandidate(0, neutrino, math::XYZPoint())));
  
  // -----------------------------------------------------
  // save the solution with the largest top mass difference
  // -----------------------------------------------------
  if( maxNComb_ == 2 ){
    match2[TtSemiLepEvtPartons::LightQ   ] = closestToWMassIndices[0];
    match2[TtSemiLepEvtPartons::LightQBar] = closestToWMassIndices[1];
    match2[TtSemiLepEvtPartons::HadB     ] = lepB;
    match2[TtSemiLepEvtPartons::LepB     ] = hadB;
    pOut->push_back( match2 );
    
    pPartonsHadP->push_back(pat::Particle(reco::LeafCandidate(0, light_q, math::XYZPoint())));
    pPartonsHadQ->push_back(pat::Particle(reco::LeafCandidate(0, light_qbar, math::XYZPoint())));
    pPartonsHadB->push_back(pat::Particle(reco::LeafCandidate(0, (*jets)[lepB].p4(), math::XYZPoint())));
    pPartonsLepB->push_back(pat::Particle(reco::LeafCandidate(0, (*jets)[hadB].p4(), math::XYZPoint())));
    pLeptons->push_back(pat::Particle(reco::LeafCandidate(0, leps->front().p4(), math::XYZPoint())));
    pNeutrinos->push_back(pat::Particle(reco::LeafCandidate(0, neutrino, math::XYZPoint())));
  

    // Trying to make chi2 to be saved into an event tree
    std::unique_ptr< std::vector<double>            > pChi2  ( new std::vector<double> );
    pChi2->push_back( 1. );
    pChi2->push_back( 2. );
    evt.put(std::move(pChi2       ), "Chi2"       );
    // std::cout << "Neutrino pt = " << pNeutrinos->at(1).pt() << std::endl;
    // std::cout << "Neutrino eta = " << pNeutrinos->at(1).eta() << std::endl;
    // std::cout << "Neutrino phi = " << pNeutrinos->at(1).phi() << std::endl;
  }

  // std::cout << "Neutrino pt =" << pNeutrinos->at(0).pt() << std::endl;                                                       
  // std::cout << "Neutrino eta =" << pNeutrinos->at(0).eta() << std::endl;                                                   
  // std::cout << "Neutrino phi =" << pNeutrinos->at(0).phi() << std::endl;

  evt.put(std::move(pOut));
  evt.put(std::move(pPartonsHadP), "PartonsHadP");
  evt.put(std::move(pPartonsHadQ), "PartonsHadQ");
  evt.put(std::move(pPartonsHadB), "PartonsHadB");
  evt.put(std::move(pPartonsLepB), "PartonsLepB");
  evt.put(std::move(pLeptons    ), "Leptons"    );
  evt.put(std::move(pNeutrinos  ), "Neutrinos"  );

                                                   
  // std::cout << "Adding my verbosity = "  << std::endl;
  // evt.isHypoValid("kWMassDeltaTopMass",1)
  // evt.print(1)

}
