import FWCore.ParameterSet.Config as cms

from Configuration.Generator.Pythia8CommonSettings_cfi import *
from Configuration.Generator.MCTunes2017.PythiaCP5Settings_cfi import *
from GeneratorInterface.EvtGenInterface.EvtGenSetting_cff import *

_generator = cms.EDFilter("Pythia8GeneratorFilter",
                         comEnergy = cms.double(13000.0),
                         pythiaHepMCVerbosity = cms.untracked.bool(False),
                         maxEventsToPrint = cms.untracked.int32(0),
                         pythiaPylistVerbosity = cms.untracked.int32(0),
                         ExternalDecays = cms.PSet(
        EvtGen130 = cms.untracked.PSet(
            decay_table = cms.string('GeneratorInterface/EvtGenInterface/data/DECAY_2014_NOLONGLIFE.DEC'),
            particle_property_file = cms.FileInPath('GeneratorInterface/EvtGenInterface/data/evt_2014.pdl'),
            user_decay_embedded= cms.vstring(
'#',
'Particle   pi+           1.3957061e-01   0.0000000e+00',
'Particle   pi-           1.3957061e-01   0.0000000e+00',
'Particle   K_S0          4.9761100e-01   0.0000000e+00',
'Particle   K*+           8.9176000e-01   5.0300000e-02',
'Particle   K*-           8.9176000e-01   5.0300000e-02',
'Particle   K*0           8.9555000e-01   4.7300000e-02',
'Particle   anti-K*0      8.9555000e-01   4.7300000e-02',
'Particle   rho0          7.7526000e-01   1.4910000e-01',
'Particle   phi           1.0194610e+00   4.2490000e-03', ## 333
'Particle   B-            5.2793200e+00   0.0000000e+00',
'Particle   B+            5.2793200e+00   0.0000000e+00',
'Particle   B0            5.2796300e+00   0.0000000e+00',
'Particle   anti-B0       5.2796300e+00   0.0000000e+00',
'Particle   B_s0          5.3668900e+00   0.0000000e+00', ## id 531
'Particle   anti-B_s0     5.3668900e+00   0.0000000e+00',
'Particle   J/psi         3.0969000e+00   9.2900006e-05', ## 443
'Particle   psi(2S)       3.6860970e+00   2.9400000e-04',
'#',
'#',
'Alias      MyB0   B0',
'Alias      Myanti-B0   anti-B0',
'ChargeConj Myanti-B0   MyB0 ',
'#',
'Alias      MyPsiT   psi(2S)',
'ChargeConj MyPsiT   MyPsiT',
'#',
'Alias      MyJ/psi  J/psi',
'ChargeConj MyJ/psi  MyJ/psi',
'#',
'Alias      MyPhi    phi',
'ChargeConj MyPhi    MyPhi',
'#',
'Alias      MyK*0    K*0',
'Alias      Myanti-K*0    anti-K*0',
'ChargeConj MyK*0    Myanti-K*0',
'#',
'#',
'Decay MyB0',
'  1.000         MyPsiT      MyK*0     SVV_HELAMP 0.159 1.563 0.775 0.0 0.612 2.712;',
'Enddecay',
'#',
'Decay Myanti-B0',
'  1.000         MyPsiT      Myanti-K*0   SVV_HELAMP 0.159 1.563 0.775 0.0 0.612 2.712;',
'Enddecay',
'#',
'Decay MyPsiT',
'1.000   MyJ/psi  pi+   pi-         VVPIPI;',
'Enddecay',
'#',
'Decay MyJ/psi',
'  1.000         mu+         mu-          PHOTOS VLL;',
'Enddecay',
'#',
'Decay MyK*0',
'  1.000         K+         pi-          VSS;',
'Enddecay',
'#',
'Decay Myanti-K*0',
'  1.000         K-         pi+          VSS;',
'Enddecay',
'End'
), 
            list_forced_decays = cms.vstring('MyB0', 'Myanti-B0'),
            operates_on_particles = cms.vint32(),
            convertPythiaCodes = cms.untracked.bool(False)
            ),
        parameterSets = cms.vstring('EvtGen130')
        ),
        PythiaParameters = cms.PSet(
        pythia8CommonSettingsBlock,
        pythia8CP5SettingsBlock,
        processParameters = cms.vstring(
            "SoftQCD:nonDiffractive = on",
            "511:m0=5.2796300",       ## changing also Xi_b- mass in pythia
            'PTFilter:filter = on', # this turn on the filter
            'PTFilter:quarkToFilter = 5', # PDG id of q quark
            'PTFilter:scaleToFilter = 1.0'),
        parameterSets = cms.vstring(
            'pythia8CommonSettings',
            'pythia8CP5Settings',
            'processParameters',
        )
    )
)

from GeneratorInterface.Core.ExternalGeneratorFilter import ExternalGeneratorFilter
generator = ExternalGeneratorFilter(_generator)

###########
# Filters #
###########
#

bfilter = cms.EDFilter(
    "PythiaFilter",
    MaxEta = cms.untracked.double(9999.),
    MinEta = cms.untracked.double(-9999.),
    ParticleID = cms.untracked.int32(511)
)

jpsifilter = cms.EDFilter(
    "PythiaDauVFilter",
    MotherID = cms.untracked.int32(100443),
    ParticleID = cms.untracked.int32(443),
    NumberDaughters = cms.untracked.int32(2),
    DaughterIDs = cms.untracked.vint32(13, -13),
    MinPt = cms.untracked.vdouble(3.0, 3.0),
    MinEta = cms.untracked.vdouble(-2.5, -2.5),
    MaxEta = cms.untracked.vdouble(2.5, 2.5),
    verbose = cms.untracked.int32(1)
)

psi2sfilter = cms.EDFilter(
        "PythiaDauVFilter",
    verbose         = cms.untracked.int32(0),
    NumberDaughters = cms.untracked.int32(3),
    MotherID        = cms.untracked.int32(511),
    ParticleID      = cms.untracked.int32(100443),
    DaughterIDs     = cms.untracked.vint32(443, 211, -211),
    MinPt           = cms.untracked.vdouble(5.0, 0.5, 0.5),
    MinEta          = cms.untracked.vdouble(-9999, -3.0, -3.0),
    MaxEta          = cms.untracked.vdouble( 9999,  3.0,  3.0)
        )

kstarfilter = cms.EDFilter(
    "PythiaDauVFilter",
    MotherID = cms.untracked.int32(511),
    ParticleID = cms.untracked.int32(313),
    NumberDaughters = cms.untracked.int32(2),
    DaughterIDs = cms.untracked.vint32(321, -211),
    MinPt = cms.untracked.vdouble(0.3, 0.3),
    MinEta = cms.untracked.vdouble(-2.5, -2.5),
    MaxEta = cms.untracked.vdouble(2.5, 2.5),
    verbose = cms.untracked.int32(1)
)

ProductionFilterSequence = cms.Sequence(generator*bfilter*psi2sfilter*jpsifilter*kstarfilter)