/cvmfs/atlas.cern.ch/repo/sw/ASG/AnalysisBase/2.4.31/PileupReweighting/PileupReweighting/TPileupReweighting.h

// Dear emacs, this is -*-c++-*-

#ifndef __TPILEUPREWEIGHTING__
#define __TPILEUPREWEIGHTING__

/* Developer notes:
* Could have incorporated the pileup histos in to the generic custom maps
* but decided to keep them seperate for probable slight performance improvement
* Allow the user to change the histo binning through AddBinning with name="pileup"
* */

#include "TNamed.h"
#include <TFile.h>
#include <TString.h>
#include "TVectorD.h"
#include <map>
#include <vector>
#include <TRandom3.h>

// STL includes
#include <iostream>
#include <stdexcept>

#include "TH1.h"


class TH1;
class TH1D;
class TH2D;
class TTree;
class TFile;



namespace CP {
  class TPileupReweighting : public TNamed {

  public: 
      TPileupReweighting(const char* name="TPileupReweighting");

      ~TPileupReweighting();


  public:

      Int_t UsePeriodConfig(const TString& configName);
      Int_t SetBinning(Int_t nbinsx, Double_t* xbins, Int_t nbinsy=0, Double_t* ybins=0);
      Int_t SetUniformBinning(Int_t nbinsx, Double_t xlow, Double_t xup, Int_t nbinsy=0, Double_t ylow=0, Double_t yup=0);
      Int_t SetBinning(TH1* hist);

      void SetDefaultChannel(Int_t channel, Int_t mcRunNumber=-1);
      Int_t GetDefaultChannel(Int_t mcRunNumber=-1);

      Double_t GetIntegratedLumi(const TString& trigger="");
      inline Double_t GetNumberOfEvents(Int_t channel) {
         Period* global = m_periods[-1];
         if(!global) return 0;
         if(global->numberOfEntries.find(channel)==global->numberOfEntries.end()) {
            Error("GetNumberOfEvents", "Unknown channel: %d",channel);
            return 0;
         }
         return global->numberOfEntries[channel];
      }
      inline Double_t GetSumOfEventWeights(Int_t channel) {
         Period* global = m_periods[-1];
         if(!global) return 0;
         if(global->sumOfWeights.find(channel)==global->sumOfWeights.end()) {
            Error("GetSumOfEventWeights", "Unknown channel: %d",channel);
            return 0;
         }
         return global->sumOfWeights[channel];
      }

      void RemapPeriod(Int_t periodNumber1, Int_t periodNumber2);

      
      Double_t GetIntegratedLumiFraction(Int_t periodNumber, UInt_t start, UInt_t end);
      Double_t GetIntegratedLumiFraction(Int_t periodNumber, Double_t mu, UInt_t start, UInt_t end);

      inline Double_t GetIntegratedLumi(UInt_t start, UInt_t end) {    return GetIntegratedLumi(-1,start,end); }
      Double_t GetIntegratedLumi(Int_t periodNumber, UInt_t start, UInt_t end);

      Double_t GetLumiBlockIntegratedLumi(Int_t runNumber, UInt_t lb);

      Float_t GetLumiBlockMu(Int_t runNumber, UInt_t lb);

      //-----------------------------------------------------
      //General Tool Configuring methods
      //-----------------------------------------------------
      inline void DisableWarnings(Bool_t in) { m_SetWarnings = !in;}
      inline void EnableDebugging(Bool_t in) { m_debugging = in;}

      inline void SetUnrepresentedDataAction(Int_t action, Double_t tolerance=0.05) {
         if(action<0 || action>3) {
            Fatal("SetUnrepresentedDataAction","Set action=%d - I'm afraid I can't let you do that, Dave",action);
            throw std::runtime_error("Throwing 2001");
         } 
         m_unrepresentedDataAction=action;
         m_unrepDataTolerance=tolerance; //applicable for action=2 mode. Default is 5%
      }
      Double_t GetUnrepresentedDataFraction(Int_t periodNumber,Int_t channel);
      inline Float_t GetUnrepresentedDataWeight(Int_t periodNumber,Int_t channel) { 
         if(m_unrepresentedDataAction!=2) {
            Warning("GetUnrepresentedDataWeight","You should not be applying this weight unless the UnrepresentedDataAction=2");
         }
         return 1./(1. - GetUnrepresentedDataFraction(periodNumber,channel));
      }
      //-----------------------------------------------------
      //Methods to veto data in Action=1 mode
      //-----------------------------------------------------
      Bool_t IsUnrepresentedData(Int_t runNumber, Float_t x, Float_t y=0.);

      inline void IgnoreConfigFilePeriods(Bool_t in) { m_ignoreFilePeriods=in; }

      Int_t AddPeriod(Int_t periodNumber, UInt_t start, UInt_t end);
      Int_t GetFirstFoundPeriodNumber(UInt_t runNumber);



      inline void SetDataScaleFactors(Float_t x,Float_t y=1.) { m_dataScaleFactorX=x;m_dataScaleFactorY=y; }
      inline void SetMCScaleFactors(Float_t x,Float_t y=1.) { m_mcScaleFactorX=x;m_mcScaleFactorY=y;}



      //-----------------------------------------------------
      //Methods to load config files
      //-----------------------------------------------------
      Int_t AddConfigFile(const TString& fileName);
      Int_t AddLumiCalcFile(const TString& fileName, const TString& trigger="None");
      Int_t AddMetaDataFile(const TString& fileName,const TString& channelBranchName="mc_channel_number");

      Bool_t RemoveChannel(int chanNum);

      Int_t Initialize();


      //-----------------------------------------------------
      //Methods to get the various weights
      //-----------------------------------------------------
      Float_t GetCombinedWeight(Int_t periodNumber, Int_t channelNumber,Float_t x,Float_t y=0.);
      Float_t GetPeriodWeight(Int_t periodNumber, Int_t channelNumber);
      Float_t GetPrimaryWeight(Int_t periodNumber, Int_t channelNumber,Float_t x);
      Float_t GetSecondaryWeight(Int_t periodNumber, Int_t channelNumber,Float_t x,Float_t y);

      //-----------------------------------------------------
      //Methods to work with the metadata
      //-----------------------------------------------------
      Double_t GetMetaData(const TString& metadataName,Int_t channelNumber) {
         if(m_metadata.find(metadataName)==m_metadata.end()) {
            Error("GetMetaData","Metadata %s not known",metadataName.Data());
            return 0;
         }
         if(m_metadata[metadataName].find(channelNumber)==m_metadata[metadataName].end()) {
            Error("GetMetaData","Metadata %s not known in channel %d",metadataName.Data(), channelNumber);
            return 0;
         }
         return m_metadata[metadataName][channelNumber];
      }
      TTree* GetMetaDataTree(); 
      Int_t GenerateMetaDataFile(const TString& fileName,const TString& channelBranchName="mc_channel_number");


      //-----------------------------------------------------
      //Methods to generate config files
      //-----------------------------------------------------
      Int_t Fill(Int_t runNumber,Int_t channelNumber,Float_t w,Float_t x, Float_t y=0.);
      Int_t WriteToFile(const TString& filename=""); //if no name given, will use tool name
      Int_t WriteToFile(TFile* outFile);


      //
      //Methods used to do the prwTree - for fast recalculation of prw weights when processing ntuples
      //
      ULong64_t GetPRWHash(Int_t periodNumber, Int_t channelNumber, Float_t x, Float_t y=0.);
      Bool_t MakeWeightTree(TString channelNumbers, TString outFile, TString hashBranch="PRWHash", TString weightBranch="PileupWeight");

      //-----------------------------------------------------
      //RandomDataPeriod functionality stuff
      //numbers generated seperately for each mc period
      //-----------------------------------------------------
      void SetRandomSeed(int seed) {m_random3->SetSeed(seed);}
      int GetRandomSeed() {return m_random3->GetSeed();}
      UInt_t GetRandomRunNumber(Int_t mcRunNumber);
      UInt_t GetRandomRunNumber(Int_t mcRunNumber, Double_t mu);
      Int_t GetRandomPeriodNumber(Int_t mcRunNumber);
      UInt_t GetRandomLumiBlockNumber(UInt_t runNumber);


      //-----------------------------------------------------
      //Methods for PROOF cluster merging of generated histos
      //-----------------------------------------------------
      Int_t Merge(TCollection *coll);
      //*std::map<Int_t,std::map<Int_t, TH1*> > & GetInputHistograms() { return m_inputHistograms;}


      //-----------------------------------------------------
      //Methods to inspect the input and weighting histograms
      //-----------------------------------------------------
      TH1* GetInputHistogram(Int_t channelNumber, Int_t periodNumber) {
         if(m_periods.find(periodNumber)==m_periods.end()||
            m_periods[periodNumber]->inputHists.find(channelNumber)==m_periods[periodNumber]->inputHists.end()) {
            return 0;
         }
         return m_periods[periodNumber]->inputHists[channelNumber];
      }

      TH1D* GetPrimaryDistribution(Int_t channelNumber, Int_t periodNumber) {
         if(m_periods.find(periodNumber)==m_periods.end()||
            m_periods[periodNumber]->primaryHists.find(channelNumber)==m_periods[periodNumber]->primaryHists.end()) {
            return 0;
         }
         return m_periods[periodNumber]->primaryHists[channelNumber];
      }

      TH1D* GetPrimaryTriggerDistribution(const TString& trigger, Int_t periodNumber, long triggerBits)  {
         if(m_triggerObjs.find(trigger)==m_triggerObjs.end() || m_triggerObjs[trigger]->triggerHists.find(periodNumber) ==m_triggerObjs[trigger]->triggerHists.end() ||
            m_triggerObjs[trigger]->triggerHists[periodNumber].find(triggerBits) == m_triggerObjs[trigger]->triggerHists[periodNumber].end()) return 0;
      /*
         if(m_triggerObjs.find(trigger)==m_triggerObjs.end() || m_triggerObjs.at(trigger).second->triggerHists.find(periodNumber)==m_triggerObjs.at(trigger).second->->triggerHists.end() ||
            m_triggerObjs.at(trigger).second->triggerHists.at(periodNumber).find(triggerBits) == m_triggerObjs.at(trigger).second.at(periodNumber).end()) return 0;
      */
         
         return m_triggerObjs[trigger]->triggerHists[periodNumber][triggerBits];
      }

      Double_t GetDataWeight(Int_t runNumber, TString trigger, Double_t x);
      Double_t GetDataWeight(Int_t runNumber, TString trigger);//version without mu dependence


      // other methods 
      Bool_t IsInitialized() { return m_isInitialized; }
// 
      Int_t AddDistribution(TH1* hist, Int_t runNumber, Int_t channelNumber);

      void ResetCountingMode() { m_countingMode=true; }

      
      void SetParentTool(TPileupReweighting* tool) { m_parentTool = tool; }

      Float_t GetDataScaleFactor() const { return m_dataScaleFactorX; }

      void SetTriggerBit(const TString& trigger, bool in=true) { m_triggerPassBits[trigger]=in; }
      void ResetTriggerBits() { m_triggerPassBits.clear(); }

      double GetRunAverageMu(int run) { return m_runs[run].inputHists["None"]->GetMean(); }

  protected:
      virtual bool runLbnOK(Int_t /*runNbr*/, Int_t /*lbn*/) { return true; } //override in the ASG tool
      virtual bool passTriggerBeforePrescale(const TString& trigger) const { 
        if(m_triggerPassBits.size()==0) return true;
        try {
          return m_triggerPassBits.at(trigger); 
        } catch(...) { return true; }
      } //override in the ASG tool
      std::map<TString, bool> m_triggerPassBits;
      
      
      Int_t GetNearestGoodBin(Int_t thisMCRunNumber, Int_t bin);
      Int_t IsBadBin(Int_t thisMCRunNumber, Int_t bin);


      TH1* CloneEmptyHistogram(Int_t runNumber, Int_t channelNumber);
      void normalizeHistogram(TH1* histo);
      void AddDistributionTree(TTree *tree, TFile *file);
      //*Int_t FactorizeDistribution(TH1* hist, const TString weightName, Int_t channelNumber, Int_t periodNumber,bool includeInMCRun,bool includeInGlobal);

      void CalculatePrescaledLuminosityHistograms(const TString& trigger);

      //********** Private members*************************
      TPileupReweighting* m_parentTool; //points to self if not a 'systematic varion' tool instance
      Bool_t m_SetWarnings;
      Bool_t m_debugging;
      Bool_t m_countingMode;
      Int_t m_unrepresentedDataAction;                         //used as property in wrapper tool
      Bool_t m_isInitialized;Bool_t m_lumiVectorIsLoaded;
      Float_t m_dataScaleFactorX;Float_t m_dataScaleFactorY;   //first is used as property in wrapper tool
      Float_t m_mcScaleFactorX;Float_t m_mcScaleFactorY;
      Int_t m_nextPeriodNumber;
      Bool_t m_ignoreFilePeriods;
      TTree *m_metadatatree;
      Double_t m_unrepDataTolerance;
      Bool_t m_doGlobalDataWeight; //used in GetDataWeight to flag mu-independent version
      Int_t m_lumicalcRunNumberOffset; //used for 'faking' a lumicalc file for run2

      std::map<TString,std::vector<TString> > m_lumicalcFiles;

      //-----------------------------------------------------
      //Shared private data members
      //-----------------------------------------------------



      TH1* m_emptyHistogram; 



      std::map<TString, std::map<Int_t, Double_t> > m_metadata; 

public:
      struct CompositeTrigger {
         int op;
         CompositeTrigger* trig1;
         CompositeTrigger* trig2;
         TString val;
         std::vector<TString> subTriggers; //only set for top-level object
         CompositeTrigger() : op(0),trig1(0),trig2(0),val("") { }
         ~CompositeTrigger() { if(trig1) delete trig1; if(trig2) delete trig2; }
         double eval(std::map<TString, std::map<Int_t, std::map<Int_t, Float_t> > >& m, int run, int lbn, const TPileupReweighting* tool) {
            switch(op) {
               case 0: if(m[val][run].find(lbn)==m[val][run].end() || !m[val][run][lbn] || !tool->passTriggerBeforePrescale(val)) return 0; /*trigger/failed disabled, so cannot contribute*/   
                        return 1./m[val][run][lbn];
               case 1: return 1. - (1. - trig1->eval(m,run,lbn,tool))*(1.-trig2->eval(m,run,lbn,tool)); //OR
               case 2: return trig1->eval(m,run,lbn,tool)*trig2->eval(m,run,lbn,tool);
               default: return 1;
            }
         }
         void getTriggers(std::vector<TString>& s) {
            if(trig1==0&&trig2==0&&val.Length()>0) s.push_back(val);
            else { trig1->getTriggers(s); trig2->getTriggers(s); }
         }
         
         long getBits(const TPileupReweighting* tool) {
          long out(0);
          for(uint i=0;i<subTriggers.size();i++) out += (tool->passTriggerBeforePrescale(subTriggers[i]) << i);
          return out;
         }
         
         std::map<int, std::map<long, TH1D*> > triggerHists; //unnormalized ... i.e. integral should be equal to the lumi! ... indexed by PeriodID,tbits
         
      };
protected:

      std::map<TString, CompositeTrigger*> m_triggerObjs; //map from trigger string to composite trigger object

      CompositeTrigger* makeTrigger(const TString& s);
      void calculateHistograms(CompositeTrigger* trigger);

public:
      inline void PrintPeriods() { for(auto p : m_periods) {std::cout << p.first << " -> "; p.second->print("");} }
      struct Period {
         Period(Int_t _id, UInt_t _start, UInt_t _end, Int_t _defaultChannel) : id(_id),start(_start),end(_end),defaultChannel(_defaultChannel) {};
         Int_t id;
         UInt_t start;
         UInt_t end;
         Int_t defaultChannel;
         std::map<Int_t, Int_t> inputBinRedirect;
         std::map<Int_t, Double_t> unrepData; //indexed by channel number
         std::vector<Period*> subPeriods; 
         std::vector<UInt_t> runNumbers; //populated with runNumbers that actually had some data available
         std::map<Int_t, TH1*> inputHists;
         std::map<Int_t, Double_t> sumOfWeights; 
         std::map<Int_t, Int_t> numberOfEntries;
         std::map<Int_t, TH1D*> primaryHists; //normalized histograms, indexed by channelNumber. -1 holds the data
         std::map<Int_t, TH2D*> secondaryHists; //semi-normalized histograms, only used in 2D reweighting
         //std::map<TString, TH1D*> triggerHists; //unnormalized ... i.e. integral should be equal to the lumi! ... should really only be filled in the channel=-1 case (i.e. data)
         bool contains(unsigned int runNumber) {
            if(runNumber >= start && runNumber <= end) return true;
            for(auto p : subPeriods) if(p->contains(runNumber)) return true;
            return false;
         };
         void SetDefaultChannel(Int_t channel) {
            defaultChannel=channel;
            for(auto p : subPeriods) p->SetDefaultChannel(channel);
         };
         void print(const char* prefix) { 
            std::cout << prefix << id << "[" << start << "," << end << "] : ";
            for(auto hist : inputHists) std::cout << hist.first << " , ";
            std::cout << std::endl;
            for(auto p : subPeriods) p->print(TString::Format(" %s",prefix).Data()); };
      };
      struct Run {
         std::map<TString,TH1*> inputHists; //key is the 'trigger' 
         std::map<Int_t,Double_t> badBins; 
         Double_t lumi; //total data in run
         std::map<UInt_t, std::pair<Double_t,Double_t> > lumiByLbn; //key=lbn, value = <lumi,mu>
         TH1D* muDist; //mu distribution for this run
      };
protected:
      std::map<Int_t, Period*> m_periods; //periods mapped by id. -1 = the "global" period (0->9999999). Uses a pointer so can easily implement remap as two entries pointing at same period
      std::map<UInt_t, Run> m_runs; //runs mapped by runNumber

      std::map<Int_t, Double_t> unrepDataByChannel; // -1 holds the total unrepData!

      std::string m_prwFilesPathPrefix;

      //-----------------------------------------------------
      //RandomDataPeriod functionality stuff
      //numbers generated seperately for each mc period
      //-----------------------------------------------------
      TRandom3 *m_random3;

      Bool_t m_ignoreBadChannels; //if true, will print a warning about any channels with too much unrepresented data, but will then just ignore them
      Bool_t m_useMultiPeriods = true; //if true, will allow for runDependentMC 
      
public:
      //this method is a convenience for copying over properties to the tools that are used for systematic variations
      inline void CopyProperties(const TPileupReweighting* in) {
         m_prwFilesPathPrefix = in->m_prwFilesPathPrefix;
         m_ignoreBadChannels = in->m_ignoreBadChannels;
         m_unrepresentedDataAction = in->m_unrepresentedDataAction;
         m_unrepDataTolerance= in->m_unrepDataTolerance;
      }

      ClassDef(TPileupReweighting,0)


  }; // End: class definition

} // End: namespace Root

#endif

---