import FWCore.ParameterSet.Config as cms
from Configuration.Generator.Pythia8CommonSettings_cfi import *
from Configuration.Generator.MCTunes2017.PythiaCP5Settings_cfi import *


FourMuonFilter = cms.EDFilter("FourLepFilter", # require 4-mu in the final state
        MinPt = cms.untracked.double(1.8),
        MaxPt = cms.untracked.double(4000.0),
        MaxEta = cms.untracked.double(2.5),
        MinEta = cms.untracked.double(0.),
        ParticleID = cms.untracked.int32(13)
)

DiJpsiFilter = cms.EDFilter("MCParticlePairFilter",  # require 2-mu mass to be 8.5 - 11.5 GeV
        MinPt = cms.untracked.vdouble(1.0,1.0),
        MaxPt = cms.untracked.vdouble(4000.0,4000.0),
        MaxEta = cms.untracked.vdouble( 2.5,2.5),
        MinEta = cms.untracked.vdouble(-2.5,-2.5),
        ParticleID1 = cms.untracked.vint32(443, -443),
        ParticleID2 = cms.untracked.vint32(443, -443),
        MinInvMass = cms.untracked.double(1.0),
        MaxInvMass = cms.untracked.double(16.0),
)

generator = cms.EDFilter("Pythia8ConcurrentGeneratorFilter",
    crossSection = cms.untracked.double(1.0),
    PythiaParameters = cms.PSet(
        pythia8CommonSettingsBlock,
        pythia8CP5SettingsBlock,
        processParameters = cms.vstring(
            'Bottomonium:states(3S1) = 551,100553,200553',
            'Bottomonium:gg2bbbar(3S1)[3S1(1)]g = on,off,off',
            'Bottomonium:gg2bbbar(3S1)[3S1(1)]gm = on,off,off',
            'Bottomonium:gg2bbbar(3S1)[3S1(8)]g = on,off,off',
            'Bottomonium:qg2bbbar(3S1)[3S1(8)]q = on,off,off',
            'Bottomonium:qqbar2bbbar(3S1)[3S1(8)]g = on,off,off',
            'Bottomonium:gg2bbbar(3S1)[1S0(8)]g = on,off,off',
            'Bottomonium:qg2bbbar(3S1)[1S0(8)]q = on,off,off',
            'Bottomonium:qqbar2bbbar(3S1)[1S0(8)]g = on,off,off',
            'Bottomonium:gg2bbbar(3S1)[3PJ(8)]g = on,off,off',
            'Bottomonium:qg2bbbar(3S1)[3PJ(8)]q = on,off,off',
            'Bottomonium:qqbar2bbbar(3S1)[3PJ(8)]g = on,off,off',
            'PartonLevel:MPI = on',
            'PartonLevel:ISR = on',
            'PartonLevel:FSR = on',
            'HadronLevel:all = on',
            'HadronLevel:Hadronize = on',
            'HadronLevel:Decay = on',
            '551:oneChannel 1 1.00 100 443 443',
            '551:onMode = off',
            '551:onIfMatch 443 443',
            '443:onMode = off',
            '443:onIfMatch 13 13',
            '333:onMode = off',
            '333:onIfMatch 13 13',
            '553:onMode = off',
            '553:onIfMatch 13 13',

        ),
        parameterSets = cms.vstring(
            'pythia8CommonSettings',
            'pythia8CP5Settings',
            'processParameters'
        ),
    ),
    comEnergy = cms.double(13000.0),
    maxEventsToPrint = cms.untracked.int32(0),
    pythiaHepMCVerbosity = cms.untracked.bool(False),
    pythiaPylistVerbosity = cms.untracked.int32(1)
)

ProductionFilterSequence = cms.Sequence(generator*FourMuonFilter*DiJpsiFilter)