import sys, os
#from pymacs_beta import *
#from pymacs_beta.parser import *
#from pymacs_beta.forcefield import *
#from pymacs_beta.ndx import *
from pmx import *
from pmx.parser import *
from pmx.forcefield import *
from pmx.ndx import *


class FFatom:
    def __init__(self,list):
        self.type = list[0]
	self.sigmaA = list[1]
	self.epsA = list[2]
	self.A = list[3]
	self.sigmaB = list[4]
	self.epsB = list[5]

class FFfile:
    def __init__(self, fname=None):
	if fname is not None:
	    foo = fname.split('.')
	    bar = re.sub('ff','',foo[0]) 
	    self.name = bar
            self.atoms = []
	    self.read_ffitp(fname)	

    def read_ffitp(self,file):
	l = open(file).readlines()
	toSkip = "atomtypes"
	for line in l:
	    if toSkip not in line:
		self.atoms.append(FFatom(line.split()))

def get_eq_atoms(ff1,ff2):
    eq_at = []
    diff_at = []
    for at1 in ff1.atoms:
	switch = 0
	for at2 in ff2.atoms:
	    if (at1.type == at2.type and at1.sigmaA == at2.sigmaA and at1.epsA == at2.epsA and at1.sigmaB == at2.sigmaB and at1.epsB == at2.epsB):
		eq_at.append(at1)
		switch = 1
		break
	if switch == 0:
	    diff_at.append(at1)

    for at1 in ff2.atoms:
        switch = 0
        for at2 in eq_at:
            if (at1.type == at2.type and at1.sigmaA == at2.sigmaA and at1.epsA == at2.epsA and at1.sigmaB == at2.sigmaB and at1.epsB == at2.epsB):
	        switch = 1
                break
        if switch == 0:
  	    diff_at.append(at1)

    return (eq_at,diff_at)


def ffwrite(eq_at,diff_at,file):
    fp = open(file,'w')
    print >>fp, '[ atomtypes ]'
    for at in eq_at:
	print >>fp, '	%s	%s	%s	%s	%s	%s' % (at.type,at.sigmaA,at.epsA,at.A,at.sigmaB,at.epsB)
    for at in diff_at:
	print >>fp, '	%s	%s	%s	%s	%s	%s' % (at.type,at.sigmaA,at.epsA,at.A,at.sigmaB,at.epsB)

def main(argv):

        # define input/output files
        files= [
           FileOption("-ffitp1", "r",["itp"],"ffMOL1.itp",""),
           FileOption("-ffitp2", "r",["itp"],"ffMOL2.itp",""),
           FileOption("-ffitp_out", "w",["itp"],"ffMOL_out.itp",""),
	]
        # define options

        options=[]

        help_text = ()

        # pass options, files and the command line to pymacs

        cmdl = Commandline( argv, options = options,
                       fileoptions = files,
                       program_desc = help_text,
                       check_for_existing_files = False )

	ff1 = FFfile(cmdl['-ffitp1'])
	ff2 = FFfile(cmdl['-ffitp2'])

	#get 2 lists of: 1) equivalent atoms in ffMOL1.itp and ffMOL2.itp 2) different atoms
	(eq_at,diff_at) = get_eq_atoms(ff1,ff2)

	#write new ffMOL.itp
	ffwrite(eq_at,diff_at,cmdl['-ffitp_out'])


main( sys.argv )