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 *
from GeneratorInterface.Core.ExternalGeneratorFilter import ExternalGeneratorFilter

_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(
'#',
'# Particles updated from PDG2022 https://journals.aps.org/prd/abstract/10.1103/PhysRevD.98.030001',
'Particle   B+            5.2793300e+00   0.0000000e+00', ## 5279.33 521
'Particle   B-            5.2793300e+00   0.0000000e+00', ## 
'Particle   B*+           5.3247800e+00   0.0000000e+00', ## 5279.33 + 45.45 523
'Particle   B*-           5.3247800e+00   0.0000000e+00',
'Particle   B0            5.2796400e+00   0.0000000e+00', ## 5279.64 511 
'Particle  anti-B0        5.2796400e+00   0.0000000e+00', ## 5279.64 511 
'Particle   B*0           5.3253800e+00   0.0000000e+00', ## 5279.64 + 45.74 513
'Particle  anti-B*0       5.3253800e+00   0.0000000e+00', 
'Particle   B_s0          5.3669100e+00   0.0000000e+00', ## 531
'Particle  anti-B_s0      5.3669100e+00   0.0000000e+00',
'Particle   B_s*0         5.4167100e+00   0.0000000e+00', ## 5366.91 + 49.80 533
'Particle  anti-B_s*0     5.4167100e+00   0.0000000e+00', ## 
'Particle   J/psi         3.0969000e+00   9.2600000e-05', ## 443
'Particle   psi(2S)       3.68610000+00   2.9400000e-04', ## 100443 ## chi_c0,1,2 are ok

'Alias      MyB+             B+',
'Alias      MyB-             B-',
'ChargeConj MyB-             MyB+',
'Alias      MyB*+            B*+',
'Alias      MyB*-            B*-',
'ChargeConj MyB*-            MyB*+',

'Alias      MyJ/psi  J/psi',
'ChargeConj MyJ/psi  MyJ/psi',
'Alias       Mypsi(2S)  psi(2S)',
'ChargeConj  Mypsi(2S)  Mypsi(2S)',

'Decay MyB*+',
'1.000        MyB+         gamma    VSP_PWAVE;',
'Enddecay',
'CDecay MyB*-',

'Decay MyB+',
'1.000  Mypsi(2S)              K+        SVS;',
'Enddecay',
'CDecay MyB-',

'Decay Mypsi(2S)',
'1.000         MyJ/psi  pi+   pi-       VVPIPI;',
'Enddecay',

'Decay MyJ/psi',
'1.000         mu+         mu-      PHOTOS VLL;',
'Enddecay',

'End'), 
            list_forced_decays = cms.vstring('MyB*+','MyB*-'),
            operates_on_particles = cms.vint32(523, -523),
            convertPythiaCodes = cms.untracked.bool(False)
            ),
        parameterSets = cms.vstring('EvtGen130')
        ),
        PythiaParameters = cms.PSet(
        pythia8CommonSettingsBlock,
        pythia8CP5SettingsBlock,
        processParameters = cms.vstring(
            "SoftQCD:nonDiffractive = on",
            "523:m0=5.32478",     ## changing also B mass in pythia
            'PTFilter:filter = on', # this turn on the filter
            'PTFilter:quarkToFilter = 5', # PDG id of q quark
            'PTFilter:scaleToFilter = 3.0'),
        parameterSets = cms.vstring(
            'pythia8CommonSettings',
            'pythia8CP5Settings',
            'processParameters',
        )
    )
)

generator = ExternalGeneratorFilter(_generator)
###########
# Filters #
###########
# Filter only pp events which produce needed b hadron
bfilter = cms.EDFilter("PythiaFilter", ParticleID = cms.untracked.int32(523))

decayfilter = cms.EDFilter(
        "PythiaDauVFilter",
  verbose         = cms.untracked.int32(0),
  NumberDaughters = cms.untracked.int32(2),
  MotherID        = cms.untracked.int32(0),
  ParticleID      = cms.untracked.int32(523),
  DaughterIDs     = cms.untracked.vint32(521, 22),
  MinPt           = cms.untracked.vdouble(12.0, 0.25),
  MinEta          = cms.untracked.vdouble(-3.0, -3.0),
  MaxEta          = cms.untracked.vdouble( 3.0,  3.0)
        )

bgroundfilter = cms.EDFilter(
        "PythiaDauVFilter",
  verbose         = cms.untracked.int32(0),
  NumberDaughters = cms.untracked.int32(2),
  MotherID        = cms.untracked.int32(523),
  ParticleID      = cms.untracked.int32(521),
  DaughterIDs     = cms.untracked.vint32(100443, 321),
  MinPt           = cms.untracked.vdouble(6.9, 0.7),
  MinEta          = cms.untracked.vdouble(-3.0, -3.0),
  MaxEta          = cms.untracked.vdouble( 3.0,  3.0)
        )

pjppfilter = cms.EDFilter(
        "PythiaDauVFilter",
  verbose         = cms.untracked.int32(0),
  NumberDaughters = cms.untracked.int32(3),
  MotherID        = cms.untracked.int32(521),
  ParticleID      = cms.untracked.int32(100443),
  DaughterIDs     = cms.untracked.vint32(443, 211, -211),
  MinPt           = cms.untracked.vdouble(6.9, 0.4, 0.4),
  MinEta          = cms.untracked.vdouble(-3.0, -3.0, -3.0),
  MaxEta          = cms.untracked.vdouble( 3.0,  3.0,  3.0)
        )

jmmfilter = cms.EDFilter(
        "PythiaDauVFilter",
  verbose         = cms.untracked.int32(0),
  NumberDaughters = cms.untracked.int32(2),
  MotherID        = cms.untracked.int32(100443),
  ParticleID      = cms.untracked.int32(443),
  DaughterIDs     = cms.untracked.vint32(13, -13),
  MinPt           = cms.untracked.vdouble(2.8, 2.8),
  MinEta          = cms.untracked.vdouble(-2.6, -2.6),
  MaxEta          = cms.untracked.vdouble( 2.6,  2.6)
        )


ProductionFilterSequence = cms.Sequence(generator*bfilter*decayfilter*bgroundfilter*pjppfilter*jmmfilter)