#!/usr/bin/env startpython
from __future__ import print_function
# This tool takes plain text as input.
# Usually the sup.info. of a paper only contains a PDF copy of the forcefield, you can convert 
# this to plain text with this website: http://www.convertpdftotext.net/
# after this you can try this cleanup tool.
# If this tool crashes (usually with an out-of-range problem), you will need to check the 
# last section mentioned in the output (for example: "Found 51 sets of torsion parameters" means
# you need to check below the " 51    ! Nr of torsions" line). Usually there will be a missing
# space between some of the index numbers in this section. Add the spaces and try again.


import os
import argparse
import re
import sys
import copy
import fnmatch
import traceback

parser = argparse.ArgumentParser()
parser.add_argument('--input', '-i', help="input forcefield file, does not need to be the fixed-format required by reaxff (OCR-scanned pdf file for example)", dest='inputfile', required=True)
parser.add_argument('--output', '-o', help="output forcefied file, will be properly formatted", dest='outputfile')

options = parser.parse_args()

try:
	infile = open(options.inputfile, 'r')
	intext = infile.readlines()
	infile.close()
except:
	"failed to read the input file!"
	traceback.print_exc()

nlines = len(intext)
print ("Forcefield File contains %i lines" %(nlines))
###### clean the general parameters #######
# start by locating the number of general parameters
ngeneralparams = None
generalparamline = None
for i in range(nlines):
	if "number of general parameters" in intext[i].lower():
		ngeneralparams = int(intext[i].split()[0])
		generalparamline = i
		break
if ngeneralparams and generalparamline:
	print ("Found %i general parameters, starting on line %i" %(ngeneralparams, generalparamline) )
else:
	print ('failed to locate the number of general parameters, make sure the line contains"! Number of general parameters"')
	sys.exit(1)

# strip off junk
if generalparamline > 1:
	# keep the first line, this is usually a comment that starts with "Reactive MD Forcefield"
	intext = intext[generalparamline-1:]
nlines = len(intext)

# clean the header line
items = intext[1].lstrip().split(" ", 1)
items[0] = items[0].rjust(3).ljust(9)
items[1] = items[1].lstrip()
intext[1] = " ".join(items)

lineoffset = 2
# clean the parameter lines
for i in range(lineoffset, ngeneralparams+lineoffset):
	items = intext[i].lstrip().split(" ", 1)
	items[0] = items[0].rjust(10)
	items[1] = items[1].lstrip()
	intext[i] = " ".join(items)

lineoffset = lineoffset + ngeneralparams

###### clean the atomic parameters #######
natomparams = int(intext[lineoffset].split()[0])
print ("Found %i sets of atomic parameters" %(natomparams) )
items = intext[lineoffset].lstrip().split(" ", 1)
items[0] = items[0].lstrip().rjust(3).ljust(7)
items[1] = items[1].lstrip()
intext[lineoffset] = "".join(items)
intext[lineoffset+1] = "            " + intext[lineoffset+1].lstrip()
intext[lineoffset+2] = "            " + intext[lineoffset+2].lstrip()
intext[lineoffset+3] = "            " + intext[lineoffset+3].lstrip()
lineoffset = lineoffset + 4

for i in range(natomparams):
	tempoffset = lineoffset + 4*i
	items = intext[tempoffset].split()
	items[0] = " " + items[0].ljust(2)
	for k in range(1,9):
		items[k] = items[k].rjust(9)
	intext[tempoffset] = "".join(items)
	for j in range(1,4):
		items = intext[tempoffset + j].split()
		for k in range(8):
			items[k] = items[k].rjust(9)
		intext[tempoffset + j] = "   " + "".join(items)

lineoffset = lineoffset + 4*natomparams

###### clean the bonds parameters #######
nbondparams = int(intext[lineoffset].split()[0])
print "Found %i sets of bond parameters" %(nbondparams)
items = intext[lineoffset].lstrip().split(" ", 1)
items[0] = items[0].lstrip().rjust(3).ljust(7)
items[1] = items[1].lstrip()
intext[lineoffset] = "".join(items)
intext[lineoffset+1] = "                        " + intext[lineoffset+1].lstrip()
lineoffset = lineoffset + 2

for i in range(nbondparams):
	tempoffset = lineoffset + 2*i
	items = intext[tempoffset].split()
	items[0] = items[0].rjust(3)
	items[1] = items[1].rjust(3)
	for k in range(2,10):
		items[k] = items[k].rjust(9)
	intext[tempoffset] = "".join(items)
	
	items = intext[tempoffset + 1].split()
	for k in range(8):
		items[k] = items[k].rjust(9)
	intext[tempoffset + 1] = "      " + "".join(items)

lineoffset = lineoffset + 2*nbondparams

###### clean the off-diagonal parameters #######
ndiagparams = int(intext[lineoffset].split()[0])
print "Found %i sets of off-diagonal parameters" %(ndiagparams)
items = intext[lineoffset].lstrip().split(" ", 1)
items[0] = items[0].lstrip().rjust(3).ljust(7)
items[1] = items[1].lstrip()
intext[lineoffset] = "".join(items)
lineoffset = lineoffset + 1

for i in range(ndiagparams):
	tempoffset = lineoffset + i
	items = intext[tempoffset].split()
	items[0] = items[0].rjust(3)
	items[1] = items[1].rjust(3)
	for k in range(2,8):
		items[k] = items[k].rjust(9)
	intext[tempoffset] = "".join(items)

lineoffset = lineoffset + ndiagparams

###### clean the angle parameters #######
nangleparams = int(intext[lineoffset].split()[0])
print "Found %i sets of angle parameters" %(nangleparams)
items = intext[lineoffset].lstrip().split(" ", 1)
items[0] = items[0].lstrip().rjust(3).ljust(7)
items[1] = items[1].lstrip()
intext[lineoffset] = "".join(items)
lineoffset = lineoffset + 1

for i in range(nangleparams):
	tempoffset = lineoffset + i
	items = intext[tempoffset].split()
	items[0] = items[0].rjust(3)
	items[1] = items[1].rjust(3)
	items[2] = items[2].rjust(3)
	for k in range(3,10):
		items[k] = items[k].rjust(9)
	intext[tempoffset] = "".join(items)

lineoffset = lineoffset + nangleparams

###### clean the torsion parameters #######
ntorsionparams = int(intext[lineoffset].split()[0])
print "Found %i sets of torsion parameters" %(ntorsionparams)
items = intext[lineoffset].lstrip().split(" ", 1)
items[0] = items[0].lstrip().rjust(3).ljust(7)
items[1] = items[1].lstrip()
intext[lineoffset] = "".join(items)
lineoffset = lineoffset + 1

for i in range(ntorsionparams):
	tempoffset = lineoffset + i
	items = intext[tempoffset].split()
	items[0] = items[0].rjust(3)
	items[1] = items[1].rjust(3)
	items[2] = items[2].rjust(3)
	items[3] = items[3].rjust(3)
	for k in range(4,11):
		items[k] = items[k].rjust(9)
	intext[tempoffset] = "".join(items)

lineoffset = lineoffset + ntorsionparams

###### clean the hydrogen parameters #######
nhydrogenparams = int(intext[lineoffset].split()[0])
print "Found %i sets of hydrogen bond parameters" %(nhydrogenparams)
items = intext[lineoffset].lstrip().split(" ", 1)
items[0] = items[0].lstrip().rjust(3).ljust(7)
items[1] = items[1].lstrip()
intext[lineoffset] = "".join(items)
lineoffset = lineoffset + 1

for i in range(nhydrogenparams):
	tempoffset = lineoffset + i
	items = intext[tempoffset].split()
	items[0] = items[0].rjust(3)
	items[1] = items[1].rjust(3)
	items[2] = items[2].rjust(3)
	for k in range(3,7):
		items[k] = items[k].rjust(9)
	intext[tempoffset] = "".join(items)

lineoffset = lineoffset + nhydrogenparams

for i in range(len(intext)):
	# clean up any trailing spaces
	intext[i] = intext[i].rstrip() + '\n'

forcefiedtext = "".join(intext[:lineoffset])

if options.outputfile:
	outfile = open(options.outputfile, 'w')
	outfile.write(forcefiedtext)
	outfile.close()
else:
	print forcefiedtext