from MadGraphControl.MadGraphUtils import *

mode=0
mhh=2.700000e+02

#---------------------------------------------------------------------------------------------------
# Setting mHH and WHH for param_card.dat
#---------------------------------------------------------------------------------------------------
parameters={'1560':'2.700000e+02', #MHH
            '1561':'1.000000e-02'} #WHH

#---------------------------------------------------------------------------------------------------
# Setting higgs mass to 125 GeV for param_card.dat
#---------------------------------------------------------------------------------------------------
higgsMass={'25':'1.250000e+02'} #MH

#---------------------------------------------------------------------------------------------------
# Setting some parameters for run_card.dat: Parton shower is HERWIG7
#---------------------------------------------------------------------------------------------------
extras = { 'lhe_version':'2.0', 
           'cut_decays':'F', 
           'pdlabel':"'nn23lo1'",
           'fixed_ren_scale': 'True', # 
           'fixed_fac_scale': 'True', #
           'scale':'135', #
           'dsqrt_q2fact1':'135', #
           'dsqrt_q2fact2':'135', # 
#           'parton_shower':'Herwig7',
           'ptj':'0',
           'ptb':'0',
           'pta':'0',
           'ptjmax':'-1',
           'ptbmax':'-1',
           'ptamax':'-1',
           'etaj':'-1',
           'etab':'-1',
           'etaa':'-1',
           'etajmin':'0',
           'etabmin':'0',
           'etaamin':'0',
           'mmaa':'0',
           'mmaamax':'-1',
           'mmbb':'0',
           'mmbbmax':'-1',
           'drjj':'0',
           'drbb':'0',
           'draa':'0',
           'drbj':'0',
           'draj':'0',
           'drab':'0',
           'drjjmax':'-1',
           'drbbmax':'-1',
           'draamax':'-1',
           'drbjmax':'-1',
           'drajmax':'-1',
           'drabmax':'-1' }

#---------------------------------------------------------------------------------------------------
# Generating di-higgs through Heavy Higgs resonance with MadGraph
#---------------------------------------------------------------------------------------------------
fcard = open('proc_card_mg5.dat','w')
if (runArgs.runNumber == 451892):
    fcard.write("""
    import model sm
    define p = g u c d s u~ c~ d~ s~
    define j = g u c d s u~ c~ d~ s~
    import model HeavyHiggsTHDM
    generate p p > hh > h h 
    output -f""")
    fcard.write('set Mhh %s \n'%(mhh))
    fcard.close()
else: 
    raise RuntimeError("runNumber %i not recognised in these jobOptions."%runArgs.runNumber)

beamEnergy=-999
if hasattr(runArgs,'ecmEnergy'):
    beamEnergy = runArgs.ecmEnergy / 2.
else: 
    raise RuntimeError("No center of mass energy found.")

#---------------------------------------------------------------------------------------------------
# Filter efficiency is low
# Thus, setting the number of generated events to 100 times maxEvents,
# to avoid crashing due to not having enough events
# Also putting protection to avoid from crashing when maxEvents=-1
#---------------------------------------------------------------------------------------------------
safefactor=100
nevents=10000*safefactor
if runArgs.maxEvents > 0:
    nevents=runArgs.maxEvents*safefactor

runName='run_01'
process_dir = new_process(card_loc='proc_card_mg5.dat')

#---------------------------------------------------------------------------------------------------
# Using the helper function from MadGraphControl for setting up the param_card
# Build a new param_card.dat from an existing one
# Used values given in "parameters" for MHH and WHH, if not set there, default values are used
# Higgs mass is set to 125 GeV by "higgsMass"
#---------------------------------------------------------------------------------------------------

build_param_card(param_card_old='param_card.HeavyScalar.dat',param_card_new='param_card_new.dat',masses=higgsMass,extras=parameters)

#---------------------------------------------------------------------------------------------------
# Using the helper function from MadGraphControl for setting up the run_card
# Build a new run_card.dat from an existing one
# Using the values given in "extras" above for the selected parameters when setting up the run_card
# If not set in "extras", default values are used 
#---------------------------------------------------------------------------------------------------

build_run_card(run_card_old=get_default_runcard(proc_dir=process_dir), run_card_new='run_card.dat', xqcut=0, nevts=nevents, rand_seed=runArgs.randomSeed, beamEnergy=beamEnergy, extras=extras)

print_cards()

generate(run_card_loc='run_card.dat',param_card_loc='param_card_new.dat',mode=mode,proc_dir=process_dir,run_name=runName)

outputDS=arrange_output(run_name=runName,proc_dir=process_dir,outputDS=runName+'._00001.events.tar.gz')
runArgs.inputGeneratorFile=outputDS

#--------------------------------------------------------------
# Herwig 7 (H7UE) showering setup
#--------------------------------------------------------------
# initialize Herwig7 generator configuration for showering of LHE files
include("MC15JobOptions/Herwig7_LHEF.py")

# configure Herwig7
Herwig7Config.add_commands("set /Herwig/Partons/RemnantDecayer:AllowTop Yes")
Herwig7Config.me_pdf_commands(order="LO", name="NNPDF23_lo_as_0130_qed")
Herwig7Config.tune_commands()
Herwig7Config.lhef_mg5amc_commands(lhe_filename=runArgs.inputGeneratorFile, me_pdf_order="LO")
evgenConfig.tune="H7.1-Default"

# add EvtGen
include("MC15JobOptions/Herwig71_EvtGen.py")

#HW7 settings and Higgs BR
Herwig7Config.add_commands ("""
#set /Herwig/Shower/KinematicsReconstructor:ReconstructionOption General
#set /Herwig/Shower/KinematicsReconstructor:InitialInitialBoostOption LongTransBoost
set /Herwig/Shower/LtoLGammaSudakov:pTmin 0.000001
set /Herwig/Shower/QtoGammaQSudakov:Alpha /Herwig/Shower/AlphaQED
cd /Herwig/Particles
create /ThePEG/ParticleData XH
setup XH 1560  XH MASS WIDTH MAX_WIDTH 0 0 0 1 0
#set /Herwig/EventHandlers/LHEReader:Decayer /Herwig/Decays/Mambo
do /Herwig/Particles/h0:SelectDecayModes h0->gamma,gamma; h0->b,bbar;
do /Herwig/Particles/h0:PrintDecayModes
set /Herwig/Particles/h0/h0->gamma,gamma;:BranchingRatio  0.5
set /Herwig/Particles/h0/h0->b,bbar;:BranchingRatio  0.5
""")

# run Herwig7
Herwig7Config.run()


#---------------------------------------------------------------------------------------------------
# EVGEN Configuration
#---------------------------------------------------------------------------------------------------
evgenConfig.generators = ["MadGraph", "Herwig7", "EvtGen"]
if (runArgs.runNumber == 451892):
    evgenConfig.description = "Di-Higgs production through 270 GeV Heavy Higgs resonance which decays to bbyy."
    evgenConfig.keywords = ["hh","BSM", "BSMHiggs", "resonance", "diphoton", "bottom"]

evgenConfig.contact = ['Florian Beisiegel <florian.beisiegel@cern.ch>']
evgenConfig.inputfilecheck = runName
#runArgs.inputGeneratorFile=runName+'._00001.events.tar.gz'

#---------------------------------------------------------------------------------------------------
# Filter for bbyy
#---------------------------------------------------------------------------------------------------
from GeneratorFilters.GeneratorFiltersConf import ParentChildFilter
filtSeq += ParentChildFilter("HbbFilter", PDGParent = [25], PDGChild = [5])
filtSeq += ParentChildFilter("HyyFilter", PDGParent = [25], PDGChild = [22])
filtSeq.Expression = "HbbFilter and HyyFilter"