#! /usr/bin/python3
from __future__ import division
from __future__ import print_function
import os, sys, argparse, re, string

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()

with open(args.slhafile) as f:
	currentblock = None
	for line in f:
		line = comment_pat.sub('',line.rstrip())
		if not line: continue
		m = block_pat.match(line)
		d = decay_pat.match(line)
		if m:
			currentblock = m.group(1).upper()
		elif d:
			currentblock = None
			label = 'DECAY_%s' % d.group(1)
			data = float(d.group(2))
			PARAMS[label] = data
		elif currentblock is not None:
			d = data_pat.match(line)
			if d:
				index = whitespace.sub('_',d.group(1).rstrip())
				try:
					data  = float(d.group(3))
				except ValueError:
					continue
				label = '%s_%s' % (currentblock, index)
				#if label in PARAMS:
				PARAMS[label] = data


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]
			if key in PARAMS :
				outputLine += str(PARAMS[key]) + split2[1]
			else :
				name, suffix = key.rsplit('_',1)
				if suffix == 'ABS':
					mass = PARAMS[name]
					try:
						mass = mass.real
					except:
						pass
					outputLine += str(abs(mass)) + split2[1]
				elif suffix == 'CTAU':
					hbarc = 197.3269631e-15
					if(name in PARAMS) :
						width = PARAMS[name]
						ctau = (hbarc / width) if width != 0 else 0
					else :
						ctau = 0
					outputLine += str(ctau) + split2[1]
				elif suffix == 'WCUT':
					if(name in PARAMS) :
						width = PARAMS[name]
					else :
						width = 0.
					wcut = 10.0 * width
					outputLine += str(wcut) + split2[1]
				elif name == 'DECAY':
					outputLine += str(0.) + split2[1]
				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)
