'''

Class which figures out the correct PU weight.
Takes as input a MC PU tag (i.e. 'S7') and a set of .root files with the
pileup distributions (generated by pileupCalc.py [1])

[1] https://twiki.cern.ch/twiki/bin/view/CMS/PileupJSONFileforData

Author: Evan K. Friis, UW

'''

import array
import ROOT

# MC distributions (built at bottom of file)
_MC_PU_DISTRIBUTIONS = {}

class PileupWeight(object):
    def __init__(self, mctag, *datafiles):
        '''
        Build a PU weight object.

        [mctag] must reflect the PU distribution used in simulation and
        currently be either 'S6' or 'S7'.

        [datafiles] is a list of .root files generated a la

        https://twiki.cern.ch/twiki/bin/view/CMS/PileupJSONFileforData

        '''
        self.data = None
        for filename in datafiles:
            file = ROOT.TFile.Open(filename)
            pu = file.Get('pileup')
            if self.data is None:
                self.data = pu.Clone()
                self.data.SetDirectory(0)
            else:
                self.data.Add(pu)
            file.Close()

        # Normalize
        self.data.Scale(1./self.data.Integral())

        if not mctag in _MC_PU_DISTRIBUTIONS:
            raise KeyError("Unknown PU distribution %s, allowed: %s" %
                           (mctag, " ".join(_MC_PU_DISTRIBUTIONS.keys())))

        # Get same bin structure
        self.mc = self.data.Clone()
        # Clear bins
        self.mc.Reset()
        for i, bin in enumerate(_MC_PU_DISTRIBUTIONS[mctag]):
            self.mc.SetBinContent(i+1, bin)

        # Normalize MC
        self.mc.Scale(1./self.mc.Integral())

    def __call__(self, ntruepu):
        '''
        Get the PU weight given the true number of interactions
        '''
        bin = self.data.FindBin(ntruepu)
        data = self.data.GetBinContent(bin)
        mc = self.mc.GetBinContent(bin)
        if mc:
            return data/mc
        else:
            return 1.0

_MC_PU_DISTRIBUTIONS['S7'] = [
    2.344E-05,
    2.344E-05,
    2.344E-05,
    2.344E-05,
    4.687E-04,
    4.687E-04,
    7.032E-04,
    9.414E-04,
    1.234E-03,
    1.603E-03,
    2.464E-03,
    3.250E-03,
    5.021E-03,
    6.644E-03,
    8.502E-03,
    1.121E-02,
    1.518E-02,
    2.033E-02,
    2.608E-02,
    3.171E-02,
    3.667E-02,
    4.060E-02,
    4.338E-02,
    4.520E-02,
    4.641E-02,
    4.735E-02,
    4.816E-02,
    4.881E-02,
    4.917E-02,
    4.909E-02,
    4.842E-02,
    4.707E-02,
    4.501E-02,
    4.228E-02,
    3.896E-02,
    3.521E-02,
    3.118E-02,
    2.702E-02,
    2.287E-02,
    1.885E-02,
    1.508E-02,
    1.166E-02,
    8.673E-03,
    6.190E-03,
    4.222E-03,
    2.746E-03,
    1.698E-03,
    9.971E-04,
    5.549E-04,
    2.924E-04,
    1.457E-04,
    6.864E-05,
    3.054E-05,
    1.282E-05,
    5.081E-06,
    1.898E-06,
    6.688E-07,
    2.221E-07,
    6.947E-08,
    2.047E-08
]

_MC_PU_DISTRIBUTIONS['S6'] = [
     0.003388501,
    0.010357558,
    0.024724258,
    0.042348605,
    0.058279812,
    0.068851751,
    0.072914824,
    0.071579609,
    0.066811668,
    0.060672356,
    0.054528356,
    0.04919354,
    0.044886042,
    0.041341896,
    0.0384679,
    0.035871463,
    0.03341952,
    0.030915649,
    0.028395374,
    0.025798107,
    0.023237445,
    0.020602754,
    0.0180688,
    0.015559693,
    0.013211063,
    0.010964293,
    0.008920993,
    0.007080504,
    0.005499239,
    0.004187022,
    0.003096474,
    0.002237361,
    0.001566428,
    0.001074149,
    0.000721755,
    0.000470838,
    0.00030268,
    0.000184665,
    0.000112883,
    6.74043E-05,
    3.82178E-05,
    2.22847E-05,
    1.20933E-05,
    6.96173E-06,
    3.4689E-06,
    1.96172E-06,
    8.49283E-07,
    5.02393E-07,
    2.15311E-07,
    9.56938E-08
]

