coul_tt

#!/usr/bin/env python

import argparse

def ReadData(dfile):
    file = open(dfile, 'r')
    atoms = list()
    line = file.readline()
    while line :
        if 'Masses' in line:
            break
        line = file.readline()
    line = file.readline()
    line = file.readline()
    while line :
        if 'Bond Coeffs' in line:
            del atoms[-1]
            break
        atoms.append(line.split())
        line = file.readline()

    cores = list()

    for a in atoms:
        if 'DC' in a[-1]:
            cores.append(int(a[0]))
        elif 'DP' in a[-1]:
            dp = int(a[0])
            break
    
    return (cores, dp)

def WritePairFile(pfile, hatoms, cores, dp):
    res = []
    res.append('# Tang-Toennies charge-dipole damping')

    for h in hatoms:
        for c in cores:
            if int(h) < c:
                res.append('pair_coeff {0:4d} {1:4d} coul/tt 4.5 1.0'.format(int(h), c))
            else:
                res.append('pair_coeff {0:4d} {1:4d} coul/tt 4.5 1.0'.format(c, int(h)))
        res.append('pair_coeff {0:4d} {1:3d}* coul/tt 4.5 1.0'.format(int(h), dp))

    with open(pfile, 'w') as f:
        for line in res:
            f.write(line+'\n')

def WriteInputFile(pfile):
    print('To inlcude in in-p.lmp:')
    print('    pair_style hybrid/overlay [...] coul/tt 4 12.0')
    print('    include {0}'.format(pfile))

def main():
    parser = argparse.ArgumentParser( description = 'Automatic coul/tt pair style generator')
    parser.add_argument('-d', '--dfile', dest='dfile', default='data-p.lmp', type=str, help='Data file with atomic indexes and labels of polarisable system [default: %(default)s]')
    parser.add_argument('-a', '--hatoms', dest='hatoms', type=str, nargs='+', help='Atomic indexes of naked charge hydrogen atoms')
    args = parser.parse_args()

    pfile = "pair-tt.lmp"

    (cores,dp) = ReadData(args.dfile)
    WritePairFile(pfile, args.hatoms, cores, dp)
    WriteInputFile(pfile)


if __name__ == '__main__':
    main()