import sys, os
from pmx import *
from pmx.options import *
from pmx.parser import *
from pmx.forcefield import *
from pmx.ndx import *
from string import digits
import copy as cp
import random
import numpy as np
import re

def split_itp_ffitp( itpfname, ffitpfname, itpoutfname ):
    itp = TopolBase( itpfname )
    fp = open(ffitpfname,'w')
    fp.write('[ atomtypes ]\n')
    for at in itp.atomtypes:
        fp.write('%8s %12.6f %12.6f %3s %12.6f %12.6f\n' % (at['name'],at['mass'],at['charge'],at['ptype'],at['sigma'],at['epsilon']))
    fp.close()
    itp.atomtypes = []
    itp.write(itpoutfname)

def structure_from_xyz( fname ):
    fp = open(fname,'r')
    lines = fp.readlines()
    fp.close()
    m = Model()
    counter = 0
    for l in lines[2:]:
        l = l.rstrip().lstrip().split()
        # add atom 
        a = Atom()
        a.name = l[0]
        a.resname = 'MOL'
        a.resnr = 1
        a.id = counter+1
        a.x[0] = float(l[1])
        a.x[1] = float(l[2])
        a.x[2] = float(l[3])
        m.atoms.append(a)
        counter+=1
    return(m)

def get_overall_charge_gaussian( fname ):
    fp = open(fname,'r')
    lines = fp.readlines()
    fp.close()
    for l in lines:
        if 'Charge' in l:
            l = l.rstrip()
            l = l.lstrip()
            foo = l.split()
            return(int(foo[2]))

def read_itp_top(fname):
    """ this function reads topology files from gaff """

    lines = open(fname).readlines()
    lines = kickOutComments(lines,';')
    itp = TopolBase(fname)#, ff = 'amber03')
    itp = Topology(fname,is_itp=True,assign_types=True,self_contained=True)#, ff = 'amber03')
#    atypes = read_atomtypes(lines, ff = 'amber03')
#    itp.atomtypes = atypes
    return itp

def write_ff(atypes, fname, ff='amber99sb'):
    fp = open(fname,'w')
    fp.write('[ atomtypes ]\n')
    for a in atypes:
        fp.write('%8s %12.6f %12.6f %3s %12.6f %12.6f\n' % (a['name'],a['mass'],a['charge'],a['ptype'],a['sigma'],a['epsilon']))
#    for atkey in atypes[0].keys():
#        at = atypes[atkey]

def set_ligname( ligname, m, itp ):
    # pdb
    for a in m.atoms:
        a.resname = ligname
    # itp
    for a in itp.atoms:
        a.resname = ligname
    itp.name = ligname

def gen_itp( fname_top, fname_gro, randnum ):
#    fname_itp = 'MOL.itp'
#    fname_ffitp = 'ffMOL.itp'
    itp = read_itp_top(fname_top)
#    itp.write(fname_itp)
#    write_ff(itp.atomtypes, fname_ffitp)
    return(itp)

def run_acpype_from_pdb( finp, ff, chargeMethod='bcc', charge=42 ):
    if charge==42: # guess the charge
        cmd = 'acpype -a {0} -o gmx -i {1} -c {2}'.format(ff,finp,chargeMethod)
    else:
        cmd = 'acpype -a {0} -o gmx -i {1} -c {2} -n {3}'.format(ff,finp,chargeMethod,charge)
    os.system(cmd)

def run_acpype( fname_prmtop, fname_inpcrd, charge, randnum, ff ):
    fname_top = 'MOL_GMX.top'
    fname_gro = 'MOL_GMX.gro'
    cmd = 'acpype -p '+fname_prmtop+' -x '+fname_inpcrd+' -o gmx -a '+ff+' -n '+str(charge)+' -c user -b MOL'
    os.system(cmd)
    return(fname_top,fname_gro)

def run_teleap( fname_mol2, fname_frcmod, randnum, ff ):
    fname_prmtop = 'prmtop_'+str(randnum)
    fname_inpcrd = 'inpcrd_'+str(randnum)
    # input file
    fp = open('leap.in','w')
#    fp.write("source %s/dat/leap/cmd/leaprc.%s\n" % (os.environ['AMBERHOME'],ff) )
    fp.write("loadamberparams %s/dat/leap/parm/%s.dat\n" % (os.environ['AMBERHOME'],ff) )
    fp.write("M = loadmol2 %s\n" % fname_mol2)
    fp.write("mods = loadamberparams %s\n" % fname_frcmod)
    fp.write("saveamberparm M %s %s\n" % (fname_prmtop,fname_inpcrd) )
    fp.write("quit\n")
    fp.close()
    # run teleap
    cmd = "teLeap -f leap.in"
    os.system(cmd)
    return(fname_prmtop,fname_inpcrd)

def generate_mol2( fname_ac_resp, halogens, halogens_nn, sigmaholes, randnum, ff, charge ):
    fname_mol2 = 'mol2_'+str(randnum)+'.mol2'
    # generate mol2
    # 1. attempt
    cmd = 'antechamber -i '+fname_ac_resp+' -fi ac -o '+fname_mol2+' -fo mol2 -at '+ff+' -nc '+str(charge)+' -j 4 -dr no >> /dev/null 2>&1'
    os.system(cmd)
    # 2. attempt (if the first attempt failed, try flag -j 5)
    if os.path.isfile(fname_mol2)==False:
        cmd = 'antechamber -i '+fname_ac_resp+' -fi ac -o '+fname_mol2+' -fo mol2 -at '+ff+' -nc '+str(charge)+' -j 5 -dr no >> /dev/null 2>&1'
        os.system(cmd)

    # add bonds
    fp = open(fname_mol2,'r')
    lines = fp.readlines()
    fp.close()
    mol2 = {}
    mol2['header'] = readSection(lines,'@<TRIPOS>MOLECULE','@')
    mol2['atoms'] = readSection(lines,'@<TRIPOS>ATOM','@')
    mol2['bonds'] = readSection(lines,'@<TRIPOS>BOND','@')
    mol2['substr'] = readSection(lines,'@<TRIPOS>SUBSTRUCTURE','@')
    fp = open(fname_mol2,'w')
    fp.write('@<TRIPOS>MOLECULE\n')
    for l in mol2['header']:
        fp.write(l)
    fp.write('@<TRIPOS>ATOM\n')
    for l in mol2['atoms']:
        foo = l.split()
        # check if sigmahole
        bSH = False
        for sh in sigmaholes:
            if sh.id==int(foo[0]):
                bSH = True
                break
        if bSH==True:
            fp.write("%7s%5s%15s%11s%11s%3s%10s%4s%15s\n" % (foo[0],foo[1],foo[2],foo[3],foo[4],'DU',foo[6],foo[7],foo[8]) )
        else: 
            fp.write(l)
    fp.write('@<TRIPOS>BOND\n')
    foo = ""
    for l in mol2['bonds']:
        fp.write(l)
        foo = l.split()
    bondnum = int(foo[0])
#    for h,hnn,sh in zip(halogens,halogens_nn,sigmaholes):
#        fp.write("%6s%6s%6s%2s\n" % (bondnum,h.id,sh.id,1) )
#        bondnum+=1
    fp.write('@<TRIPOS>SUBSTRUCTURE\n')
    for l in mol2['substr']:
        fp.write(l)
    fp.close()
    return(fname_mol2)

def read_frcmod( fname ):
    fp = open(fname,'r')
    lines = fp.readlines()
    fp.close()
    out = {}
    out['MASS'] = readSection(lines,'MASS','BOND')
    out['BOND'] = readSection(lines,'BOND','ANGLE')
    out['ANGLE'] = readSection(lines,'ANGLE','DIHE')
    out['DIHE'] = readSection(lines,'DIHE','IMPROPER')
    out['IMPROPER'] = readSection(lines,'IMPROPER','NONBON')
    out['NONBON'] = readSection(lines,'NONBON','*')
    return(out)

def gen_parmchk_dum( randnum ):
    fname = 'parmchk_'+str(randnum)
    fp = open(fname,'w')
    fp.write('PARM    DU      1               1        0.00   0       0\n')
    fp.close()
    return( fname )

def generate_frcmod( fname_ac_resp, halogens, halogens_nn, randnum, rule, ff, scaleD=1.0 ):
    fname_frcmod = 'frcmod_'+str(randnum)
    # from AmberTools20 need to generate an additional parmchk.dat file for dummies
    add_parmchk_fname = gen_parmchk_dum( randnum )
    # run frcmod
    cmd = 'parmchk2 -i {0} -f ac -o {1} -s {2} -atc {3}'.format(fname_ac_resp,fname_frcmod,ff,add_parmchk_fname)
    os.system(cmd)
    # correct frcmod
    frcmod = read_frcmod( fname_frcmod )
    fp = open(fname_frcmod,'w')
    fp.write('remark\n')
    # mass
    fp.write('MASS\n')
    for l in frcmod['MASS']:
        fp.write(l)
    # bonds
    fp.write('BOND\n')
    for l in frcmod['BOND']:
        if 'DU' in l:
            foo = l[0:5]#l.split()[0]
            l = foo+'       600.0 '+str(get_hal_sigma_distance( rule, foo, scaleD=scaleD ))+'\n'
        fp.write(l)
    # angles
    fp.write('\nANGLE\n')
    for l in frcmod['ANGLE']:
        if 'DU' in l:
            foo = l[0:8]#l.split()[0]
            l = foo+'           150.0     180.00\n'
        fp.write(l)
    # dihedrals
    fp.write('\nDIHE\n')
    for l in frcmod['DIHE']:
        if 'DU' in l:
            foo = l[0:11]#l.split()[0]
            l = foo+'     1    0.000         0.000           0.000\n'
        fp.write(l)
    # impropers
    fp.write('\nIMPROPER\n')
    for l in frcmod['IMPROPER']:
        fp.write(l)
    # nonbondeds
    fp.write('\nNONBON\n')
    for l in frcmod['NONBON']:
        if 'DU' in l:
            l = 'DU 1.0 0.0\n'
        fp.write(l)
    fp.close()
    return(fname_frcmod)

def antechamber_ac_resp( fname_ac, fname_resp, halogens, halogens_nn, sigmaholes, randnum, ff, charge ):
    fname_ac_resp = 'resp_'+str(randnum)+'.ac'
    # incorporate charges into .ac
    # 1. attempt
    cmd = 'antechamber -i '+fname_ac+' -fi ac -o '+fname_ac_resp+' -fo ac -c rc -cf '+fname_resp+' -j 4 -dr no -at '+ff+' -nc '+str(charge)+' > /dev/null 2>&1'
    os.system(cmd)
    # 2. attempt (if the first attempt failed, try flag -j 5)
    if os.path.isfile(fname_ac_resp)==False:
        cmd = 'antechamber -i '+fname_ac+' -fi ac -o '+fname_ac_resp+' -fo ac -c rc -cf '+fname_resp+' -j 5 -dr no -at '+ff+' -nc '+str(charge)+' > /dev/null 2>&1'
        os.system(cmd)

    # add bonds
#    add_bonds_ac( fname_ac_resp, halogens, sigmaholes )

    # in the ac file, sigmaholes have type LP, change it to DU (otherwise teleap will complain)
    ac_change_LP_to_DU( fname_ac_resp )

    return(fname_ac_resp)

def generate_respfiles( fname_ac, fname_esp, randnum ):
    fname_resp = 'respfile_'+str(randnum)
    # respgen
    cmd = 'respgen -i '+fname_ac+' -f resp1 -o respin1_'+str(randnum)+'.respin'
    os.system(cmd)
    cmd = 'respgen -i '+fname_ac+' -f resp2 -o respin2_'+str(randnum)+'.respin'
    os.system(cmd)
    # fit resp charges
    cmd = 'resp -O -i respin1_'+str(randnum)+'.respin -o respout1_'+str(randnum)+'.respout -e '+fname_esp+' -t qout_stage1'
    os.system(cmd)
    cmd = 'resp -O -i respin2_'+str(randnum)+'.respin -o respout2_'+str(randnum)+'.respout -e '+fname_esp+' -q qout_stage1 -t '+fname_resp
    os.system(cmd)
    return(fname_resp)

def generate_esp( fname_log, halogens, halogens_nn, m, randnum, rule, ff, scaleD=1.0, bESP=False ):
    fname_esp = 'espgen_'+str(randnum)+'.esp'

    # run espgen
    if bESP==False:
        cmd = 'espgen -i '+fname_log+' -o '+fname_esp
    else: # esp is provided as input
        cmd = 'cp {0} {1}'.format(fname_log,fname_esp)
    os.system(cmd)

    # insert sigmahole into esp
    fp = open(fname_esp,'r')
    lines = fp.readlines()
    fp.close()

    # extract coordinates from esp
    espcoord = extract_esp_coords( lines )

    # find halogen coords
    hcoords = []
    hnncoords = []
    for ah,ahnn in zip(halogens,halogens_nn):
        hcoords.append(espcoord[ah.id-1])
        hnncoords.append(espcoord[ahnn.id-1])

    # write esp file with inserted sigmahole
    fp = open(fname_esp,'w')
    bohr = 0.529177249
    i = 0
    bHalFinished = False
    for l in lines:
        lsplit = l.split()
        if i==0:
            foo = l[:5] # first 5 elements
            bar = l[5:]
            foo = foo.lstrip()
            foo = int(foo)+len(halogens)
            fp.write("%5d%s" % (foo,bar) )
            i+=1
            continue
        elif i>0 and len(lsplit)==4 and bHalFinished==False: # write halogens
            hnum = 0
            for ah,ahnn in zip(halogens,halogens_nn):
                ch = [ float(hcoords[hnum][0]),float(hcoords[hnum][1]),float(hcoords[hnum][2]) ]
                chnn = [ float(hnncoords[hnum][0]),float(hnncoords[hnum][1]),float(hnncoords[hnum][2]) ]
                vec = [ch[0]-chnn[0],ch[1]-chnn[1],ch[2]-chnn[2]]
                vecnorm = np.sqrt(vec[0]*vec[0]+vec[1]*vec[1]+vec[2]*vec[2])
                vec = vec/vecnorm
                halname = ah.name.lower()
                dist = get_hal_sigma_distance( rule, halname, scaleD=scaleD )
                newpos = [ch[0]+vec[0]*dist/bohr , ch[1]+vec[1]*dist/bohr, ch[2]+vec[2]*dist/bohr ]

                fp.write("%32.7E%16.7E%16.7E\n" % (newpos[0],newpos[1],newpos[2]) )
                hnum += 1
            bHalFinished = True
        fp.write(l)
        i+=1
    fp.close()
    return(fname_esp)

def extract_esp_coords( lines ):
    coord = []
    for l in lines[1:]:
        foo = l.split()
        if len(foo)>3:
            return(coord)
        coord.append(foo)

def get_halogen_position( ah, ahnn, rule, scaleD=1.0 ): 
    vec = [ah.x[0]-ahnn.x[0],ah.x[1]-ahnn.x[1],ah.x[2]-ahnn.x[2]]
    vecnorm = np.sqrt(vec[0]*vec[0]+vec[1]*vec[1]+vec[2]*vec[2])
    vec = vec/vecnorm
    halname = ah.name.lower()
    dist = get_hal_sigma_distance( rule, halname, scaleD=scaleD )
    newpos = [ah.x[0]+vec[0]*dist , ah.x[1]+vec[1]*dist, ah.x[2]+vec[2]*dist ]
    return(newpos)

def get_ac_atom_dict( lines ):
    out = {}
    for l in lines:
        if l.startswith('ATOM'):
            l = l.rstrip()
            foo = l.split()
            out[int(foo[1])] = foo[2]
    return(out)

def add_bonds_ac( fname_ac, halogens, sigmaholes):
    # read
    fp = open(fname_ac,'r')
    lines = fp.readlines()
    fp.close()
    # get bond number
    foo = lines[-1].split()
    bondnum = int(foo[1])
    # get atom name dictionary: dict[id] = at_name
    ac_atom_dict = get_ac_atom_dict( lines )
    # append new bonds
    fp = open(fname_ac,'a')
    i = 0
    for h,sh in zip(halogens,sigmaholes):
        bondnum+=1
        halname = ac_atom_dict[h.id]
        shname = ac_atom_dict[sh.id]
        fp.write("%4s%5d%5d%5d%5d%7s%5s\n" % ("BOND",bondnum,h.id,sh.id,1,halname,shname) )
        i+=1
    fp.close()

def ac_change_LP_to_DU( fname_ac ):
    regexp = re.compile('LP')
    # read
    fp = open(fname_ac,'r')
    lines = fp.readlines()
    fp.close()
    # write
    fp = open(fname_ac,'w')
    for l in lines:
        if l.startswith('ATOM'):
            if 'EP' in l:
                l = regexp.sub('DU',l)
        fp.write(l)
    fp.close()

def ac_change_DU_to_LP( fname_ac ):
    regexp = re.compile('DU')
    # read
    fp = open(fname_ac,'r')
    lines = fp.readlines()
    fp.close()
    # write
    fp = open(fname_ac,'w')
    for l in lines:
        if l.startswith('ATOM'):
            if 'EP' in l:
                l = regexp.sub('LP',l)
        fp.write(l)
    fp.close()

def antechamber_ac( fname_pdb, m, halogens, sigmaholes, charge, randnum, ff ):
    fname_ac = 'sigmahole_'+str(randnum)+'.ac'
    resname = 'MOL' #m.residues[0].resname
    # 1. attempt
    cmd = 'antechamber -i '+fname_pdb+' -fi pdb -o '+fname_ac+' -fo ac -j 4 -rn '+resname+' -nc '+str(charge)+' -dr no'+' -at '+ff+' > /dev/null 2>&1'
    os.system(cmd)
    # 2. attempt (if the first attempt failed, try flag -j 5)
    if os.path.isfile(fname_ac)==False:
        cmd = 'antechamber -i '+fname_pdb+' -fi pdb -o '+fname_ac+' -fo ac -j 5 -rn '+resname+' -nc '+str(charge)+' -dr no'+' -at '+ff+' > /dev/null 2>&1'
        os.system(cmd)

    # add bonds
    add_bonds_ac( fname_ac, halogens, sigmaholes )

#    ac_change_DU_to_LP( fname_ac )

    return(fname_ac)

def antechamber_pdb( fname, charge, randnum, ff ):
    fname_pdb = 'lig_gaussian_'+str(randnum)+'.pdb'
    resname = 'MOL' #m.residues[0].resname
    cmd = 'antechamber -i '+fname+' -fi gout -o '+fname_pdb+' -fo pdb -j 4 -rn '+resname+' -nc '+str(charge)+' -dr no'+' -at '+ff+' > /dev/null 2>&1'
    os.system(cmd)
    return(fname_pdb)

def add_sigmahole_pdb( halogens, halogens_nn, m, randnum, rule, ff, scaleD=1.0 ):
    sigmahole_num = 1
    sigmaholes = []
    fname_pdb = 'str_'+str(randnum)+'.pdb'
    for ah,ahnn in zip(halogens,halogens_nn):
        newpos = get_halogen_position( ah, ahnn, rule, scaleD=scaleD ) 

        newatom = cp.deepcopy(ah)
        newatom.name = 'EP'+str(sigmahole_num)
        newatom.x = newpos
        newatom.id = len(m.atoms)+1
        m.atoms.append(newatom)
        sigmaholes.append(newatom)

        fname_pdb = 'sigmahole_'+str(randnum)+'.pdb'
        sigmahole_num += 1

    m.write(fname_pdb)
    return(fname_pdb,sigmaholes)

def add_ffitp_sigmahole( ffitp, hid, hnnid, num, ff ):
    bFound = False
    for a in ffitp.atomtypes:
        if 'DU' in a['name']:
            bFound = True

    if bFound==False:
        atomtype = dict()
        atomtype['name'] = str('DU')
        atomtype['bond_type'] = str('DU')
        atomtype['mass'] = float(0.0)
        atomtype['charge'] = float(0.0)
        atomtype['ptype'] = str('A')
        atomtype['sigma'] = float(0.0)
        atomtype['epsilon'] = float(0.0)
        ffitp.atomtypes.append(atomtype)

#    if 'DU' not in ffitp.atomtypes.keys():
#        if 'opls' in ff:
#            ffitp.atomtypes['DU'] = [0.0,0.0,'A',0.0,0.0]
#        if 'cgenff' in ff:
#            ffitp.atomtypes['DU'] = [0.0,0.0,'A',0.0,0.0]
            # the other parameters are already incorporated into CGenFF: nbfix and LJ changes for halogens
#        if 'gaff' in ff:
#            ffitp.atomtypes['DU'] = [0.0,0.0,'A',0.0,0.0]

def add_itp_sigmahole( itp, hid, hnnid, num, rule, ff, scaleD=1.0 ):
    ah = itp.atoms[hid-1] # halogen
    ahnn = itp.atoms[hnnid-1] # atom to which halogen is bound (halogen nearest neighbor)
    halname = ah.name.lower()
    newatom = cp.deepcopy(ah)
    newatom.name = 'EP'+str(num)
    newatom.cgnr = len(itp.atoms)+1

    ########## atom ###########
    if 'cl' in halname:        
        if 'jorgensen' in rule:
            newatom.q = 0.075
            ah.q -= newatom.q
        elif 'gutierrez' in rule:
            newatom.q = 0.05
            ah.q -= 0.08
            ahnn.q += 0.03
        if 'ibrahim' in rule:
            newatom.q = 0.035
            ah.q -= newatom.q
    elif 'br' in halname:
        if 'jorgensen' in rule:
            newatom.q = 0.1
            ah.q -= newatom.q
        elif 'gutierrez' in rule:
            newatom.q = 0.05
            ah.q -= 0.08
            ahnn.q += 0.03
        if 'ibrahim' in rule:
            newatom.q = 0.04
            ah.q -= newatom.q
    elif 'i' in halname:
        if 'jorgensen' in rule:
            newatom.q = 0.11
            ah.q -= newatom.q
        elif 'gutierrez' in rule:
            newatom.q = 0.05
            ah.q -= 0.09
            ahnn.q += 0.04
        if 'ibrahim' in rule:
            newatom.q = 0.055
            ah.q -= newatom.q
    newatom.m = 0.0

    newatom.id = len(itp.atoms)+1
    newatom.atomtype = 'DU'
    itp.atoms.append(newatom)

    ######### bond ##########
    d = get_hal_sigma_distance( rule, halname, scaleD=scaleD )/10.0 # in nm

#    newbond = [hid,newatom.id,1,d,0.0]
    newbond = [ah,newatom,1,[np.round(d,4),0.0]]
    itp.bonds.append(newbond)
    # also add a bond to the halogen nearest neighbor:
    # this is done to exclude 1-2 and 1-3 interactions for the virtual particle
    # in the same way as it is done for the halogen atom,
    # because the virtual particle is considered a part of halogen
    bondlength = get_bond_length( itp, hid, hnnid )
#    newbond2 = [hnnid,newatom.id,1,bondlength+d,0.0]
    newbond2 = [ahnn,newatom,1,[np.round(bondlength+d,4),0.0]]
    itp.bonds.append(newbond2)

    ######### pairs #########
#    gen_pairs( itp, hid, hnnid, newatom.id )
    gen_pairs_alt( itp, hid, newatom ) # generate pairs as if sigma hole is the same atom as halogen

    ######### virtual site ##
    hnnid2 = 1
    dvsite = get_vsite_param( halname, ff, itp, hid, hnnid, rule, scaleD=scaleD ) 
    vsite = [newatom,ahnn,ah,1,dvsite]
    itp.has_vsites2 = True
    itp.virtual_sites2.append(vsite)

def check_and_add_bond( itp, a1, a2, d ):
    ids = [a1.id,a2.id]
    for b in itp.bonds:
        if (b[0].id in ids) and (b[1].id in ids):
            return(True)
    # bond not found
    newbond = [a1,a2,1,[np.round(d,4),999.9]] 
    itp.bonds.append(newbond)

def add_vsite_respfit( ff, itp, halogens, halogens_nn, sigmaholes, rule, scaleD=1.0 ):
    for ah,ahnn,sh in zip(halogens,halogens_nn,sigmaholes):
        #d_hal_sigma = (ah-sh)/10.0 # convert to nm
        d_hal_sigma = get_hal_sigma_distance( rule, ah.name.lower(), scaleD=scaleD )/10.0 # in nm
        # check if this bond is already in .itp, if not, add it
        check_and_add_bond( itp, ah, sh, d_hal_sigma )
#        d_hal_hnn = (ah-ahnn)/10.0 # convert to nm
        d_hal_hnn = get_bond_length( itp, ah.id, ahnn.id )
        dvsite = (d_hal_hnn+d_hal_sigma)/d_hal_hnn
        vsite = [sh,ahnn,ah,1,dvsite]
        itp.has_vsites2 = True
        itp.virtual_sites2.append(vsite)
#    return(itp)

def get_vsite_nums( itp ):
    dums = []
    for a in itp.atoms:
        if a.name.startswith('EP'):
            dums.append(a.id)
    return(dums)

def remove_angles_dihedrals( itp ):
    dums = get_vsite_nums( itp )
    
    # remove angles
    newangles = []
    for ang in itp.angles:
        a1 = ang[0].id
        a2 = ang[1].id
        a3 = ang[2].id
        if (a1 in dums) or (a2 in dums) or (a3 in dums):
            continue
        else:
            newangles.append(ang)
    itp.angles = newangles

    # remove dihedrals
    newdihedrals = []
    for dih in itp.dihedrals:
        a1 = dih[0].id
        a2 = dih[1].id
        a3 = dih[2].id
        a4 = dih[3].id
        if (a1 in dums) or (a2 in dums) or (a3 in dums) or (a4 in dums):
            continue
        else:
            newdihedrals.append(dih)
    itp.dihedrals = newdihedrals

def add_sh_hnn_bond_respfit( ff, itp, halogens, halogens_nn, sigmaholes ):
    for ah,ahnn,sh in zip(halogens,halogens_nn,sigmaholes):
#        d_hnn_sigma = (ahnn-sh)/10.0 # convert to nm
        # extract these distances from .itp, as the structure may not match equilibrium bond
        d_hal_hnn = get_bond_length( itp, ah.id, ahnn.id ) 
        d_hal_sh = get_bond_length( itp, ah.id, sh.id ) 
        d_hnn_sigma = d_hal_hnn + d_hal_sh
#        newbond = [ahnn.id,sh.id,1,d_hnn_sigma,999.99]
        newbond = [ahnn,sh,1,[d_hnn_sigma,999.99]]
        itp.bonds.append(newbond)
#    return(itp)

def add_pairs_respfit( itp, halogens, sigmaholes ):
    newpairs = cp.deepcopy(itp.pairs)
    pairstoremove = []
    pairstoadd = []
    for hal,sh in zip(halogens,sigmaholes):
        for p in newpairs:
            if p[0].id==sh.id or p[1].id==sh.id:
                pairstoremove.append(p)
                continue
            if p[0].id==hal.id:
                newpair = [sh,p[1],p[2]]
            elif p[1].id==hal.id:
                newpair = [p[0],sh,p[2]]
            else:
                continue
            pairstoadd.append(newpair)
    # remove
    for p in pairstoremove:
        newpairs.remove(p)
    # add
    for p in pairstoadd:
        newpairs.append(p)
    itp.pairs = newpairs

def get_bond_length( itp, hid, hnnid ):
    for b in itp.bonds:
        if b[0].id==hnnid and b[1].id==hid:
            return(b[3][0])
        elif b[0].id==hid and b[1].id==hnnid:
            return(b[3][0])
    return(0.0)

def gen_pairs_alt( itp, hid, newat ):
    pairat = []
    for p in itp.pairs:
        if p[0].id==hid:
            pairat.append(p[1])
        if p[1].id==hid:
            pairat.append(p[0])
    for p in pairat:
        newpair = [newat, p, 1]
        itp.pairs.append(newpair)

def gen_pairs( itp, hid, hnnid, newid ):
    pairid = []
    for b in itp.bonds:
        if b[0]==hnnid and b[1]!=hid:
            pairid.append(b[1])
        elif b[0]!=hid and b[1]==hnnid:
            pairid.append(b[0])
    for p in pairid:
        newpair = [newid, p, 1]
        itp.pairs.append(newpair)
        

def get_vsite_param( halname, ff, itp, hid, hnnid, rule, scaleD=1.0 ):
    vs = 0.0
    d_hal_sigma = get_hal_sigma_distance( rule, halname, scaleD=scaleD )/10.0 # converting to nm
    d_c_hal = 0.0
    for b in itp.bonds:
        if (b[0].id==hid and b[1].id==hnnid) or (b[0].id==hnnid and b[1].id==hid):
            d_c_hal = b[3][0]
    if d_c_hal==0.0:
        print("Something went wrong: no halogen binding atom was found")
        sys.exit(0)
    vs = (d_c_hal+d_hal_sigma)/d_c_hal 
    return(vs)

def get_hal_sigma_distance( rule, halname, scaleD=1.0 ):
    # all in Angstroms
    d = 0.0
    if 'cl' in halname:
        if 'jorgensen' in rule:
            d = 1.6
        elif 'gutierrez' in rule:
            d = 1.64
        if 'ibrahim' in rule:
            d = 1.945
    elif 'br' in halname:
        if 'jorgensen' in rule:
            d = 1.6
        elif 'gutierrez' in rule:
            d = 1.89
        if 'ibrahim' in rule:
            d = 2.02
    elif 'i' in halname:
        if 'jorgensen' in rule:
            d = 1.8
        elif 'gutierrez' in rule:
            d = 2.20
        if 'ibrahim' in rule:
            d = 2.15
    d = scaleD*d
    return(d)

def add_pdb_sigmahole( m, hid, hnnid, num, rule, ff, scaleD=1.0 ):
    ah = m.atoms[hid-1]
    ahnn = m.atoms[hnnid-1]
    vec = [ah.x[0]-ahnn.x[0],ah.x[1]-ahnn.x[1],ah.x[2]-ahnn.x[2]]
    vecnorm = np.sqrt(vec[0]*vec[0]+vec[1]*vec[1]+vec[2]*vec[2])
    vec = vec/vecnorm
    halname = ah.name.lower()

    dist = get_hal_sigma_distance( rule, halname, scaleD=scaleD )
    newpos = [ah.x[0]+vec[0]*dist , ah.x[1]+vec[1]*dist, ah.x[2]+vec[2]*dist ]

    newatom = cp.deepcopy(ah)
    newatom.name = 'EP'+str(num)
    newatom.x = newpos   
    newatom.id = len(m.atoms)+1
    m.atoms.append(newatom)

def gen_sigmaholes( itp, m, ffitp, halogens, halogens_nn, rule, ff='oplsaa', scaleD=1.0 ):
    for i in range(0,len(halogens)):
        hal = halogens[i]
        halnn = halogens_nn[i]
        hal_id = hal.id
        halnn_id = halnn.id
        sigmahole_num = i+1

        add_pdb_sigmahole( m, hal_id, halnn_id, sigmahole_num, rule, ff, scaleD=scaleD )
        add_itp_sigmahole( itp, hal_id, halnn_id, sigmahole_num, rule, ff, scaleD=scaleD )
        add_ffitp_sigmahole( ffitp, hal_id, halnn_id, sigmahole_num, ff )
    

def find_halogens( itp, m=False, bSH=True ):
    halogens = []
    halogens_nn = []
    if bSH==False:
        return(halogens,halogens_nn)

    if m==False:
        atoms = itp.atoms
    else:
        atoms = m.atoms

    for at in atoms:
        foo = at.name.lower()
        foo = foo.translate(str.maketrans('','',digits))
        if foo.startswith('cl') or foo.startswith('br') or foo.startswith('i'):
            halogens.append( at )

            if m==False:
                for b in itp.bonds:
                    a1 = b[0]
                    a2 = b[1]
                    if a1.id==at.id:
                        halogens_nn.append( itp.atoms[a2.id-1] )
                    elif a2.id==at.id:
                        halogens_nn.append( itp.atoms[a1.id-1] )
            else: # find neighbor in pdb based on distance
                mind = 9999.999
                nn = ""
                for a in atoms:
                    aname = a.name.lower()
                    aname = aname.translate(str.maketrans('','',digits))
                    if aname.startswith('h') or (aname==foo):
                        continue
                    d = (at.x[0]-a.x[0])**2 + (at.x[1]-a.x[1])**2 + (at.x[2]-a.x[2])**2
                    if d<mind:
                        mind = d
                        nn = a
                halogens_nn.append( nn )

    return( halogens,halogens_nn )

def write_ffitp(atypes, itp, fname, ff='opls'):
    fp = open(fname,'w')
    fp.write('[ atomtypes ]\n')
    ## output ffitp, if ffitp was read ##
    for atkey in atypes.keys():
        at = atypes[atkey]
        if 'opls' in ff:
            fp.write('%8s %12.6f %12.6f %3s %12.6f %12.6f\n' % (atkey,at[0], at[1], at[2], at[3], at[4]) )
        elif 'cgenff' in ff:
            fp.write('%8s %12.6f %12.6f %3s %12.6f %12.6f\n' % (atkey,at[0], at[1], at[2], at[3], at[4]) )
        elif 'gaff' in ff:
            fp.write('%8s %12.6f %12.6f %3s %12.6f %12.6f\n' % (atkey,at[0], at[1], at[2], at[3], at[4]) )
  #      else:
   #         print >>fp, '%8s %12.6f %12.6f %3s %12.6f %12.6f' % (at[0], at[1], at[2], at[3], at[4], at[5])
    ## also output sigmahole dummies, if present ##
#    for a in itp.atoms:
#        if 'DU' in a.atomtype:
#            if 'opls' in ff:
#                print >>fp, '%8s %12.6f %12.6f %3s %12.6f %12.6f' % (atkey,0.0, 0.0, 0.0, 0.0, 0.0)
#            elif 'cgenff' in ff:
#                print >>fp, '%8s %12.6f %12.6f %3s %12.6f %12.6f' % (atkey,0.0, 0.0, 0.0, 0.0, 0.0)
#            elif 'gaff' in ff:
#                print >>fp, '%8s %12.6f %12.6f %3s %12.6f %12.6f' % (atkey,0.0, 0.0, 0.0, 0.0, 0.0)

################################################################################33


def main(argv):

    version = "1.1"

	# define input/output files
    files= [
	   FileOption("-itp", "r/o",["itp"],"MOL.itp",""),
	   FileOption("-ffitp", "r/o",["itp"],"ffMOL.itp",""),
	   FileOption("-pdb", "r/o",["pdb"],"mol.pdb",""),
	   FileOption("-xyz", "r/o",["xyz"],"orca.xyz",""),
	   FileOption("-log", "r/o",["log"],"gaussian.log",""),
	   FileOption("-esp", "r/o",["esp"],"esp.esp",""),
	   FileOption("-oitp", "o",["itp"],"outMOL.itp",""),
	   FileOption("-offitp", "o",["itp"],"outffMOL.itp" ,""),
	   FileOption("-opdb", "o",["pdb"],"outpdb.pdb","" ),
	    ]

	# define options
    options=[
           Option( "-ff", "string", "opls", "force-field: opls, cgenff, gaff"),
           Option( "-rule", "string", "default", "ibrahim, jorgensen, gutierrez"),
           Option( "-ligname", "string", "", "set ligand name (if none given, will not change the name)"),
           Option( "-q", "int", 0, "charge of the molecule"),
           Option( "-scaleD", "float", 1.0, "additionally can scale halogen-sigmahole distance which is set by a rule (may be needed for stability)"),
           Option( "-sh", "bool", True, "if set to False, no sigmahole will be generated"),
           Option( "-clean", "bool", False, "clean working files"),
            ]

    help_text = ('OPLS sigma hole based on Jorgensen and Schyman, JCTC, 2012',
                     'Distances between halogen and sigma hole are fixed.',
                     'Charges for the sigma holes are fixed as well.',
                     'Only charges of halogen atoms are adjusted.',
                     '',
		     'CGenFF sigma hole: Gutierrez et al, Bioorg and Med Chem, 2016',
                     'Distances and sigma hole charges fixed.',
                     'The charges of halogen and first carbon atoms are changed (Table 1).',
		     '(NOTE: paramchem does not change charge of the first carbon)',
                     'LJ parameters and NBFIX are needed, but they are already',
                     'incorporated in the later versions of Charmm36m with CGenFF.',
                     '',
                     'GAFF sigma hole: Ibrahim, JCC, 2011',
                     'Distances between halogen and sigma hole are fixed.',
                     'Sigma hole charges set by approximating values observed for various ligands.',
                     'Only charges of halogen atoms are adjusted.',
                     'If Gaussian .log file is provided, proper RESP fitting is performed.',
                     'This way, sigma hole charges are fit and not pre-set.',
		     'If gaussian.log is provided, RESP fit will be performed',
                     '',
                     'The flag -rule allows mixing the rules to generate virtual particles.',
                     '"default" keeps the above described rules for the respective force-fields.',
                     'The rules may arbitrarily be used with any force-field by setting the -rule flag.',
                     ''
                  )

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

    ff = cmdl['-ff'].lower()
    if 'opls' in ff:
        ff = 'oplsaa'
    if ('charmm' in ff) or ('cgenff' in ff):
        ff = 'cgenff'
    if ('gaff2' in ff):
        ff = 'gaff2'
    elif ('gaff' in ff) or ('amber' in ff):
        ff = 'gaff'

    bSH = cmdl['-sh']

    bITPinp = cmdl.opt['-itp'].is_set
    bFFITPinp = cmdl.opt['-ffitp'].is_set
    itpfname = cmdl['-itp']
    ffitpfname = cmdl['-ffitp']

    ##########################
    ### identify the rules ###
    rule = cmdl['-rule'].lower()
    if rule.startswith('default'):
        if 'opls' in ff:
            rule = 'jorgensen'
        elif ('charmm' in ff) or ('cgenff' in ff):
            rule = 'gutierrez'
        elif 'gaff' in ff:
            rule = 'ibrahim'
    elif rule.startswith('jorgensen'):
        rule = 'jorgensen'
    elif rule.startswith('gutierrez'):
        rule = 'gutierrez'
    elif rule.startswith('ibrahim'):
        rule = 'ibrahim'
    else:
        sys.stdout.write('Rule set with the -rule flag is unknown. Exiting...\n')
        sys.exit(1)
    ########################### 

    if cmdl.opt['-log'].is_set==False and cmdl.opt['-esp'].is_set==False:
        if cmdl.opt['-pdb'].is_set==False:
            print('Need to provide Gaussian output or ESP file or PDB structure file')
            sys.exit(0)
        if cmdl.opt['-esp'].is_set==True and cmdl.opt['-xyz'].is_set==False and cmdl.opt['-pdb'].is_set==False:
            print('Need to provide XYZ or PBD when using ESP file')
            sys.exit(0)
    if cmdl.opt['-esp'].is_set==True and cmdl.opt['-xyz'].is_set==False and cmdl.opt['-pdb'].is_set==False:
        print('Need to provide XYZ or PDB when using ESP file')
        sys.exit(0)

    if cmdl.opt['-pdb'].is_set:
        # just read pdb
        if bITPinp or cmdl.opt['-log'].is_set or cmdl.opt['-esp'].is_set:
            m = Model().read(cmdl['-pdb'])
        # need to generate pdb first, then read it (only works for gaff now)
        elif( 'gaff' in ff):
            if cmdl.opt['-q'].is_set:
                run_acpype_from_pdb( cmdl['-pdb'],ff,charge=cmdl['-q'] )
            else:
                run_acpype_from_pdb( cmdl['-pdb'],ff )
            m = Model().read(cmdl['-pdb'])
            pdbname = cmdl['-pdb'].split('/')[-1].split('.')[-2]
            acpypepath = ''.join( cmdl['-pdb'].split('.')[0:-1] )+'.acpype'
            itpfname = '{0}/{1}_GMX.itp'.format(acpypepath,pdbname)
            ffitpfname = cmdl['-offitp']
            split_itp_ffitp( itpfname, ffitpfname, cmdl['-oitp'] )
            bITPinp = True
            bFFITPinp = True
            itpfname = cmdl['-oitp']
        else:
            print('For this ff cannot generate topology straight away from pdb')
            sys.exit(0)
#    sys.exit(0)

    scaleD = cmdl['-scaleD']
    bClean = cmdl['-clean']
    charge = cmdl['-q']

    bRESP = False
    if bITPinp:
        itp = TopolBase(itpfname)#, ff=ff)
    elif cmdl.opt['-log'].is_set:
        if 'gaff' not in ff:
            print('RESP fit can be done only for GAFF force field')
            sys.exit(0)
        gaussianLogFile = cmdl['-log']
        bRESP = True
    if bFFITPinp:
        ffitp = TopolBase( ffitpfname )#ff=ff )
    else:
        ffitp = TopolBase( filename=None )#ff=ff )

    if bITPinp:
        halogens = []
        halogens_nn = []
        if bSH==True:
            (halogens, halogens_nn) = find_halogens( itp,m=False )
        gen_sigmaholes( itp, m, ffitp, halogens, halogens_nn, rule, ff=ff, scaleD=scaleD )
#        write_ffitp( ffitp.atomtypes, itp, cmdl['-offitp'], ff=ff )
        write_ff( ffitp.atomtypes, cmdl['-offitp'], ff=ff )
    elif cmdl.opt['-log'].is_set or cmdl.opt['-esp'].is_set:
        randnum = random.randrange(1000,9999) # needed for temporary files
        # get charge
        if cmdl.opt['-log'].is_set:
            # get overall charge from gaussian output file
            charge = get_overall_charge_gaussian( cmdl['-log'] )
        elif cmdl.opt['-q'].is_set:
            charge = cmdl.opt['-q']
        else: # if nothing is provided, set charge to 0.0 
            charge = 0
        # get pdb
        if cmdl.opt['-log'].is_set:
            # extract pdb from Gaussian output
            fname_pdb = antechamber_pdb( cmdl['-log'], charge, randnum, ff=ff )
            m = Model().read(fname_pdb)
        elif cmdl.opt['-xyz'].is_set:
            m = structure_from_xyz( cmdl['-xyz'] )
        elif cmdl.opt['-pdb'].is_set:
            fname_pdb = cmdl['-pdb']
            m = Model().read(fname_pdb)
        else:
            print('Structure file (pdb or xyz) is needed when using -esp option')
            sys.exit(0)
        # find halogens
        (halogens, halogens_nn) = find_halogens( itp=False,m=m, bSH=bSH )
        # add sigmahole to pdb
        fname_pdb,sigmaholes = add_sigmahole_pdb( halogens, halogens_nn, m, randnum, rule, ff=ff, scaleD=scaleD )
        # create .ac file with sigmahole
        fname_ac = antechamber_ac( fname_pdb, m, halogens, sigmaholes, charge, randnum, ff )
        if cmdl.opt['-log'].is_set:
            # create esp with sigmahole
            fname_esp = generate_esp( cmdl['-log'], halogens, halogens_nn, m, randnum, rule, ff, scaleD=scaleD )
        else:
            fname_esp = generate_esp( cmdl['-esp'], halogens, halogens_nn, m, randnum, rule, ff, scaleD=scaleD, bESP=True )
#            fname_esp = cmdl['-esp']
        # perform resp fit
        fname_resp = generate_respfiles( fname_ac, fname_esp, randnum )
        # generate .ac with resp charges
        fname_ac_resp = antechamber_ac_resp( fname_ac, fname_resp, halogens, halogens_nn, sigmaholes, randnum, ff, charge )
        # generate frcmod
        fname_frcmod = generate_frcmod( fname_ac_resp, halogens, halogens_nn, randnum, rule, ff, scaleD=scaleD )
        # get mol2
        fname_mol2 = generate_mol2( fname_ac_resp, halogens, halogens_nn, sigmaholes, randnum, ff, charge )
        # run teLeap
        fname_prmtop,fname_inpcrd = run_teleap( fname_mol2, fname_frcmod, randnum, ff )
        # run acpype
        fname_top,fname_gro = run_acpype( fname_prmtop, fname_inpcrd, charge, randnum, ff )
        # top2itp
        itp = gen_itp( fname_top, fname_gro, randnum )
        # vsites and additional bonds for exclusions
        add_vsite_respfit( ff, itp, halogens, halogens_nn, sigmaholes, rule, scaleD=scaleD )
        # and additional bonds for exclusions
        add_sh_hnn_bond_respfit( ff, itp, halogens, halogens_nn, sigmaholes )
        # generate pairs as if sigma hole is the same atom as halogen (also remove unnecessary pairs)
        add_pairs_respfit( itp, halogens, sigmaholes )
        # remove any angles and dihedrals that might have been generated for the dummy vsites
        remove_angles_dihedrals( itp )
        # TODO: LJ radii (if RESP to be used with charmm)
        # write ffitp
        write_ff( itp.atomtypes, cmdl['-offitp'] )
        # remove atomtypes from .itp (they are already written to ffitp)
        itp.atomtypes = []
        # clean
        if bClean==True:
            cmd = "rm "+fname_pdb+" "+fname_ac+" "+fname_esp+" "+fname_resp+" "+fname_ac_resp+" "+fname_frcmod+" "+fname_mol2+" "+fname_prmtop+" "+fname_inpcrd+" "+fname_top+" ANTECHAMBER* ATOMTYPE* leap.in leap.log"
            os.system(cmd)
    else:
        print('Need to provide -itp or -log')
        sys.exit(0)
#        print halogens[0].name
#        print halogens_nn[0].name
#        sys.exit(0)

    if len(cmdl['-ligname'])>0:
        set_ligname( cmdl['-ligname'], m, itp )

    m.write(cmdl['-opdb'])
    itp.write( cmdl['-oitp'], stateBonded='A' )

main( sys.argv )
