#!/usr/bin/env python
"""
Script for converting structures in xyz format to SQM compatible input.

Author: Brian Radak <radakb@biomaps.rutgers.edu>
"""
import sys
import os
import re
import optparse
import string

DEFAULT_HEADER=(' &qmmm\n'
                '  qm_theory=\'AM1\', qmcharge = 0, maxcyc = 0,\n'
                '  tight_p_conv = 1, scfconv = 1.0e-10, pseudo_diag = 0, '
                'errconv = 1.0e-10\n'
                ' /\n')

def get_atomic_number(atomic_symbol):
    """
    Return the atomic number given an atomic symbol. 

    Since this may fail, say because the symbol is not an element but an MM 
    atom type, then it is checked against a list of known MM types. If that 
    fails, characters will be stripped off the right until an element is 
    recognized (e.g. 'CT' -> 'C').  Total failure of all of this results in an
    atomic number of 0.
    """
    # This is not complete, I got bored halfway through typing all of these
    # out. In practice I only use groups I-III anyway. Besides, very few NDDO
    # methods get that far.
    periodic_table = {'H': 1, 'He': 2, 'Li': 3, 'Be': 4, 'B': 5, 'C': 6, 
                      'N': 7, 'O': 8, 'F': 9, 'Ne': 10, 'Na': 11, 'Mg': 12, 
                      'Al': 13, 'Si': 14, 'P': 15, 'S': 16, 'Cl': 17, 'Ar': 18,
                      'K': 19, 'Ca': 20, 'Sc': 21, 'Ti': 22, 'V': 23, 'Cr': 24,
                      'Mn': 25, 'Fe': 26, 'Co': 27, 'Ni': 28, 'Cu': 29, 
                      'Zn': 30, 'Ga': 31, 'Ge': 32, 'As': 33, 'Se': 34, 
                      'Br': 35, 'Kr': 36, 'Rb': 37, 'Sr': 38, 'Y': 39,
                      'Zr': 40, 'Nb': 41, 'Mo': 42, 'Tc': 43, 'Ru': 44, 
                      'Rh': 45, 'Pd': 46, 'Ag': 47, 'Cd': 48, 'In': 49, 
                      'Sn': 50, 'Sb': 51, 'Te': 52,'I': 53, 'Xe': 54, 'Cs': 55,
                      'Ba': 56, 'La': 57
                      }
    mm_types = {'SOD' : 11, 'CLA' : 17} 

    atomic_symbol = str(atomic_symbol)
    # Check if the symbol is simply an element (capitalize first letter)
    try:
        return int(periodic_table[atomic_symbol.capitalize()])
    except KeyError:
        pass
    # Failing that, check if it is a known MM type (use all caps). These is
    # really just a hack to work with strange types encountered in CHARMM.
    try:
        return int(mm_types[atomic_symbol.upper()])
    except KeyError:
        pass
    # As a last ditch effort, strip off the end characters until it is either
    # recognized as an element or none are left. In the latter case, return 0.
    new_len = len(atomic_symbol)
    atomic_symbol = atomic_symbol.capitalize()
    while new_len > 0 and atomic_symbol not in periodic_table.keys():
        atomic_symbol = atomic_symbol[:new_len].capitalize()
        new_len -= 1
    try:
        return int(periodic_table[atomic_symbol])
    except KeyError:
        return 0

usage = ('usage: %prog xyzfilename [sqmfilename]\n'
         'xyzfilename - xyz file to convert\n'
         'sqmfilename - sqm file to write (default: use .xyz prefix + .inp)')
parser = optparse.OptionParser(usage)
parser.add_option('-q','--charge',type=int,dest='charge',default=None,
                  help='system charge, may be read from xyz file')
parser.add_option('--header',dest='header',default=DEFAULT_HEADER,
                  help='file containt an sqm header')
parser.add_option('--title',dest='title',default=None,
                  help='title to be set in sqm input (default: use xyz title)')

(opt,args) = parser.parse_args()
argc = len(args)
if argc == 0 or argc > 2:
    parser.print_help()
    print
    print ('Keywords recognized in .xyz title lines:\n'
           'charge=<charge> - set the system charge to <charge>\n'
           )
    sys.exit(0)
if argc > 0:
    xyzfilename = args[0]
if argc > 1:
    sqmfilename = args[1]
else:
    sqmfilename = '%s.inp'%os.path.splitext(xyzfilename)[0]

#
# Read the xyz file.
#
xyz = open(xyzfilename,'r')
natoms = int(xyz.readline().strip().split()[0])
title = xyz.readline().strip()
q = opt.charge
if q is None:
    m = re.search('charge *= *([0-9\-]+)',title)
    if m is not None:
        q = int(m.group(1))
        title = title.replace(m.group(0),'')
    else:
        q = 0
        print 'Warning! No charge on command line or in xyz file, assuming 0.'
anames = ['' for n in range(natoms)]
coords = [0. for i in range(3*natoms)]
for n in range(natoms):
    tokens = xyz.readline().strip().split()
    anames[n] = tokens[0]
    coords[3*n+0] = float(tokens[1])
    coords[3*n+1] = float(tokens[2])
    coords[3*n+2] = float(tokens[3])
xyz.close()

#
# Write the sqm file.
#
if os.path.exists('%s/%s'%(os.getcwd(),opt.header)):
    headerlines = open(opt.header,'r').readlines()
    header = ''.join(headerlines)
else:
    header = opt.header
header = header.replace('qmcharge = 0','qmcharge = %d'%q)
sqm_input = open(sqmfilename,'w')
if opt.title is not None:
    title = opt.title
sqm_input.write(title + '\n')
sqm_input.write(header)
for n in range(natoms):
    anum,aname = get_atomic_number(anames[n]),anames[n]
    x,y,z = coords[3*n],coords[3*n+1],coords[3*n+2]
    sqm_input.write('%2d %-4s %12.7f %12.7f %12.7f\n'%(anum,aname,x,y,z))
sqm_input.close()