#! /usr/bin/env python
import os
import random
import sys
import shutil
import fileinput
from math import sqrt
# This script helps you merge the results of many MC generation jobs
# submitted to the SPRACE cluster.

# Get parameters from command line
if len(sys.argv) < 2:  # the program name and the config file
	# stop the program and print an error message
	sys.exit("Usage: sprace_ALPGEN_merge.py <merge option>")

# Set parameters to the job
###############################################
mergeoption = sys.argv[1]
dryrun = 0
###############################################

if dryrun != 0:
	print "*** DRY RUN *** (change to dryrun = 0 to really run)"
	
# Get the directory contents, and only the "dir"
################################################
pwd = os.getcwd()
dirlist = os.listdir(pwd)
truedirlist = []
for entry in dirlist:
	if entry.find("dir") != -1:
		truedirlist.append(entry)

taskname = truedirlist[0].split("_dir_")[0]

# First part: concatenate all the unw files into
# a big one.
mergedunwfilename = taskname+".unw"
print mergedunwfilename

unwfilelist = []
for line in truedirlist:
	unwfile = line+"/"+taskname+".unw"
	unwfilelist.append(unwfile)

# See if we are actually doing this
if mergeoption == "1":
	print "Doing actual merging..."
	mergedunwfile = open(mergedunwfilename,'w') 
	for line in fileinput.input(unwfilelist):
		mergedunwfile.write(line)

# Second part: get all the unw_par and write a coherent 
# one, with the average crossection and error.  
mergedunwparfilename = taskname+"_unw.par"
unwparfilelist = []
for line in truedirlist:
	unwparfile = line+"/"+taskname+"_unw.par"
	unwparfilelist.append(unwparfile)

crosssections = []
weights = []
totalevents = 0

for line in fileinput.input(unwparfilelist):
	if line.find("Crosssection") != -1: 
		xsec = float(line.split()[0])
		error = float(line.split()[1])
		weight = 1.0/(error*error)
		crosssections.append(xsec)
		weights.append(weight)
	if line.find("unwtd") != -1:
		events = int(line.split()[0])
		totalevents = totalevents + events

numer = 0.0
denom = 0.0 
for xsec,wgt in zip(crosssections,weights):
	numer = numer+(xsec*wgt)
	denom = denom+wgt

avgxsec = numer/denom
erroravg = sqrt(1/denom)
totallumi = totalevents/avgxsec

fin = open(truedirlist[0]+"/"+taskname+"_unw.par")
mergedunwparfile = open(mergedunwparfilename,'w')
for line in fin:
	newline = line
	if line.find("Crosssection") != -1:
		newline = str(avgxsec)+"\t"+str(erroravg)+"\t! Crosssection +- error (pb)\n"
	if line.find("unwtd") != -1:
		newline = str(totalevents)+"\t\t"+str(totallumi)+"\t! unwtd events, lum (pb-1)\n"
	mergedunwparfile.write(newline)

# Third part: move all the .wgt and .par files to a single directory,
# and zip it.
if mergeoption == "1":
	zipdir = pwd+"/original"
	if not os.path.exists(zipdir):
		os.mkdir(zipdir)
		
	for line in truedirlist:
		wgtfile = line+"/"+taskname+".wgt"
		parfile = line+"/"+taskname+".par"
	       	shutil.copy(wgtfile,zipdir+"/"+line+".wgt")
		shutil.copy(parfile,zipdir+"/"+line+".par")
		
	relzipdir = os.path.split(zipdir)[1]
	command = "tar -czf "+taskname+"_originalFiles.tar.gz "+relzipdir
	print command