#! /usr/bin/env python

################################################################################
#
# Script intended to reweight events generated using aMC@LO and
# already reweighted to include the virtual corrections (without an
# overall alpha_s/2pi) with the corresponding Resummation
# contributions. No secondary unweighting is performed.
#
################################################################################

import os
import re
import subprocess
import sys
import time
import logging
import glob
import copy
import math
pjoin = os.path.join
proc_path = os.path.abspath(pjoin(os.path.dirname(os.path.realpath( __file__ )),os.pardir,os.pardir))
MGRootPath = pjoin(proc_path,os.pardir)
sys.path.append(MGRootPath)
sys.path.append(pjoin(proc_path,'bin'))
import madgraph.various.lhe_parser as lhe_parser
import models.import_ufo as import_ufo
import banner
import madgraph.various.misc as misc
from madgraph.iolibs.files import cp
import madgraph.various.process_checks as process_checks

logging.basicConfig(level=logging.INFO)
logging.getLogger('madgraph.model').setLevel(logging.CRITICAL)

# ==============================================
# Helper functions
# ==============================================

def setup_channel(channelPath):
    """ Compiles the Resummation fortran steering 
    code and returns the corresponding process."""

    logging.info('Setting up resummation runner in\n   %s'%channelPath)

    try:
        misc.compile(cwd=channelPath, mode='fortran', job_specs = False)
    except:
        logging.error("Could not compile resummation runner in %s/%s"%(channelPath))
        raise
    
    logging.info('  ... Done.')
    
    return subprocess.Popen([os.path.join(channelPath,'resum_driver')], 
          stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, 
                                                                   cwd=channelPath)

def kill_processes(encountered_channels):
    """ Kills the MadLoop processes open """

    for proc in encountered_channels.values():
        try:
            proc[0].stdin.write('y\n')
        except:
            proc[0].kill()

def get_value(Runner, input_str):
    """ This version of get_me_value is simplified for the purpose of this
    class. No compilation is necessary. The CT mode can be specified."""
    Runner.stdin.write(input_str)
    output = Runner.stdout.readline()  
    return output

# ==============================================
# Start program
# ==============================================

# Specify here the path of the file to reweight
if len(sys.argv)>0:
  eventFile=sys.argv[1]
else:
    logging.error("Could not process the arguments. The arguments are: ./resum_reweighter.py <eventFilePath> <outputEvtName>")
    exit()

outSpecified=False
evtFileName = '.'.join(os.path.basename(eventFile).split('.')[:-1])
evtFileExt  = os.path.basename(eventFile).split('.')[-1]
try:
  outPath=pjoin(os.path.dirname(os.path.realpath(eventFile)),sys.argv[2])
  outSpecified=True
except:
  outPath=pjoin(os.path.dirname(os.path.realpath(eventFile)),'%s_rwgt.%s'%(evtFileName,evtFileExt))

i=0
while os.path.exists(outPath):
    i=i+1    
    if outSpecified and i==1:
        logging.warning("The chosen filename %s for the output already exists."%os.path.basename(outPath))
    outPath=pjoin(os.path.dirname(os.path.realpath(eventFile)),
                  '%s_rwgt_%i.%s'%(evtFileName,i,evtFileExt))

try:
    evtBanner = banner.Banner(eventFile)    
    evtFile = lhe_parser.EventFile(eventFile)
except:
    logging.error("Could not read event file %s."%eventFile)

# Detect the total number of event in an efficient way
n_evts = int(subprocess.Popen('grep -rin "<event>" %s | wc -l'%os.path.realpath(eventFile),
                        stdout=subprocess.PIPE, shell=True).communicate()[0])

logging.info("Writing out the reweighted event file on\n   %s"%outPath)
# Write out the output event file
outputEvtFile = open(outPath,'w')
outputEvtFile.write(evtFile.banner)
# List the channel present in the process output
channel_list = [ os.path.basename(chan) for chan in \
                 misc.glob('P*', pjoin(proc_path, 'SubProcesses')) ]

# Now scan over events
# For each encountered channel, store the corresponding process runner
encountered_channels=dict([])

ebeam1=float( evtBanner.get('run_card','ebeam1'))
ebeam2=float( evtBanner.get('run_card','ebeam2'))

t_before = time.time()
i_evt=0
wgt_summed = 0.0
squared_wgt_summed = 0.0
running_spread = 0.0
max_wgt = 0.0
n_percent_monitor=1
for event in evtFile:
    i_evt = i_evt+1
    t_evt_start = time.time()    
    # The channel Folder
    chanFolder=""
    EvtChannelName="VetoPrefactors"
    # The Channel name to setup, if needed
    ChannelNameToSetup=""

    # =======================================
    # Setup the runner if not already present
    # =======================================

    try:
        runner=encountered_channels[EvtChannelName][0]
    except KeyError:
        logging.info('Found new channel %s'%EvtChannelName)
        virtFolderPath=pjoin(os.path.dirname(os.path.realpath(eventFile)),EvtChannelName)
        try:
            runner=setup_channel(virtFolderPath)
        except:
            logging.error("Could not setup resum runner %s"%virtFolderPath)
            kill_processes(encountered_channels)
            exit()
    
            # Save the runners just created for ChannelNameToSetup and possibly related to
            # EvtChannelName as well.
        encountered_channels[EvtChannelName]=(runner,virtFolderPath)

    # =======================================
    # End of virtual runner setup
    # =======================================

    momenta=[[p.E,p.px,p.py,p.pz] for p in event if p.status==-1]
    pdg=[int(p.pid) for p in event if p.status==-1]

    event.parse_reweight()
    born_wgt=event.reweight_data['1001']

# PDG codes of the incoming particles
    f1=str(pdg[0])
    f2=str(pdg[1])
# Bjorken x's of the particles
    x1=str(momenta[0][0]/ebeam1)
    x2=str(momenta[1][0]/ebeam2) 
# Q is set to \sqrt(\hat{s})
    Q =str(2*math.sqrt(momenta[0][0]*momenta[1][0]))
# renormalisation scale used to generate the events
    muMad=str(event.scale)

    input_string=f1+' '+f2+' '+x1+' '+x2+' '+Q+' '+muMad+'\n'
    trial  = 0
    max_trial = 50
    worked = False
    new_wgt = 0.0

    output_string = get_value(runner,input_string)
    mode=float(output_string.split()[0])
    BCorr0=float(output_string.split()[1])
    BCorr1=float(output_string.split()[2])
    BCorrm0=float(output_string.split()[3])
    BCorrm1=float(output_string.split()[4])
    Efull=float(output_string.split()[5])
    EfullNLL=float(output_string.split()[6])
    Hcomp_fact=float(output_string.split()[7])
    Hcomp_factm=float(output_string.split()[8])
    alphah=float(output_string.split()[9])
    alpham=float(output_string.split()[10])
    EfO=float(output_string.split()[11])
    EfONLL=float(output_string.split()[12])
    alpha=float(output_string.split()[13])
    if mode==1:
        new_wgt= (1+alpha*BCorr1+(alphah/(2.0*math.pi))*(event.wgt/born_wgt+Hcomp_fact))*BCorr0*Efull*born_wgt
    if mode==2:
        new_wgt= BCorr0*EfullNLL*born_wgt
    if mode==3:
        new_wgt= (1+BCorrm1+EfO+(alpham/(2.0*math.pi))*(event.wgt/born_wgt+Hcomp_factm))*BCorrm0*born_wgt
    if mode==4:
        new_wgt= BCorrm0*(1+EfONLL)*born_wgt

    worked = True

    if i_evt<=5:
        logging.info('Event %i processed in %.2E seconds.'%(i_evt,time.time()-t_evt_start))
    if i_evt==5:
        logging.info('Further output of timing event by event suppressed.')

    event.wgt=new_wgt
    wgt_summed = wgt_summed + new_wgt
    squared_wgt_summed = squared_wgt_summed + new_wgt**2
    max_wgt = max(max_wgt,abs(new_wgt))
    if (i_evt/float(n_evts))>n_percent_monitor*0.01:
        logging.info('%i%% of %i events processed in %.2E seconds. Cross section: %.4E pb'\
                     %(n_percent_monitor,n_evts,(time.time()-t_before),wgt_summed/i_evt))
        n_percent_monitor=n_percent_monitor+1
    # Only start computing the cumulative spread after the first 10% of events.
    if n_percent_monitor>1:
        running_spread = running_spread + (new_wgt - float(wgt_summed/i_evt))**2
    outputEvtFile.write(str(event))

# Clean up things
outputEvtFile.close()
kill_processes(encountered_channels)

# Finish
logging.info("Successfully reweighted %i events in %.2E seconds in file\n  %s."\
    %(n_evts,(time.time()-t_before),outPath))

# Give statistics on the weight
logging.info("Resummation weight statistics,\n  avg       : %.2E\n  cum. var. : %.2E\n  max       : %.2E"\
    %((wgt_summed/n_evts),math.sqrt(running_spread/n_evts),max_wgt))

# Give statistics on the weight
logging.info("Resummation cross-section obtained\n  %.6E  +/- %.2E"\
    %((wgt_summed/n_evts),math.sqrt(squared_wgt_summed)/n_evts))

29407d8823106482b75ed2050b32c3274580b83c
<mask token>


def setup_channel(channelPath):
    """ Compiles the Resummation fortran steering 
    code and returns the corresponding process."""
    logging.info('Setting up resummation runner in\n   %s' % channelPath)
    try:
        misc.compile(cwd=channelPath, mode='fortran', job_specs=False)
    except:
        logging.error('Could not compile resummation runner in %s/%s' %
            channelPath)
        raise
    logging.info('  ... Done.')
    return subprocess.Popen([os.path.join(channelPath, 'resum_driver')],
        stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.
        PIPE, cwd=channelPath)


def kill_processes(encountered_channels):
    """ Kills the MadLoop processes open """
    for proc in encountered_channels.values():
        try:
            proc[0].stdin.write('y\n')
        except:
            proc[0].kill()


def get_value(Runner, input_str):
    """ This version of get_me_value is simplified for the purpose of this
    class. No compilation is necessary. The CT mode can be specified."""
    Runner.stdin.write(input_str)
    output = Runner.stdout.readline()
    return output


<mask token>

<mask token>
sys.path.append(MGRootPath)
sys.path.append(pjoin(proc_path, 'bin'))
<mask token>
logging.basicConfig(level=logging.INFO)
logging.getLogger('madgraph.model').setLevel(logging.CRITICAL)


def setup_channel(channelPath):
    """ Compiles the Resummation fortran steering 
    code and returns the corresponding process."""
    logging.info('Setting up resummation runner in\n   %s' % channelPath)
    try:
        misc.compile(cwd=channelPath, mode='fortran', job_specs=False)
    except:
        logging.error('Could not compile resummation runner in %s/%s' %
            channelPath)
        raise
    logging.info('  ... Done.')
    return subprocess.Popen([os.path.join(channelPath, 'resum_driver')],
        stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.
        PIPE, cwd=channelPath)


def kill_processes(encountered_channels):
    """ Kills the MadLoop processes open """
    for proc in encountered_channels.values():
        try:
            proc[0].stdin.write('y\n')
        except:
            proc[0].kill()


def get_value(Runner, input_str):
    """ This version of get_me_value is simplified for the purpose of this
    class. No compilation is necessary. The CT mode can be specified."""
    Runner.stdin.write(input_str)
    output = Runner.stdout.readline()
    return output


if len(sys.argv) > 0:
    eventFile = sys.argv[1]
else:
    logging.error(
        'Could not process the arguments. The arguments are: ./resum_reweighter.py <eventFilePath> <outputEvtName>'
        )
    exit()
<mask token>
try:
    outPath = pjoin(os.path.dirname(os.path.realpath(eventFile)), sys.argv[2])
    outSpecified = True
except:
    outPath = pjoin(os.path.dirname(os.path.realpath(eventFile)), 
        '%s_rwgt.%s' % (evtFileName, evtFileExt))
<mask token>
while os.path.exists(outPath):
    i = i + 1
    if outSpecified and i == 1:
        logging.warning(
            'The chosen filename %s for the output already exists.' % os.
            path.basename(outPath))
    outPath = pjoin(os.path.dirname(os.path.realpath(eventFile)), 
        '%s_rwgt_%i.%s' % (evtFileName, i, evtFileExt))
try:
    evtBanner = banner.Banner(eventFile)
    evtFile = lhe_parser.EventFile(eventFile)
except:
    logging.error('Could not read event file %s.' % eventFile)
<mask token>
logging.info("""Writing out the reweighted event file on
   %s""" % outPath)
<mask token>
outputEvtFile.write(evtFile.banner)
<mask token>
for event in evtFile:
    i_evt = i_evt + 1
    t_evt_start = time.time()
    chanFolder = ''
    EvtChannelName = 'VetoPrefactors'
    ChannelNameToSetup = ''
    try:
        runner = encountered_channels[EvtChannelName][0]
    except KeyError:
        logging.info('Found new channel %s' % EvtChannelName)
        virtFolderPath = pjoin(os.path.dirname(os.path.realpath(eventFile)),
            EvtChannelName)
        try:
            runner = setup_channel(virtFolderPath)
        except:
            logging.error('Could not setup resum runner %s' % virtFolderPath)
            kill_processes(encountered_channels)
            exit()
        encountered_channels[EvtChannelName] = runner, virtFolderPath
    momenta = [[p.E, p.px, p.py, p.pz] for p in event if p.status == -1]
    pdg = [int(p.pid) for p in event if p.status == -1]
    event.parse_reweight()
    born_wgt = event.reweight_data['1001']
    f1 = str(pdg[0])
    f2 = str(pdg[1])
    x1 = str(momenta[0][0] / ebeam1)
    x2 = str(momenta[1][0] / ebeam2)
    Q = str(2 * math.sqrt(momenta[0][0] * momenta[1][0]))
    muMad = str(event.scale)
    input_string = (f1 + ' ' + f2 + ' ' + x1 + ' ' + x2 + ' ' + Q + ' ' +
        muMad + '\n')
    trial = 0
    max_trial = 50
    worked = False
    new_wgt = 0.0
    output_string = get_value(runner, input_string)
    mode = float(output_string.split()[0])
    BCorr0 = float(output_string.split()[1])
    BCorr1 = float(output_string.split()[2])
    BCorrm0 = float(output_string.split()[3])
    BCorrm1 = float(output_string.split()[4])
    Efull = float(output_string.split()[5])
    EfullNLL = float(output_string.split()[6])
    Hcomp_fact = float(output_string.split()[7])
    Hcomp_factm = float(output_string.split()[8])
    alphah = float(output_string.split()[9])
    alpham = float(output_string.split()[10])
    EfO = float(output_string.split()[11])
    EfONLL = float(output_string.split()[12])
    alpha = float(output_string.split()[13])
    if mode == 1:
        new_wgt = (1 + alpha * BCorr1 + alphah / (2.0 * math.pi) * (event.
            wgt / born_wgt + Hcomp_fact)) * BCorr0 * Efull * born_wgt
    if mode == 2:
        new_wgt = BCorr0 * EfullNLL * born_wgt
    if mode == 3:
        new_wgt = (1 + BCorrm1 + EfO + alpham / (2.0 * math.pi) * (event.
            wgt / born_wgt + Hcomp_factm)) * BCorrm0 * born_wgt
    if mode == 4:
        new_wgt = BCorrm0 * (1 + EfONLL) * born_wgt
    worked = True
    if i_evt <= 5:
        logging.info('Event %i processed in %.2E seconds.' % (i_evt, time.
            time() - t_evt_start))
    if i_evt == 5:
        logging.info('Further output of timing event by event suppressed.')
    event.wgt = new_wgt
    wgt_summed = wgt_summed + new_wgt
    squared_wgt_summed = squared_wgt_summed + new_wgt ** 2
    max_wgt = max(max_wgt, abs(new_wgt))
    if i_evt / float(n_evts) > n_percent_monitor * 0.01:
        logging.info(
            '%i%% of %i events processed in %.2E seconds. Cross section: %.4E pb'
             % (n_percent_monitor, n_evts, time.time() - t_before, 
            wgt_summed / i_evt))
        n_percent_monitor = n_percent_monitor + 1
    if n_percent_monitor > 1:
        running_spread = running_spread + (new_wgt - float(wgt_summed / i_evt)
            ) ** 2
    outputEvtFile.write(str(event))
outputEvtFile.close()
kill_processes(encountered_channels)
logging.info(
    """Successfully reweighted %i events in %.2E seconds in file
  %s.""" %
    (n_evts, time.time() - t_before, outPath))
logging.info(
    """Resummation weight statistics,
  avg       : %.2E
  cum. var. : %.2E
  max       : %.2E"""
     % (wgt_summed / n_evts, math.sqrt(running_spread / n_evts), max_wgt))
logging.info("""Resummation cross-section obtained
  %.6E  +/- %.2E""" % (
    wgt_summed / n_evts, math.sqrt(squared_wgt_summed) / n_evts))

<mask token>
pjoin = os.path.join
proc_path = os.path.abspath(pjoin(os.path.dirname(os.path.realpath(__file__
    )), os.pardir, os.pardir))
MGRootPath = pjoin(proc_path, os.pardir)
sys.path.append(MGRootPath)
sys.path.append(pjoin(proc_path, 'bin'))
<mask token>
logging.basicConfig(level=logging.INFO)
logging.getLogger('madgraph.model').setLevel(logging.CRITICAL)


def setup_channel(channelPath):
    """ Compiles the Resummation fortran steering 
    code and returns the corresponding process."""
    logging.info('Setting up resummation runner in\n   %s' % channelPath)
    try:
        misc.compile(cwd=channelPath, mode='fortran', job_specs=False)
    except:
        logging.error('Could not compile resummation runner in %s/%s' %
            channelPath)
        raise
    logging.info('  ... Done.')
    return subprocess.Popen([os.path.join(channelPath, 'resum_driver')],
        stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.
        PIPE, cwd=channelPath)


def kill_processes(encountered_channels):
    """ Kills the MadLoop processes open """
    for proc in encountered_channels.values():
        try:
            proc[0].stdin.write('y\n')
        except:
            proc[0].kill()


def get_value(Runner, input_str):
    """ This version of get_me_value is simplified for the purpose of this
    class. No compilation is necessary. The CT mode can be specified."""
    Runner.stdin.write(input_str)
    output = Runner.stdout.readline()
    return output


if len(sys.argv) > 0:
    eventFile = sys.argv[1]
else:
    logging.error(
        'Could not process the arguments. The arguments are: ./resum_reweighter.py <eventFilePath> <outputEvtName>'
        )
    exit()
outSpecified = False
evtFileName = '.'.join(os.path.basename(eventFile).split('.')[:-1])
evtFileExt = os.path.basename(eventFile).split('.')[-1]
try:
    outPath = pjoin(os.path.dirname(os.path.realpath(eventFile)), sys.argv[2])
    outSpecified = True
except:
    outPath = pjoin(os.path.dirname(os.path.realpath(eventFile)), 
        '%s_rwgt.%s' % (evtFileName, evtFileExt))
i = 0
while os.path.exists(outPath):
    i = i + 1
    if outSpecified and i == 1:
        logging.warning(
            'The chosen filename %s for the output already exists.' % os.
            path.basename(outPath))
    outPath = pjoin(os.path.dirname(os.path.realpath(eventFile)), 
        '%s_rwgt_%i.%s' % (evtFileName, i, evtFileExt))
try:
    evtBanner = banner.Banner(eventFile)
    evtFile = lhe_parser.EventFile(eventFile)
except:
    logging.error('Could not read event file %s.' % eventFile)
n_evts = int(subprocess.Popen('grep -rin "<event>" %s | wc -l' % os.path.
    realpath(eventFile), stdout=subprocess.PIPE, shell=True).communicate()[0])
logging.info("""Writing out the reweighted event file on
   %s""" % outPath)
outputEvtFile = open(outPath, 'w')
outputEvtFile.write(evtFile.banner)
channel_list = [os.path.basename(chan) for chan in misc.glob('P*', pjoin(
    proc_path, 'SubProcesses'))]
encountered_channels = dict([])
ebeam1 = float(evtBanner.get('run_card', 'ebeam1'))
ebeam2 = float(evtBanner.get('run_card', 'ebeam2'))
t_before = time.time()
i_evt = 0
wgt_summed = 0.0
squared_wgt_summed = 0.0
running_spread = 0.0
max_wgt = 0.0
n_percent_monitor = 1
for event in evtFile:
    i_evt = i_evt + 1
    t_evt_start = time.time()
    chanFolder = ''
    EvtChannelName = 'VetoPrefactors'
    ChannelNameToSetup = ''
    try:
        runner = encountered_channels[EvtChannelName][0]
    except KeyError:
        logging.info('Found new channel %s' % EvtChannelName)
        virtFolderPath = pjoin(os.path.dirname(os.path.realpath(eventFile)),
            EvtChannelName)
        try:
            runner = setup_channel(virtFolderPath)
        except:
            logging.error('Could not setup resum runner %s' % virtFolderPath)
            kill_processes(encountered_channels)
            exit()
        encountered_channels[EvtChannelName] = runner, virtFolderPath
    momenta = [[p.E, p.px, p.py, p.pz] for p in event if p.status == -1]
    pdg = [int(p.pid) for p in event if p.status == -1]
    event.parse_reweight()
    born_wgt = event.reweight_data['1001']
    f1 = str(pdg[0])
    f2 = str(pdg[1])
    x1 = str(momenta[0][0] / ebeam1)
    x2 = str(momenta[1][0] / ebeam2)
    Q = str(2 * math.sqrt(momenta[0][0] * momenta[1][0]))
    muMad = str(event.scale)
    input_string = (f1 + ' ' + f2 + ' ' + x1 + ' ' + x2 + ' ' + Q + ' ' +
        muMad + '\n')
    trial = 0
    max_trial = 50
    worked = False
    new_wgt = 0.0
    output_string = get_value(runner, input_string)
    mode = float(output_string.split()[0])
    BCorr0 = float(output_string.split()[1])
    BCorr1 = float(output_string.split()[2])
    BCorrm0 = float(output_string.split()[3])
    BCorrm1 = float(output_string.split()[4])
    Efull = float(output_string.split()[5])
    EfullNLL = float(output_string.split()[6])
    Hcomp_fact = float(output_string.split()[7])
    Hcomp_factm = float(output_string.split()[8])
    alphah = float(output_string.split()[9])
    alpham = float(output_string.split()[10])
    EfO = float(output_string.split()[11])
    EfONLL = float(output_string.split()[12])
    alpha = float(output_string.split()[13])
    if mode == 1:
        new_wgt = (1 + alpha * BCorr1 + alphah / (2.0 * math.pi) * (event.
            wgt / born_wgt + Hcomp_fact)) * BCorr0 * Efull * born_wgt
    if mode == 2:
        new_wgt = BCorr0 * EfullNLL * born_wgt
    if mode == 3:
        new_wgt = (1 + BCorrm1 + EfO + alpham / (2.0 * math.pi) * (event.
            wgt / born_wgt + Hcomp_factm)) * BCorrm0 * born_wgt
    if mode == 4:
        new_wgt = BCorrm0 * (1 + EfONLL) * born_wgt
    worked = True
    if i_evt <= 5:
        logging.info('Event %i processed in %.2E seconds.' % (i_evt, time.
            time() - t_evt_start))
    if i_evt == 5:
        logging.info('Further output of timing event by event suppressed.')
    event.wgt = new_wgt
    wgt_summed = wgt_summed + new_wgt
    squared_wgt_summed = squared_wgt_summed + new_wgt ** 2
    max_wgt = max(max_wgt, abs(new_wgt))
    if i_evt / float(n_evts) > n_percent_monitor * 0.01:
        logging.info(
            '%i%% of %i events processed in %.2E seconds. Cross section: %.4E pb'
             % (n_percent_monitor, n_evts, time.time() - t_before, 
            wgt_summed / i_evt))
        n_percent_monitor = n_percent_monitor + 1
    if n_percent_monitor > 1:
        running_spread = running_spread + (new_wgt - float(wgt_summed / i_evt)
            ) ** 2
    outputEvtFile.write(str(event))
outputEvtFile.close()
kill_processes(encountered_channels)
logging.info(
    """Successfully reweighted %i events in %.2E seconds in file
  %s.""" %
    (n_evts, time.time() - t_before, outPath))
logging.info(
    """Resummation weight statistics,
  avg       : %.2E
  cum. var. : %.2E
  max       : %.2E"""
     % (wgt_summed / n_evts, math.sqrt(running_spread / n_evts), max_wgt))
logging.info("""Resummation cross-section obtained
  %.6E  +/- %.2E""" % (
    wgt_summed / n_evts, math.sqrt(squared_wgt_summed) / n_evts))

import os
import re
import subprocess
import sys
import time
import logging
import glob
import copy
import math
pjoin = os.path.join
proc_path = os.path.abspath(pjoin(os.path.dirname(os.path.realpath(__file__
    )), os.pardir, os.pardir))
MGRootPath = pjoin(proc_path, os.pardir)
sys.path.append(MGRootPath)
sys.path.append(pjoin(proc_path, 'bin'))
import madgraph.various.lhe_parser as lhe_parser
import models.import_ufo as import_ufo
import banner
import madgraph.various.misc as misc
from madgraph.iolibs.files import cp
import madgraph.various.process_checks as process_checks
logging.basicConfig(level=logging.INFO)
logging.getLogger('madgraph.model').setLevel(logging.CRITICAL)


def setup_channel(channelPath):
    """ Compiles the Resummation fortran steering 
    code and returns the corresponding process."""
    logging.info('Setting up resummation runner in\n   %s' % channelPath)
    try:
        misc.compile(cwd=channelPath, mode='fortran', job_specs=False)
    except:
        logging.error('Could not compile resummation runner in %s/%s' %
            channelPath)
        raise
    logging.info('  ... Done.')
    return subprocess.Popen([os.path.join(channelPath, 'resum_driver')],
        stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.
        PIPE, cwd=channelPath)


def kill_processes(encountered_channels):
    """ Kills the MadLoop processes open """
    for proc in encountered_channels.values():
        try:
            proc[0].stdin.write('y\n')
        except:
            proc[0].kill()


def get_value(Runner, input_str):
    """ This version of get_me_value is simplified for the purpose of this
    class. No compilation is necessary. The CT mode can be specified."""
    Runner.stdin.write(input_str)
    output = Runner.stdout.readline()
    return output


if len(sys.argv) > 0:
    eventFile = sys.argv[1]
else:
    logging.error(
        'Could not process the arguments. The arguments are: ./resum_reweighter.py <eventFilePath> <outputEvtName>'
        )
    exit()
outSpecified = False
evtFileName = '.'.join(os.path.basename(eventFile).split('.')[:-1])
evtFileExt = os.path.basename(eventFile).split('.')[-1]
try:
    outPath = pjoin(os.path.dirname(os.path.realpath(eventFile)), sys.argv[2])
    outSpecified = True
except:
    outPath = pjoin(os.path.dirname(os.path.realpath(eventFile)), 
        '%s_rwgt.%s' % (evtFileName, evtFileExt))
i = 0
while os.path.exists(outPath):
    i = i + 1
    if outSpecified and i == 1:
        logging.warning(
            'The chosen filename %s for the output already exists.' % os.
            path.basename(outPath))
    outPath = pjoin(os.path.dirname(os.path.realpath(eventFile)), 
        '%s_rwgt_%i.%s' % (evtFileName, i, evtFileExt))
try:
    evtBanner = banner.Banner(eventFile)
    evtFile = lhe_parser.EventFile(eventFile)
except:
    logging.error('Could not read event file %s.' % eventFile)
n_evts = int(subprocess.Popen('grep -rin "<event>" %s | wc -l' % os.path.
    realpath(eventFile), stdout=subprocess.PIPE, shell=True).communicate()[0])
logging.info("""Writing out the reweighted event file on
   %s""" % outPath)
outputEvtFile = open(outPath, 'w')
outputEvtFile.write(evtFile.banner)
channel_list = [os.path.basename(chan) for chan in misc.glob('P*', pjoin(
    proc_path, 'SubProcesses'))]
encountered_channels = dict([])
ebeam1 = float(evtBanner.get('run_card', 'ebeam1'))
ebeam2 = float(evtBanner.get('run_card', 'ebeam2'))
t_before = time.time()
i_evt = 0
wgt_summed = 0.0
squared_wgt_summed = 0.0
running_spread = 0.0
max_wgt = 0.0
n_percent_monitor = 1
for event in evtFile:
    i_evt = i_evt + 1
    t_evt_start = time.time()
    chanFolder = ''
    EvtChannelName = 'VetoPrefactors'
    ChannelNameToSetup = ''
    try:
        runner = encountered_channels[EvtChannelName][0]
    except KeyError:
        logging.info('Found new channel %s' % EvtChannelName)
        virtFolderPath = pjoin(os.path.dirname(os.path.realpath(eventFile)),
            EvtChannelName)
        try:
            runner = setup_channel(virtFolderPath)
        except:
            logging.error('Could not setup resum runner %s' % virtFolderPath)
            kill_processes(encountered_channels)
            exit()
        encountered_channels[EvtChannelName] = runner, virtFolderPath
    momenta = [[p.E, p.px, p.py, p.pz] for p in event if p.status == -1]
    pdg = [int(p.pid) for p in event if p.status == -1]
    event.parse_reweight()
    born_wgt = event.reweight_data['1001']
    f1 = str(pdg[0])
    f2 = str(pdg[1])
    x1 = str(momenta[0][0] / ebeam1)
    x2 = str(momenta[1][0] / ebeam2)
    Q = str(2 * math.sqrt(momenta[0][0] * momenta[1][0]))
    muMad = str(event.scale)
    input_string = (f1 + ' ' + f2 + ' ' + x1 + ' ' + x2 + ' ' + Q + ' ' +
        muMad + '\n')
    trial = 0
    max_trial = 50
    worked = False
    new_wgt = 0.0
    output_string = get_value(runner, input_string)
    mode = float(output_string.split()[0])
    BCorr0 = float(output_string.split()[1])
    BCorr1 = float(output_string.split()[2])
    BCorrm0 = float(output_string.split()[3])
    BCorrm1 = float(output_string.split()[4])
    Efull = float(output_string.split()[5])
    EfullNLL = float(output_string.split()[6])
    Hcomp_fact = float(output_string.split()[7])
    Hcomp_factm = float(output_string.split()[8])
    alphah = float(output_string.split()[9])
    alpham = float(output_string.split()[10])
    EfO = float(output_string.split()[11])
    EfONLL = float(output_string.split()[12])
    alpha = float(output_string.split()[13])
    if mode == 1:
        new_wgt = (1 + alpha * BCorr1 + alphah / (2.0 * math.pi) * (event.
            wgt / born_wgt + Hcomp_fact)) * BCorr0 * Efull * born_wgt
    if mode == 2:
        new_wgt = BCorr0 * EfullNLL * born_wgt
    if mode == 3:
        new_wgt = (1 + BCorrm1 + EfO + alpham / (2.0 * math.pi) * (event.
            wgt / born_wgt + Hcomp_factm)) * BCorrm0 * born_wgt
    if mode == 4:
        new_wgt = BCorrm0 * (1 + EfONLL) * born_wgt
    worked = True
    if i_evt <= 5:
        logging.info('Event %i processed in %.2E seconds.' % (i_evt, time.
            time() - t_evt_start))
    if i_evt == 5:
        logging.info('Further output of timing event by event suppressed.')
    event.wgt = new_wgt
    wgt_summed = wgt_summed + new_wgt
    squared_wgt_summed = squared_wgt_summed + new_wgt ** 2
    max_wgt = max(max_wgt, abs(new_wgt))
    if i_evt / float(n_evts) > n_percent_monitor * 0.01:
        logging.info(
            '%i%% of %i events processed in %.2E seconds. Cross section: %.4E pb'
             % (n_percent_monitor, n_evts, time.time() - t_before, 
            wgt_summed / i_evt))
        n_percent_monitor = n_percent_monitor + 1
    if n_percent_monitor > 1:
        running_spread = running_spread + (new_wgt - float(wgt_summed / i_evt)
            ) ** 2
    outputEvtFile.write(str(event))
outputEvtFile.close()
kill_processes(encountered_channels)
logging.info(
    """Successfully reweighted %i events in %.2E seconds in file
  %s.""" %
    (n_evts, time.time() - t_before, outPath))
logging.info(
    """Resummation weight statistics,
  avg       : %.2E
  cum. var. : %.2E
  max       : %.2E"""
     % (wgt_summed / n_evts, math.sqrt(running_spread / n_evts), max_wgt))
logging.info("""Resummation cross-section obtained
  %.6E  +/- %.2E""" % (
    wgt_summed / n_evts, math.sqrt(squared_wgt_summed) / n_evts))

#! /usr/bin/env python

################################################################################
#
# Script intended to reweight events generated using aMC@LO and
# already reweighted to include the virtual corrections (without an
# overall alpha_s/2pi) with the corresponding Resummation
# contributions. No secondary unweighting is performed.
#
################################################################################

import os
import re
import subprocess
import sys
import time
import logging
import glob
import copy
import math
pjoin = os.path.join
proc_path = os.path.abspath(pjoin(os.path.dirname(os.path.realpath( __file__ )),os.pardir,os.pardir))
MGRootPath = pjoin(proc_path,os.pardir)
sys.path.append(MGRootPath)
sys.path.append(pjoin(proc_path,'bin'))
import madgraph.various.lhe_parser as lhe_parser
import models.import_ufo as import_ufo
import banner
import madgraph.various.misc as misc
from madgraph.iolibs.files import cp
import madgraph.various.process_checks as process_checks

logging.basicConfig(level=logging.INFO)
logging.getLogger('madgraph.model').setLevel(logging.CRITICAL)

# ==============================================
# Helper functions
# ==============================================

def setup_channel(channelPath):
    """ Compiles the Resummation fortran steering 
    code and returns the corresponding process."""

    logging.info('Setting up resummation runner in\n   %s'%channelPath)

    try:
        misc.compile(cwd=channelPath, mode='fortran', job_specs = False)
    except:
        logging.error("Could not compile resummation runner in %s/%s"%(channelPath))
        raise
    
    logging.info('  ... Done.')
    
    return subprocess.Popen([os.path.join(channelPath,'resum_driver')], 
          stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, 
                                                                   cwd=channelPath)

def kill_processes(encountered_channels):
    """ Kills the MadLoop processes open """

    for proc in encountered_channels.values():
        try:
            proc[0].stdin.write('y\n')
        except:
            proc[0].kill()

def get_value(Runner, input_str):
    """ This version of get_me_value is simplified for the purpose of this
    class. No compilation is necessary. The CT mode can be specified."""
    Runner.stdin.write(input_str)
    output = Runner.stdout.readline()  
    return output

# ==============================================
# Start program
# ==============================================

# Specify here the path of the file to reweight
if len(sys.argv)>0:
  eventFile=sys.argv[1]
else:
    logging.error("Could not process the arguments. The arguments are: ./resum_reweighter.py <eventFilePath> <outputEvtName>")
    exit()

outSpecified=False
evtFileName = '.'.join(os.path.basename(eventFile).split('.')[:-1])
evtFileExt  = os.path.basename(eventFile).split('.')[-1]
try:
  outPath=pjoin(os.path.dirname(os.path.realpath(eventFile)),sys.argv[2])
  outSpecified=True
except:
  outPath=pjoin(os.path.dirname(os.path.realpath(eventFile)),'%s_rwgt.%s'%(evtFileName,evtFileExt))

i=0
while os.path.exists(outPath):
    i=i+1    
    if outSpecified and i==1:
        logging.warning("The chosen filename %s for the output already exists."%os.path.basename(outPath))
    outPath=pjoin(os.path.dirname(os.path.realpath(eventFile)),
                  '%s_rwgt_%i.%s'%(evtFileName,i,evtFileExt))

try:
    evtBanner = banner.Banner(eventFile)    
    evtFile = lhe_parser.EventFile(eventFile)
except:
    logging.error("Could not read event file %s."%eventFile)

# Detect the total number of event in an efficient way
n_evts = int(subprocess.Popen('grep -rin "<event>" %s | wc -l'%os.path.realpath(eventFile),
                        stdout=subprocess.PIPE, shell=True).communicate()[0])

logging.info("Writing out the reweighted event file on\n   %s"%outPath)
# Write out the output event file
outputEvtFile = open(outPath,'w')
outputEvtFile.write(evtFile.banner)
# List the channel present in the process output
channel_list = [ os.path.basename(chan) for chan in \
                 misc.glob('P*', pjoin(proc_path, 'SubProcesses')) ]

# Now scan over events
# For each encountered channel, store the corresponding process runner
encountered_channels=dict([])

ebeam1=float( evtBanner.get('run_card','ebeam1'))
ebeam2=float( evtBanner.get('run_card','ebeam2'))

t_before = time.time()
i_evt=0
wgt_summed = 0.0
squared_wgt_summed = 0.0
running_spread = 0.0
max_wgt = 0.0
n_percent_monitor=1
for event in evtFile:
    i_evt = i_evt+1
    t_evt_start = time.time()    
    # The channel Folder
    chanFolder=""
    EvtChannelName="VetoPrefactors"
    # The Channel name to setup, if needed
    ChannelNameToSetup=""

    # =======================================
    # Setup the runner if not already present
    # =======================================

    try:
        runner=encountered_channels[EvtChannelName][0]
    except KeyError:
        logging.info('Found new channel %s'%EvtChannelName)
        virtFolderPath=pjoin(os.path.dirname(os.path.realpath(eventFile)),EvtChannelName)
        try:
            runner=setup_channel(virtFolderPath)
        except:
            logging.error("Could not setup resum runner %s"%virtFolderPath)
            kill_processes(encountered_channels)
            exit()
    
            # Save the runners just created for ChannelNameToSetup and possibly related to
            # EvtChannelName as well.
        encountered_channels[EvtChannelName]=(runner,virtFolderPath)

    # =======================================
    # End of virtual runner setup
    # =======================================

    momenta=[[p.E,p.px,p.py,p.pz] for p in event if p.status==-1]
    pdg=[int(p.pid) for p in event if p.status==-1]

    event.parse_reweight()
    born_wgt=event.reweight_data['1001']

# PDG codes of the incoming particles
    f1=str(pdg[0])
    f2=str(pdg[1])
# Bjorken x's of the particles
    x1=str(momenta[0][0]/ebeam1)
    x2=str(momenta[1][0]/ebeam2) 
# Q is set to \sqrt(\hat{s})
    Q =str(2*math.sqrt(momenta[0][0]*momenta[1][0]))
# renormalisation scale used to generate the events
    muMad=str(event.scale)

    input_string=f1+' '+f2+' '+x1+' '+x2+' '+Q+' '+muMad+'\n'
    trial  = 0
    max_trial = 50
    worked = False
    new_wgt = 0.0

    output_string = get_value(runner,input_string)
    mode=float(output_string.split()[0])
    BCorr0=float(output_string.split()[1])
    BCorr1=float(output_string.split()[2])
    BCorrm0=float(output_string.split()[3])
    BCorrm1=float(output_string.split()[4])
    Efull=float(output_string.split()[5])
    EfullNLL=float(output_string.split()[6])
    Hcomp_fact=float(output_string.split()[7])
    Hcomp_factm=float(output_string.split()[8])
    alphah=float(output_string.split()[9])
    alpham=float(output_string.split()[10])
    EfO=float(output_string.split()[11])
    EfONLL=float(output_string.split()[12])
    alpha=float(output_string.split()[13])
    if mode==1:
        new_wgt= (1+alpha*BCorr1+(alphah/(2.0*math.pi))*(event.wgt/born_wgt+Hcomp_fact))*BCorr0*Efull*born_wgt
    if mode==2:
        new_wgt= BCorr0*EfullNLL*born_wgt
    if mode==3:
        new_wgt= (1+BCorrm1+EfO+(alpham/(2.0*math.pi))*(event.wgt/born_wgt+Hcomp_factm))*BCorrm0*born_wgt
    if mode==4:
        new_wgt= BCorrm0*(1+EfONLL)*born_wgt

    worked = True

    if i_evt<=5:
        logging.info('Event %i processed in %.2E seconds.'%(i_evt,time.time()-t_evt_start))
    if i_evt==5:
        logging.info('Further output of timing event by event suppressed.')

    event.wgt=new_wgt
    wgt_summed = wgt_summed + new_wgt
    squared_wgt_summed = squared_wgt_summed + new_wgt**2
    max_wgt = max(max_wgt,abs(new_wgt))
    if (i_evt/float(n_evts))>n_percent_monitor*0.01:
        logging.info('%i%% of %i events processed in %.2E seconds. Cross section: %.4E pb'\
                     %(n_percent_monitor,n_evts,(time.time()-t_before),wgt_summed/i_evt))
        n_percent_monitor=n_percent_monitor+1
    # Only start computing the cumulative spread after the first 10% of events.
    if n_percent_monitor>1:
        running_spread = running_spread + (new_wgt - float(wgt_summed/i_evt))**2
    outputEvtFile.write(str(event))

# Clean up things
outputEvtFile.close()
kill_processes(encountered_channels)

# Finish
logging.info("Successfully reweighted %i events in %.2E seconds in file\n  %s."\
    %(n_evts,(time.time()-t_before),outPath))

# Give statistics on the weight
logging.info("Resummation weight statistics,\n  avg       : %.2E\n  cum. var. : %.2E\n  max       : %.2E"\
    %((wgt_summed/n_evts),math.sqrt(running_spread/n_evts),max_wgt))

# Give statistics on the weight
logging.info("Resummation cross-section obtained\n  %.6E  +/- %.2E"\
    %((wgt_summed/n_evts),math.sqrt(squared_wgt_summed)/n_evts))
