import MadGraphControl.MadGraphUtils
from MadGraphControl.MadGraphUtils import *
from MadGraphControl.MadGraphParamHelpers import set_top_params

#Job bookkeping infos
evgenConfig.description = 'aMcAtNlo ttW+0,1j@NLO FxFx'
evgenConfig.contact = ["marcos.miralles.lopez@cern.ch"]
evgenConfig.keywords+=['ttW','jets']
#evgenConfig.inputconfcheck= '/afs/cern.ch/work/m/mmiralle/ATHENA/QT/MG5_ttX_FxFx_myGridpack/550001.mg5_aMC_Py8.ttW_FxFx_01jNLO/550001'

# General settings
evgenConfig.nEventsPerJob = 10000
nevents = runArgs.maxEvents*1.1 if runArgs.maxEvents>0 else 1.1*evgenConfig.nEventsPerJob
if is_gen_from_gridpack():
	nevents=int(2.5*runArgs.maxEvents)

#Madgraph run card and shower settings
# Shower/merging settings
maxjetflavor=5
parton_shower='PYTHIA8'
nJetMax=1
qCut=30.

gridpack_mode=True
# gridpack_dir='madevent/'

# if not is_gen_from_gridpack():
#     process = """
#     import model loop_sm-no_b_mass
#     define p = g u c d s b u~ c~ d~ s~ b~
#     define j = g u c d s b u~ c~ d~ s~ b~
#     define w = w+ w-
#     define l+ = e+ mu+ ta+
#     define l- = e- mu- ta-
#     generate p p > t t~ w [QCD] @0
#     add process p p > t t~ w j [QCD] @1
#     output -f"""

#     process_dir = str(new_process(process))
# else:
#     process_dir = str(MADGRAPH_GRIDPACK_LOCATION)

process = """
    import model loop_sm-no_b_mass
    define p = g u c d s b u~ c~ d~ s~ b~
    define j = g u c d s b u~ c~ d~ s~ b~
    define w = w+ w-
    define l+ = e+ mu+ ta+
    define l- = e- mu- ta-
    generate p p > t t~ w [QCD] @0
    add process p p > t t~ w j [QCD] @1
    output -f"""

process_dir = str(new_process(process))

MadGraphControl.MadGraphUtils.MADGRAPH_PDFSETTING={
    'central_pdf':303600, # the lhapf id of the central pdf, see https://lhapdf.hepforge.org/pdfsets
    'pdf_variations':[303600], # list of pdfs ids for which all variations (error sets) will be included as weights
    'alternative_pdfs':[325100,13000,25300,14000,14100,14200,61200,42560], 
    'scale_variations':[0.5,1,2], # variations of muR and muF wrt the central scale, all combinations of muF and muR will be evaluated
}

#Fetch default run_card.dat and set parameters
settings = {
            'maxjetflavor'  : int(maxjetflavor),
            'parton_shower' : parton_shower,
            'nevents'       : int(nevents),
            'ickkw'         : 3,
            'jetradius'     : 1.0,
            'ptj'           : 10,
            'etaj'          : 10,
            'bwcutoff'      : 50
        }

modify_run_card(process_dir=process_dir,runArgs=runArgs,settings=settings)

set_top_params(process_dir=process_dir,mTop=172.5,FourFS=False)

input_events=process_dir+'/Events/GridRun_'+str(runArgs.randomSeed)+'/events.lhe.gz'

madspin_card=process_dir+'/Cards/madspin_card.dat'
if os.access(madspin_card,os.R_OK):
	os.unlink(madspin_card)

mscard = open(madspin_card,'w')
mscard.write("""#************************************************************
#*                        MadSpin                           *
#*                                                          *
#*    P. Artoisenet, R. Frederix, R. Rietkerk, O. Mattelaer *
#*                                                          *
#*    Part of the MadGraph5_aMC@NLO Framework:              *
#*    The MadGraph5_aMC@NLO Development Team - Find us at   *
#*    https://server06.fynu.ucl.ac.be/projects/madgraph     *
#*                                                          *
#************************************************************
#Some options (uncomment to apply)
#
# set seed 1
# set Nevents_for_max_weigth 75 # number of events for the estimate of the max. weight
 set max_weight_ps_point 400  # number of PS to estimate the maximum for each event
#
set seed %i
set BW_cut 50 # cut on how far the particle can be off-shell
# specify the decay for the final state particles
decay t > w+ b, w+ > all all
decay t~ > w- b~, w- > all all
decay w+ > all all
decay w- > all all
decay z > all all
# running the actual code
launch"""%runArgs.randomSeed)
mscard.close()    


generate(process_dir=process_dir,runArgs=runArgs,grid_pack=gridpack_mode)
arrange_output(process_dir=process_dir,runArgs=runArgs,lhe_version=3,saveProcDir=True)

# if not is_gen_from_gridpack():
#     # Helper for resetting process number
#     check_reset_proc_number(opts)
check_reset_proc_number(opts)

#### Shower: Py8 with A14 Tune, with modifications to make it simil-NLO
include("Pythia8_i/Pythia8_A14_NNPDF23LO_EvtGen_Common.py") 
include("Pythia8_i/Pythia8_aMcAtNlo.py")

#simil-NLO parameters
genSeq.Pythia8.Commands += ["JetMatching:merge          = on", 
                            "SpaceShower:alphaSuseCMW    = on",              
                            "SpaceShower:alphaSorder     = 2",               
                            "TimeShower:alphaSuseCMW     = on",              
                            "TimeShower:alphaSorder      = 2",               
                            "SpaceShower:alphaSvalue    = 0.118",
                            "TimeShower:alphaSvalue     = 0.118" 
                             ] 

# FxFx Matching settings, according to authors prescriptions (NB: it changes tune pars)
PYTHIA8_nJetMax=nJetMax
PYTHIA8_qCut=qCut
print "PYTHIA8_nJetMax = %i"%PYTHIA8_nJetMax
print "PYTHIA8_qCut = %i"%PYTHIA8_qCut

genSeq.Pythia8.Commands += ["JetMatching:merge          = on", 
                            "JetMatching:scheme         = 1",  
                            "JetMatching:setMad         = off",
                            "SpaceShower:rapidityOrder  = off",
                            "SpaceShower:pTmaxFudge     = 1.0",   
                            "JetMatching:qCut            = %f"%PYTHIA8_qCut,
                            "JetMatching:coneRadius      = 1.0",
                            "JetMatching:etaJetMax       = 10.0",
                            "JetMatching:doFxFx          = on",
                            "JetMatching:qCutME          = 10.0",
                            "JetMatching:nJetMax         = %i"%PYTHIA8_nJetMax,
                            'JetMatching:jetAlgorithm = 2', #explicit setting of kt-merging for FxFx (also imposed by Py8-FxFx inteface)
                            'JetMatching:slowJetPower = 1', #explicit setting of kt-merging for FxFx (also imposed by Py8-FxFx inteface)
                            'JetMatching:nQmatch = 5', #4 corresponds to 4-flavour scheme (no matching of b-quarks), 5 for 5-flavour scheme
                            "JetMatching:eTjetMin = %f"%PYTHIA8_qCut #This is 20 in the Pythia default, it should be <= qCut 
                             ]

genSeq.Pythia8.UserHooks = ['JetMatchingMadgraph'] 
genSeq.Pythia8.FxFxXS = True