// -*- C++ -*-
//
// Package:    DimuonVtxProducer
// Class:      DimuonVtxProducer
// 
/**\class DimuonVtxProducer DimuonVtxProducer.cc MuonAnalysis/DimuonVtxProducer/src/DimuonVtxProducer.cc

 Description: [one line class summary]

 Implementation:
     [Notes on implementation]
*/
//
// Original Author:  Cesar Bernardes,40 2-B15,+41227671625,
//         Created:  Fri Oct  4 19:44:56 CEST 2013
// $Id$
//
//


// system include files
#include <memory>

// user include files
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Framework/interface/EDProducer.h"

#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/MakerMacros.h"

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


//my includes
#include "FWCore/Utilities/interface/InputTag.h"
#include "DataFormats/Common/interface/ValueMap.h"
#include "DataFormats/Common/interface/View.h"
#include "DataFormats/Candidate/interface/CompositeCandidate.h"
#include "DataFormats/MuonReco/interface/Muon.h"

//for tracks
#include "DataFormats/TrackReco/interface/Track.h"


// for vertexing
#include "FWCore/Framework/interface/ESHandle.h"
#include "TrackingTools/TransientTrack/interface/TransientTrack.h"
#include "TrackingTools/TransientTrack/interface/TransientTrackBuilder.h"
#include "TrackingTools/Records/interface/TransientTrackRecord.h"
#include "RecoVertex/KalmanVertexFit/interface/KalmanVertexFitter.h"


#include <TMath.h>

//
// class declaration
//

class DimuonVtxProducer : public edm::EDProducer {
   public:
      explicit DimuonVtxProducer(const edm::ParameterSet&);
      ~DimuonVtxProducer();

      static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
      reco::TrackRef mu_track(const reco::Muon & mu);
      //TransientVertex dimuon_vtxfit(const reco::Muon & mu1, const reco::Muon & mu2, edm::ESHandle<TransientTrackBuilder> trkBuild);

   private:
      virtual void beginJob() ;
      virtual void produce(edm::Event&, const edm::EventSetup&);
      virtual void endJob() ;
      
      virtual void beginRun(edm::Run&, edm::EventSetup const&);
      virtual void endRun(edm::Run&, edm::EventSetup const&);
      virtual void beginLuminosityBlock(edm::LuminosityBlock&, edm::EventSetup const&);
      virtual void endLuminosityBlock(edm::LuminosityBlock&, edm::EventSetup const&);

      // ----------member data ---------------------------

      edm::InputTag src_;

      /// Write a ValueMap<float> in the event
      void writeValueMap(edm::Event &iEvent,
              const edm::Handle<edm::View<reco::Candidate> > & handle,
              const std::vector<float> & values,
              const std::string    & label) const ;


};

//
// constants, enums and typedefs
//


//
// static data member definitions
//

//
// constructors and destructor
//
DimuonVtxProducer::DimuonVtxProducer(const edm::ParameterSet& iConfig) :
   src_(iConfig.getParameter<edm::InputTag>("src"))
{
   //register your products
/* Examples
   produces<ExampleData2>();

   //if do put with a label
   produces<ExampleData2>("label");
 
   //if you want to put into the Run
   produces<ExampleData2,InRun>();
*/
   //now do what ever other initialization is needed
  
   produces<edm::ValueMap<float> >("vtxNormQui2");
   produces<edm::ValueMap<float> >("vtxZcoordinate");
   produces<edm::ValueMap<float> >("vtxRdistance");
   produces<edm::ValueMap<float> >("tagPtAtTheVtx");
   produces<edm::ValueMap<float> >("probePtAtTheVtx");
   produces<edm::ValueMap<float> >("tagPtBefore");
   produces<edm::ValueMap<float> >("probePtBefore");

}


DimuonVtxProducer::~DimuonVtxProducer()
{
 
   // do anything here that needs to be done at desctruction time
   // (e.g. close files, deallocate resources etc.)

}


//
// member functions
//

// ------------ method called to produce the data  ------------
void
DimuonVtxProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup)
{
   using namespace edm;

/* This is an event example
   //Read 'ExampleData' from the Event
   Handle<ExampleData> pIn;
   iEvent.getByLabel("example",pIn);

   //Use the ExampleData to create an ExampleData2 which 
   // is put into the Event
   std::auto_ptr<ExampleData2> pOut(new ExampleData2(*pIn));
   iEvent.put(pOut);
*/

/* this is an EventSetup example
   //Read SetupData from the SetupRecord in the EventSetup
   ESHandle<SetupData> pSetup;
   iSetup.get<SetupRecord>().get(pSetup);
*/

   Handle<View<reco::Candidate> > src;
   iEvent.getByLabel(src_, src);
  
   size_t n = src->size();
   std::vector<float> vtxNormQui2(n,-999), vtxZcoordinate(n,-999), vtxRdistance(n,-999), tagPtAtTheVtx(n,0), probePtAtTheVtx(n,0),                                           tagPtBefore(n,0), probePtBefore(n,0);
   
   edm::ESHandle<TransientTrackBuilder> trkBuild;
   iSetup.get<TransientTrackRecord>().get("TransientTrackBuilder",trkBuild);

   for(size_t i = 0; i < n; ++i) {

      const reco::Candidate & ci = (*src)[i];
      if (ci.numberOfDaughters() != 2) throw cms::Exception("CorruptData") <<
         "DimuonVtxProducer should be used on composite candidates with two daughters, this one has " << ci.numberOfDaughters() << "\n";
      const reco::Candidate &d1 = *ci.daughter(0), &d2 = *ci.daughter(1);
      
      const reco::Muon *mu1 = dynamic_cast<const reco::Muon *>(&*d1.masterClone());
      if (mu1 == 0) throw cms::Exception("CorruptData") << "First daughter of candidate is not a ShallowClone of a reco::Muon\n";

      const reco::Muon *mu2 = dynamic_cast<const reco::Muon *>(&*d2.masterClone());
      if (mu2 == 0) throw cms::Exception("CorruptData") << "Second daughter of candidate is not a ShallowClone of a reco::Muon\n"; 

      if( mu1->globalTrack().isNull() || mu2->globalTrack().isNull() ) continue;

      //TransientVertex dimu_vtx = dimuon_vtxfit(*mu1, *mu2, trkBuild);

      reco::TrackRef muon1 = mu_track(*mu1);
      reco::TrackRef muon2 = mu_track(*mu2);
      std::vector<reco::TransientTrack> t_tks;
      t_tks.clear();
      t_tks.push_back( (*trkBuild).build(muon1.get()));
      t_tks.push_back( (*trkBuild).build(muon2.get()));     
      TransientVertex dimu_vtx;
      if(t_tks.size() > 1){
         KalmanVertexFitter kvf(true);
         dimu_vtx = kvf.vertex(t_tks);
      } 

      
      if(dimu_vtx.isValid() && dimu_vtx.totalChiSquared() >= 0. && dimu_vtx.degreesOfFreedom() > 0){

        //std::cout << "Normalized Qui2 do Vertice: \t" << dimu_vtx.totalChiSquared()/dimu_vtx.degreesOfFreedom() << std::endl;
        vtxNormQui2[i] = dimu_vtx.totalChiSquared()/dimu_vtx.degreesOfFreedom();

        //std::cout << "Normalized Qui2 do Vertice (z coordinate): \t" << dimu_vtx.position().z() << std::endl;
        vtxZcoordinate[i] = dimu_vtx.position().z();

        //std::cout << "Normalized Qui2 do Vertice (r distance): \t" << dimu_vtx.position().perp() << std::endl;
        vtxRdistance[i] = dimu_vtx.position().perp();  

        
        // important! evaluate momentum vectors AT THE VERTEX
        TrajectoryStateClosestToPoint tag_TSCP = t_tks[0].trajectoryStateClosestToPoint(dimu_vtx.position());
        TrajectoryStateClosestToPoint probe_TSCP = t_tks[1].trajectoryStateClosestToPoint(dimu_vtx.position());
        GlobalVector tag_momentum = tag_TSCP.momentum();
        GlobalVector probe_momentum = probe_TSCP.momentum();

        //Tag pt
        double tag_px = tag_momentum.x();
        double tag_py = tag_momentum.y();
        //std::cout << "tag px:  " << tag_px << "  tag py:  "  << tag_py << std::endl;
        tagPtAtTheVtx[i] = TMath::Sqrt(TMath::Power(tag_px,2) + TMath::Power(tag_py,2));
        tagPtBefore[i] = mu1->pt();

        //Probe pt
        double probe_px = probe_momentum.x();
        double probe_py = probe_momentum.y();
        //std::cout << "probe px:  " << probe_px << "  probe py:  "  << probe_py << std::endl;        
        probePtAtTheVtx[i] = TMath::Sqrt(TMath::Power(probe_px,2) + TMath::Power(probe_py,2));
        probePtBefore[i] = mu2->pt();

      }else{
        //std::cout << "Qui2 do Vertice ----- VERTEX IS NOT VALID!!!" << std::endl;
      }     
      
   }
   
   writeValueMap(iEvent, src, vtxNormQui2, "vtxNormQui2");
   writeValueMap(iEvent, src, vtxZcoordinate, "vtxZcoordinate");
   writeValueMap(iEvent, src, vtxRdistance, "vtxRdistance");
   writeValueMap(iEvent, src, tagPtAtTheVtx, "tagPtAtTheVtx");
   writeValueMap(iEvent, src, probePtAtTheVtx, "probePtAtTheVtx");
   writeValueMap(iEvent, src, tagPtBefore, "tagPtBefore");
   writeValueMap(iEvent, src, probePtBefore, "probePtBefore");
 
}

// ------------ method called once each job just before starting event loop  ------------
void 
DimuonVtxProducer::beginJob()
{
}

// ------------ method called once each job just after ending the event loop  ------------
void 
DimuonVtxProducer::endJob() {
}

// ------------ method called when starting to processes a run  ------------
void 
DimuonVtxProducer::beginRun(edm::Run&, edm::EventSetup const&)
{
}

// ------------ method called when ending the processing of a run  ------------
void 
DimuonVtxProducer::endRun(edm::Run&, edm::EventSetup const&)
{
}

// ------------ method called when starting to processes a luminosity block  ------------
void 
DimuonVtxProducer::beginLuminosityBlock(edm::LuminosityBlock&, edm::EventSetup const&)
{
}

// ------------ method called when ending the processing of a luminosity block  ------------
void 
DimuonVtxProducer::endLuminosityBlock(edm::LuminosityBlock&, edm::EventSetup const&)
{
}

// ------------ method fills 'descriptions' with the allowed parameters for the module  ------------
void
DimuonVtxProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
  //The following says we do not know what parameters are allowed so do no validation
  // Please change this to state exactly what you do use, even if it is no parameters
  edm::ParameterSetDescription desc;
  desc.setUnknown();
  descriptions.addDefault(desc);
}



void
DimuonVtxProducer::writeValueMap(edm::Event &iEvent,
        const edm::Handle<edm::View<reco::Candidate> > & handle,
        const std::vector<float> & values,
        const std::string & label) const
{
    using namespace edm;
    using namespace std;
    auto_ptr<ValueMap<float> > valMap(new ValueMap<float>());
    edm::ValueMap<float>::Filler filler(*valMap);
    filler.insert(handle, values.begin(), values.end());
    filler.fill();
    iEvent.put(valMap, label);
}


/*
TransientVertex DimuonVtxProducer::dimuon_vtxfit(const reco::Muon & mu1, const reco::Muon & mu2, edm::ESHandle<TransientTrackBuilder> trkBuild){

    reco::TrackRef muon1 = mu_track(mu1);

    reco::TrackRef muon2 = mu_track(mu2);


    std::vector<reco::TransientTrack> t_tks;

    t_tks.clear();    

    t_tks.push_back( (*trkBuild).build(muon1.get()));
 
    t_tks.push_back( (*trkBuild).build(muon2.get()));

    TransientVertex tv;

    if (t_tks.size() > 1) {

      KalmanVertexFitter kvf(true);

      tv = kvf.vertex(t_tks);

    }

   return tv;
}
*/

reco::TrackRef DimuonVtxProducer::mu_track(const reco::Muon & mu){

    reco::TrackRef to_return;

    if(mu.isStandAloneMuon()) to_return = mu.outerTrack();

    if(mu.isTrackerMuon()) to_return = mu.innerTrack();

    if(mu.isGlobalMuon()) to_return = mu.globalTrack();

    return to_return;

}


//define this as a plug-in
DEFINE_FWK_MODULE(DimuonVtxProducer);