#! /usr/bin/env python
from __future__ import division
import os, sys, argparse, re, string, pyslha

comment_pat = re.compile('\s*#.*$')
block_pat = re.compile('^\s*block\s+(\w+)',flags=re.I)
decay_pat = re.compile('^\s*decay\s+(\w+)\s+([-+\.\w]+)',flags=re.I)

data_pat = re.compile('^\s*((\d+\s+)+)(-?\d\S*)\s*$')
whitespace = re.compile('\s+')


PARAMS = {}

# set up the option parser for command line input 
parser = argparse.ArgumentParser(
	description='Modify a ThePEG model file with parameters from a matching SLHA file.'
)
parser.add_argument(
	'modelfile', 
	metavar='ThePEG_model', 
	help='ThePEG model file to use as template. Must have "# SLHA #"" annotations.'
)
parser.add_argument(
	'slhafile', 
	metavar='SLHA_file', 
	help='SLHA spectrum file.'
)
parser.add_argument(
    '-o','--output',
    default="FRModel_slha.model",
    help="Name for the output file"
)

args = parser.parse_args()

blocks_dict = {}

shla_file_path = args.slhafile

if os.path.exists(shla_file_path):

	slha_file = open(shla_file_path, 'r')
	d = pyslha.read(slha_file)
	for block in d.blocks:
		tmp_dict = {}
		for k, v in d.blocks[block].items():
			tmp_dict["{}:{}".format(block,k)]=v 
			if type(k)==tuple:
				label="{}_{}".format(block,k[0])
			else:
				label="{}_{}".format(block,k)
			PARAMS[label]=v
			print(label,block,k,v)
else:
	print("{} does not exist".format(shla_file_path))
	sys.exit(1)

template = open(args.modelfile)
output   = open(args.output,'w')

line = template.readline()
while line:
	split = line.split("${")
	if (len(split) == 1 ) :
		output.write(line)
	else :
		outputLine = split[0]
		for i in range(1,len(split)) :
			split2 = split[i].split("}")
			key = split2[0]
			print(key)
			if key in PARAMS :
				outputLine += str(PARAMS[key]) + split2[1]
				print(outputLine)
			else :
				print 'Parameter ',key,'not set in SLHA file, keeping default'
				outputLine = '#' + outputLine + "${" + split2[0] + "}" + split2[1]
			output.write(outputLine)
	line = template.readline()
output.write("\n")
print "Output written as : ",args.output
