#!/usr/bin/env python
# -*- coding: iso-8859-1 -*-

# turboESP v0.1, 2.6.2003, Mikael Johansson
# email: mikael.johansson@iki.fi
#
# This script uses a COSMO output from Turbomole
# to fit ESP charges for a molecule.
# It expects to find the standard out.cosmo file
#
# It uses the MMTK toolkit, available at
# http://starship.python.net/crew/hinsen/MMTK/
#
# Internally all atoms are handled as hydrogens. This is just
# to prevent the program from crashing when an atom is not
# found in the MMTK Atoms Database (located somewhere like
# /usr/lib/python2.2/site-packages/MMTK/Database/ )
#
# MMTK uses nm as length units, and elementary charges.
#

from MMTK import *
from MMTK.ChargeFit import ChargeFit, TotalChargeConstraint
import sys, string
split=string.split
from time import time

AAU =0.529177208319  # Ångström per a.u. (CODATA 1998)
nmAU=0.0529177208319 # nm per a.u.

ATOMS=[
# 1    2    3    4    5    6    7    8    9   10   11   12   13   14   15   16   17   18  <-IUPAC column nr.
'q' , #                                                                                     IUPAC row nr.
'h' ,                                                                                'he', # 1
'li','be',                                                  'b' ,'c' ,'n' ,'o' ,'f' ,'ne', # 2
'na','mg',                                                  'al','si','p' ,'s' ,'cl','ar', # 3
'k' ,'ca','sc','ti','v' ,'cr','mn','fe','co','ni','cu','zn','ga','ge','as','se','br','kr', # 4
'rb','sr','y' ,'zr','nb','mo','tc','ru','rh','pd','ag','cd','in','sn','sb','te','i' ,'xe', # 5
'cs','ba',                                                                                 # 6
          'la','ce','pr','nd','pm','sm','eu','gd','tb','dy','ho','er','tm','yb','ly',      # lanthanides
               'hf','ta','w' ,'re','os','ir','pt','au','hg','tl','pb','bi','po','at','rn', # 6 cont.
'fr','ra',                                                                                 # 7
          'ac','th','pa','u' ,'np','pu','am','cm','bk','cf','es','fm','md','no','lr',      # actinides
               'rf','db','sg','bh','hs','mt']                                              # 7 cont.



VERSION='v0.1'
DATE='2.6.2003'         

def header():
    print "\nturboESP "+VERSION+", Mikael Johansson (mikael.johansson@iki.fi) "+DATE+"\n"



def usage():
    print "Usage: "+sys.argv[0]+" [charge]"
    print "       [charge]: total charge of the system (default 0)"



def init():
    args=sys.argv
    silent=0
    if (len(args)>2):
        header()
        sys.stderr.write("Error: Wrong number of arguments\n")
        usage()
        sys.exit(1)

    charge=0.0
    if (len(args)==2):
        charge=float(args[1])

    cosmofile=open("out.cosmo",'r')

    return cosmofile, charge



def outit(cosmofile):
    cosmofile.close()



def getatomcoords(cosmofile):

    line=cosmofile.readline()
    found=0
    while line:
        if string.find(line,"$coord_rad")<>-1:
            found=1
            break
        line=cosmofile.readline()
    if not(found):
        sys.stderr.write("Error: unable to find keyword $coord_rad in '" \
                         +coordfile.name+"'\n")
        sys.exit(2)    

    column=split(cosmofile.readline())
    system=Collection()
    AN=0

# out.cosmo format (coordinates in bohr):
# #atom   x                  y                  z             element  radius [A]
#    1   0.00000000000000   0.00000000000000  -0.12621051616787  o      1.72000

    while column[0][0]<>'$':
        if column[0][0]<>'#':
            column[4]=string.lower(column[4])
            if ATOMS.count(column[4])==0:
                sys.stderr.write("Warning: Unknown atomtype found: " \
                                 +column[4]+", replaced with a hydrogen\n")
                column[4]='h'
            # hardcoded hydrogens, see comment at beginning:
            system.addObject(Atom("h", \
                position=Vector(float(column[1])*nmAU,float(column[2])*nmAU,\
                                float(column[3])*nmAU)))
            system[AN].atomtype=column[4]
            AN=AN+1
        column=split(cosmofile.readline())
    return system, AN



def getcosmopoints(cosmofile):

    line=cosmofile.readline()
    found=0
    while line:
        if string.find(line,"$segment_information")<>-1:
            found=1
            break
        line=cosmofile.readline()
    if not(found):
        sys.stderr.write("Error: unable to find keyword $segment_information in '"\
                         +coordfile.name+"'\n")
        sys.exit(2)    

# out.cosmo format (coordinates in bohr):
# potential     - solute potential on segment (A length scale)
#   n   atom              position (X, Y, Z)                   charge         area        charge/area     potential
#    1    1   -2.444138531   -0.061540084   -2.267949410    0.002655821    0.309230275    0.008588490   -0.082341202

    line=cosmofile.readline()
    points=[]
    
    while line:
        column=split(line)
        if column[0][0]=='$':
            break
        if column[0][0]<>'#':
            # pot converted from e/Å => MMTK units
            points.append((Vector(float(column[2])*nmAU,float(column[3])*nmAU,float(column[4])*nmAU), \
                           float(column[8])*10*Units.electrostatic_energy))

        line=cosmofile.readline()

    return points



########
# MAIN #
########

cosmofile, charge=init()    

system, atomcount=getatomcoords(cosmofile)
points=getcosmopoints(cosmofile)

constraints=[TotalChargeConstraint(system, charge)]
fit = ChargeFit(system, points, constraints)

print atomcount
print "# xyz-format, lengths in Ångström. Last column is the ESP-fitted charge"

for atom in system.atomList():
    print ("%3s %20.14f %20.14f %20.14f %12.6f" % \
           (atom.atomtype,atom.position()[0]*10,atom.position()[1]*10,atom.position()[2]*10,fit[atom]))

outit(cosmofile)