#include "caf_eff_utils/JetCorr.hpp"
#include "caf_eff_utils/JetEffTool.hpp"
#include "eff_utils/EffVal.hpp"
#include "eff_utils/BinnedEfficiency.hpp"
#include "cafe/Stat.hpp"
#include "cafe/Config.hpp"
#include "tmb_tree/TMBJet.hpp"
#include <stdexcept>
#include "caf_mc_util/MCReqID.hpp"
#include "caf_dq/CafeDataQualityProcessor.hpp"
#include "caf_dq/CafeDataQualityLuminosity.hpp"

using namespace cafe ;
using namespace std ;
using namespace eff_utils;
using namespace DataQualityChecker ;

extern caf_mc_util::MCReqID *MCREQID ;

namespace caf_util {

  // helper function based on one from RunJetSSR
  // ______________________________________________

  void SetJESType(int origActAs, const TMBJet& jet) { // yes, ActAs is a mutable. SNAFU.
    if (origActAs == TMBJet::kUnCorrected) {
      jet.ActAsUnCorrected();
    } else if (origActAs == TMBJet::kJESMUCorrected) {
      jet.ActAsJESMUCorrected();
    } else if (origActAs == TMBJet::kJESMUCorrectedShiftedPlus) {
      jet.ActAsJESMUCorrectedShiftedPlus();
    } else if (origActAs == TMBJet::kJESMUCorrectedShiftedMinus) {
      jet.ActAsJESMUCorrectedShiftedMinus();
    } else if (origActAs == TMBJet::kJESCorrected) {
      jet.ActAsJESCorrected();
    } else if (origActAs == TMBJet::kJESCorrectedShiftedPlus) {
      jet.ActAsJESCorrectedShiftedPlus();
    } else if (origActAs == TMBJet::kJESCorrectedShiftedMinus) {
      jet.ActAsJESCorrectedShiftedMinus();
    } else if (origActAs == TMBJet::kSmeared) {
      jet.ActAsSmeared();
    } else if (origActAs == TMBJet::kSmearedMU) {
      jet.ActAsSmearedMU();
    } else {
      throw std::runtime_error( Form ("JetCorr: Can't reset jet's ActAs to unknown JES type (%d)", origActAs));
    }
    return ;
  }

  // Now the class itself
  // ____________________

  JetCorr::JetCorr(const char *name) : cafe::Processor(name),
				       _eff(0), _weight(0), _rms(0),
				       _events(0),_reqid(0)
  {
    cafe::Config config(name);
    
    _fromBranch = config.get("From", "");
    if (_fromBranch.size() == 0) {
      err() << "caf_eff_utils::JetCorr named \"" << this->name()
	    <<"\": You must specify an input branch name!"
	    << endl ;
      exit(1) ;
    }

    _vtxBranch = config.get("VtxBranch", "PrimaryVertex");

    // This switch is for using with JetSSR so the JetID weight
    // will not get propagated to the global stat pointer
    _applyToGlobalWeight = config.get("ApplyToGlobalWeight",false);

    // EffInfo info ;

    // efficiency file location (the initialization does not matter)
    path = config.get("Path","./");
    
    // EffInfo specifications
    string type = config.get("EffType","");
    if(type.size() > 0) info.EffType(type) ;

    string effname = config.get("EffName","");
    if(effname.size() > 0) info.EffName(effname) ;

    string effmethod = config.get("EffMethod","");
    if(effmethod.size() > 0) info.EffMethod(effmethod) ;
    
    //not needed
    //triggerversion = config.get("TriggerVersion","");
    //if(triggerversion.size() > 0) info.TriggerVersion(triggerversion) ;
    
    vector< string > names = config.getVString("EffVarNames",", ");
    if(names.size() > 0) info.EffVarNames(names) ;
    for (vector< string >::const_iterator it = names.begin() ; it != names.end(); it++) {
      // Jet
      if ( *it == "deta" || *it == "deteta" || *it == "DetEta" )
	_vars.add("_deta") ;
      else if ( *it == "dphi" || *it == "detphi"  || *it == "DetPhi" )
	_vars.add("_dphi") ;
      //else if ( *it == "e" || *it == "E" )   _vars.add("fM") ;
      else if ( *it == "eta" || *it == "Eta" || *it == "phi" || *it == "Phi"
		|| *it == "pt" || *it == "Pt" || *it == "pT" 
		|| *it == "e" || *it == "E") {
	_vars.add("fX") ;
	_vars.add("fY") ;
	_vars.add("fZ") ;
	_vars.add("_actas");
	_vars.add("_jes_C");
	if ( *it == "e" || *it == "E" )   _vars.add("fM") ;
      }
      //else if ( *it == "edet" || *it == "Edet" )   _vars.add("_uncorr.fM") ;
      else if (	*it == "ptdet" || *it == "Ptdet" || *it == "pTdet"
		|| *it == "edet" || *it == "Edet") {
	_vars.add("fX") ;
	_vars.add("fY") ;
	_vars.add("fZ") ;
	_vars.add("_actas") ;
	_vars.add("_uncorr.fX") ;
	_vars.add("_uncorr.fY") ;
	_vars.add("_uncorr.fZ") ;
	if ( *it == "edet" || *it == "Edet" )   {
	  _vars.add("_uncorr.fM") ;
	  _vars.add("fM") ;
	}
      }
      // Primary vertex
      else if ( *it == "z" || *it == "Z" || *it == "zvtx" || *it == "Zvtx"
	 || *it == "vertexZ" || *it == "VertexZ" || *it == "pvzsign" || *it == "PVZsign"
	 || *it == "etasign" || *it == "ETAsign")
	_varsVtx.add("_VertexZ");
      //Instantaneous Luminosity
      else if ( *it == "instlumi" )
        _varsMC.add("_overlay_instlum");
      // Or else
      else {
	std::cout <<  "JetCorr Constructor ERROR: Axis variable \"" << (*it) 
		  <<"\" not recognised!" << std::endl ;
	exit(1) ;
      } 
    }
    // For all jets
    _vars.add("_actas");

    int version = config.get("EffVersion",0);    
    if(version > 0) info.EffVersion(version) ;

    string objtype = config.get("ObjType","");
    if(objtype.size() > 0) info.ObjType(objtype) ;

    string quality = config.get("ObjQuality","");
    if(quality.size() > 0) info.ObjQuality(quality) ;

    string relativeto = config.get("ObjRelativeTo","");
    if(relativeto.size() > 0) info.ObjRelativeTo(relativeto) ;

    int objversion = config.get("ObjVersion",0);
    if(objversion > 0) info.ObjVersion(objversion) ;

    long lo = config.get ("RunRangeLow",0);
    if(lo > 0) info.RunRangeLow(lo) ;

    long hi = config.get("RunRangeHigh",0);
    if(hi > 0) info.RunRangeHigh(hi) ;

    try {
      _eff = new JetEffTool(info, path) ;
    } catch (eff_utils::EffException& ex) {
      err() << "caf_eff_utils::JetCorr named \"" << this->name()
	    <<"\": caught exception " << ex._message
	    << endl ;
      err() << info << endl ;
      exit(1) ;
    }
    
    if (!_eff->isValid()) {
      err() << "caf_eff_utils::JetCorr named \"" << this->name()
	    <<"\": Could not find efficiency object with following specifications" << endl
	    << info << endl ;
      err() << info << endl ;
      exit(1) ;
    }

    bool ignoreOverflow = config.get("ignoreOverflow",false);
    bool ignoreUnderflow = config.get("ignoreUnderflow",false);
    if (ignoreOverflow || ignoreUnderflow) {
      const BinnedEfficiency* beff = dynamic_cast<const BinnedEfficiency*> (&_eff->EffObj()) ;
      if (!beff) {
        err() << "caf_eff_utils::JetCorr named \"" << this->name()
              <<"\": Could not convert efficiency to the binned efficiency" << endl ;
        exit(1) ;
      }
      if (ignoreOverflow && ignoreUnderflow)
        beff->setExeptionMode(BinnedEfficiency::IGNORE_UNDEROVERFLOW) ;
      else if (ignoreOverflow ) 
        beff->setExeptionMode(BinnedEfficiency::IGNORE_OVERFLOW) ;
      else if (ignoreUnderflow)
        beff->setExeptionMode(BinnedEfficiency::IGNORE_UNDERFLOW) ;
    }
    
    out() << " >>>>>>>> Jet Correction \"" << this->name() << "\"" << endl ;
    //out() << "Use RATIO file \"" << path << "/" << info.MakeFileName() << endl ;
    out() << "Jet branch: "<<_fromBranch<<", Primary Vertex branch: "<<_vtxBranch<<endl;
    out() << endl ;
  }

  void JetCorr::begin(){
    _eff->setDebugLevel (debug()); // this is not available in the constructor
    //get all data epochs requested by user
    const vector<string>* data_epochs = MCREQID->data_epochs();
    if (!data_epochs->size()) {
      throw runtime_error( "caf_eff_utils::JetCorr[" + this->name() 
			   + "] ERROR: no data epoch specified in MCReqID");
    }
    
    //initialize p17 flag
    _isRunIIaMC=false;
    vtx0_p17 = new TMBPrimaryVertex(0,0,0,0,0,0,0,0,0,0,NULL,0);
    
  }

  void JetCorr::finish() {
    _rms = _rms >0 ? sqrt(_rms) : 0.0 ;
    out() << "Jet Correction \"" << name() 
	  << "\" average weight = " << _weight
	  << " RMS = " << _rms ;
      if (_events > 0) 
	out() << " RMS/Sqrt(N) = " << _rms/sqrt((double)_events) ;
      out () << endl ;
  }

  bool JetCorr::processEvent(cafe::Event &event)
  {

    //get pointer to statistics collector
    StatPointer stat ;
    event.get("StatPointer", stat) ;
    
    // verify if we need to upload new correction  
    change_request(event) ;

    double weight = 1.0 ;

    vector < EffVal > corrfactors;
    vector < int    > jetindex;
    
    if(!_isRunIIaMC){
    Collection<TMBPrimaryVertex> vtxs (event.getCollection<TMBPrimaryVertex>(_vtxBranch,_varsVtx));
    const TMBPrimaryVertex *vtx0 = (vtxs.size() != 0) ? &vtxs[0] : vtx0_p17; 
    if (debug()>19) out() << "#vtxs: " << vtxs.size() <<", vtx0: " << (int)vtx0 << endl;

    const float instlumi = 36.*event.getMCEventInfo(_varsMC)->overlay_instlum();
 
    Collection<TMBJet> from(event.getJets(_fromBranch.c_str(),_vars));
    for(Collection<TMBJet>::const_iterator it = from.begin(); 
	it != from.end(); it++) {	
      
      if (debug()>1) 
      out() << "======== JetCorr \"" << name() 
	    << "\"" << endl << " READ JET: pT = " << it->Pt() 
	    << " detEta = " << 0.1 * it->detEta() 
	    << " actas = " << it->CurrentlyActAs() 
	    << endl ;
      
      // remember the ActAs
      int actas = it->CurrentlyActAs();

      EffVal eff ;
      try {
	eff = _eff->Eff(&(*it), vtx0, instlumi) ;
      } catch (eff_utils::EffException& ex) {
	err() << "caf_eff_utils::JetCorr named \"" << name()
	      <<"\": exception catch " << ex._message
	      << endl << " for jet with pT = " << it->Pt() 
	      << " detEta = " << 0.1 * it->detEta() 
	      << " actas = " << it->CurrentlyActAs() 
	      << endl ;
	SetJESType(actas, *it); // Set the ActAs back
	return true ;
      }

      SetJESType(actas, *it); // Set the ActAs back
      
      /**out() << "======== JetCorr " << name() << endl
	    << " READ JET: pT = " << it->Pt()
            << " detEta = " << 0.1 * it->detEta()
	    << " SF = " << eff.val
	    << " from path = " << path
            << endl ;*/
      
      if(it->isZBjet_MC()) {eff.val=1.; eff.elo = 0.; eff.ehi = 0.;}
      weight *= eff.val ;
      
      // we need to propagate the individual weights with their individual
      // uncertainties to analysis processors
      // so they can evaluate the correlations properly
      corrfactors.push_back(eff);
      jetindex.push_back(it->index());

      if (_applyToGlobalWeight)
	stat.applyWeight("JetCorr: " + name(), eff.val) ;
    
      if (debug()>1) 
	out() << "======== JetCorr \"" << name() 
	      << "\"" << endl << " Weight = " << eff.val 
	      << "+" << eff.ehi
	      << "-" << eff.elo
	      << " Total event weight = " << weight 
	      << " (ActAs after: " << it->CurrentlyActAs() << " =? " << actas << ")"
	      << endl ;
      if (eff.val <0) 
	err() << "======== JetCorr \"" << name() 
	      << "\"" << endl << " Weight = " << eff.val 
	      << "+" << eff.ehi
	      << "-" << eff.elo
	      << " Total event weight = " << weight
	      << endl ;
    }

    if (weight <0 ) {
      err() << "caf_eff_utils::JetCorr named \"" << name()
	    <<"\": weight = " << weight
	    << " for event with leading jet with pT = " << from[0].Pt() 
	    << " detEta = " << 0.1 * from[0].detEta() 
	    << " actas = " << from[0].CurrentlyActAs() 
	    << endl ;
      return true ;
    }

    if (from.size()==0) {
      if (_applyToGlobalWeight)
	stat.applyWeight("JetCorr: " + name(), _weight) ;
    }
    }//end RunIIb condition
    
    else{
    Collection<TMBJet> from(event.getJets(_fromBranch.c_str(),_vars));
    for(Collection<TMBJet>::const_iterator it = from.begin(); 
	it != from.end(); it++) {	
      
      // remember the ActAs
      int actas = it->CurrentlyActAs();
      
      EffVal eff ;
      try {
	eff = _eff->Eff(&(*it), vtx0_p17, 100.) ;
      } catch (eff_utils::EffException& ex) {
	err() << "caf_eff_utils::JetCorr named \"" << name()
	      <<"\": exception catch " << ex._message
	      << endl << " for jet with pT = " << it->Pt() 
	      << " detEta = " << 0.1 * it->detEta() 
	      << " actas = " << it->CurrentlyActAs() 
	      << endl ;
	SetJESType(actas, *it); // Set the ActAs back
	return true ;
      }
      
      SetJESType(actas, *it); // Set the ActAs back
      
      weight *= 1;
      
      corrfactors.push_back(eff);
      jetindex.push_back(it->index());

      if (_applyToGlobalWeight)
	stat.applyWeight("JetCorr: " + name(), weight);
    }
    if (from.size()==0) {
      if (_applyToGlobalWeight)
	stat.applyWeight("JetCorr: " + name(), _weight) ;
    }
    }//end RunIIa condition
    
    _events++ ;
    
    event.put(name()+":weight", weight) ;
    event.put(name()+":weight_vector", corrfactors);
    event.put(name()+":jet_vector", jetindex);

    if (_events == 1) { 
      _weight = weight ;
      _rms = 0 ;
    } else if (_events > 1) {
      _rms = ((double) (_events-2))/(_events-1)*_rms + pow(weight-_weight,2)/_events ;
      _weight = ((double)(_events-1))/_events*_weight + weight/_events ;
    }

    event.put(name()+":average_weight", _weight) ;
    event.put(name()+":average_weight_rms", _rms) ;
    if (debug()>1) {
      out() << "======== JetCorr \"" << name() 
	    << "\"" << endl << "Average weight = " << _weight ;
      if (_events > 0) 
	out() << " +- " << sqrt(_rms/_events) ;
      out() << " RMS = " << sqrt(_rms)
	    << endl ;
    }

    return true  ;	
  }

  void JetCorr::change_request( cafe::Event &event ) {
    // check once again if MCREQID valid
    if( !MCREQID ) throw runtime_error("ERROR: in JetCorr, processor MCReqID is not initialized.") ;
    
    // verify if request id is stayed the same. In that case do nothing.
    if( MCREQID->reqid() == _reqid ) return ;
    _reqid = MCREQID->reqid() ;
    
    //get data epochs associated to this MC
    const vector<string>* epochs =  MCREQID->current_data_epochs();
    
    string mc_version = "" ;
    if      (event.isRun2a())  mc_version = "RunIIa";
    else if (event.isRun2b1()) mc_version = "RunIIb1";
    else if (event.isRun2b2()) mc_version = "RunIIb2";
    else if (event.isRun2b3()) mc_version = "RunIIb2";//temporary fix
    else throw runtime_error( "ERROR: in JetCorr, MC version unidentified") ;
    
    string data_epochs = "" ;
    for (vector<string>::const_iterator eit=epochs->begin();
         eit!=epochs->end(); eit++){
      if(data_epochs=="") data_epochs  = *eit;
      else data_epochs  = data_epochs + "_" + *eit;
    }
    
    if(mc_version=="RunIIa" && data_epochs=="RunIIa") path="jetid_eff/jetid/p17/";
    else if(mc_version=="RunIIb1" && data_epochs=="RunIIb1") path="jetid_eff/jetid/RunIIb_Summer2011/RunIIb1/";
    else if(mc_version=="RunIIb2" && data_epochs=="RunIIb2") path="jetid_eff/jetid/RunIIb_Summer2011/RunIIb2/";    
    else if(mc_version=="RunIIb2" && data_epochs=="RunIIb3") path="jetid_eff/jetid/RunIIb_Summer2011/RunIIb3/";    
    else if(mc_version=="RunIIb2" && data_epochs=="RunIIb4") path="jetid_eff/jetid/RunIIb_Summer2011/RunIIb4/";    
    else if(mc_version=="RunIIb2" && data_epochs=="RunIIb2_RunIIb3") path="jetid_eff/jetid/RunIIb_Summer2011/RunIIb2-3/";
    else if(mc_version=="RunIIb2" && data_epochs=="RunIIb2_RunIIb3_RunIIb4") path="jetid_eff/jetid/RunIIb_Summer2011/RunIIb2-4/";
    else if(mc_version=="RunIIb2" && data_epochs=="RunIIb3_RunIIb4") path="jetid_eff/jetid/RunIIb_Summer2011/RunIIb3-4/";
//not recommended but usuable
    else if(mc_version=="RunIIb1" && data_epochs=="RunIIb1_RunIIb2") path="jetid_eff/jetid/RunIIb_Summer2010/RunIIb1-3/";
    else if(mc_version=="RunIIb1" && data_epochs=="RunIIb1_RunIIb2_RunIIb3") path="jetid_eff/jetid/RunIIb_Summer2010/RunIIb1-3/";
    else if(mc_version=="RunIIb1" && data_epochs=="RunIIb1_RunIIb2_RunIIb3_RunIIb4") path="jetid_eff/jetid/RunIIb_Summer2010/RunIIb1-3/";
    else throw runtime_error(string("ERROR: in caf_eff_utils/JetCorr, the jet Id or VC scale factors for the MC-data configuration you are trying to use do not exist. Please correct your config file.")+ "\nThe mc event was from version: "+mc_version+" while the configured data epoch was: "+data_epochs) ;
    
    out() << "===================================================" << endl;
    out() << "The MC-data configuration is now: MC " << mc_version << " for " << data_epochs << " data." << endl; 
    out() << "The corresponding SFs will be read here: " << path << endl;
    out() << endl ;
    
    if(event.isRun2b3()){
      out() << "WARNING: the jet id SFs for Run IIb3 MC are not available yet! SFs derived on Run IIb2 MC sample are used." << endl;
    }
    if(data_epochs.find("RunIIb4")!=string::npos){
      out() << "WARNING: the jet id SFs for Run IIb4 data are not available yet! SFs derived on a dataset up to Run IIb3 are used." << endl;
    }
    
    _isRunIIaMC=false;
    if(mc_version=="RunIIa") _isRunIIaMC=true;
    
    delete _eff;
    try {
      _eff = new JetEffTool(info, path) ;
    } catch (eff_utils::EffException& ex) {
      err() << "caf_eff_utils::JetCorr named \"" << this->name()
	    <<"\": caught exception " << ex._message
	    << endl ;
      err() << info << endl ;
      exit(1) ;
    }
  }
  
}

ClassImp(caf_util::JetCorr) ;