#!/source/python/IRIX-6.4/bin/python

# Orbviewer: plot jaguar orbitals using VTK toolkit and python
# Copyright 2000, Richard P. Muller and William A. Goddard, III
# All rights reserved.

from libVTKCommonPython import *
from libVTKGraphicsPython import *
import string, os, sys, re

def check_param(param,min,max):
    if param < min:
        param = min
    elif param > max:
        param = max
    return param

def help():
    print 'Orbviewer.py [options] [plt file]'
    print 'Plot out Jaguar orbitals.'
    print ''
    print 'Orbviewer reads the *.plt file generated by Jaguar and plots the '
    print 'orbitals found therein. It takes the geometry from the *.out file '
    print 'and thus must be able to find that. For the file '
    print '\'orb.01_a03MO.plt\' the program first searches for '
    print '\'orb.01_a03MO.out\', and then searches for \'orb.01.out\'.'
    print ''
    print 'Options:'
    print '-c  #  Set contour value to # (range [0.0001,0.1]).'
    print '-b  #  Set background grey-level to # (range [0,1]).'
    print '-o  #  Set opacity to # (range [0.2,1]).'
    print '-C     View molecule using cylinders rather than ball-and-stick.'
    print '-G     Only print geometry, no orbitals.'
    print '-L     View molecule using lines.'
    print '-V     Dump VRML output of image.'
    sys.exit()

def bond_center(atomi,atomj):
    atnoi,xi,yi,zi = atomi
    atnoj,xj,yj,zj = atomj
    cx = 0.5*(xi+xj) #center of bond
    cy = 0.5*(yi+yj)
    cz = 0.5*(zi+zj)
    return (cx,cy,cz)

def get_color(atomi):
    atno,x,y,z = atomi
    r = rint[atno]
    g = gint[atno]
    b = bint[atno]
    return (r,g,b)

def bond_center_mid(atomi,atomj):
    atnoi,xi,yi,zi = atomi
    atnoj,xj,yj,zj = atomj
    cx1 = xi + 0.25*(xj-xi)
    cx2 = xi + 0.75*(xj-xi)
    cy1 = yi + 0.25*(yj-yi)
    cy2 = yi + 0.75*(yj-yi)
    cz1 = zi + 0.25*(zj-zi)
    cz2 = zi + 0.75*(zj-zi)
    return ((cx1,cy1,cz1),(cx2,cy2,cz2))

def bond_unit_vector(atomi,atomj):
    atnoi,xi,yi,zi = atomi
    atnoj,xj,yj,zj = atomj
    ax = xj-xi       #bond axis
    ay = yj-yi
    az = zj-zi
    al = sqrt(ax*ax + ay*ay + az*az)
    ax = ax/al
    ay = ay/al
    az = az/al
    return (al,ax,ay,az)    

def getrot(ax,ay,az):
    from math import acos,pi
    rad2deg = 180./pi
    # Compute the rotational axis (y cross a)
    rx = az
    ry = 0
    rz = -ax
    # Compute the rotational angle (y dot a)
    theta = rad2deg*acos(ay)
    return (theta,rx,ry,rz)

def namesplit(name):
    words = string.split(name,'.')
    ext = words[-1]
    root = words[0]
    for word in words[1:-1]:
        root = root + '.' + word
    return (root,ext)

def altname(name):
    words = string.split(name,'_')
    alt = words[0]
    for word in words[1:-1]:
        alt = alt + '_' + word
    alt = alt + '.out'
    return alt

def get_bonds_by_distance(geo):
    from math import sqrt
    bondlist = []
    for atomi in geo:
        for atomj in geo:
            if atomi < atomj:
                atnoi,xi,yi,zi = atomi
                atnoj,xj,yj,zj = atomj
                r = (xi-xj)*(xi-xj) + \
                    (yi-yj)*(yi-yj) + \
                    (zi-zj)*(zi-zj)
                r = sqrt(r)
                r_model = 0.5*(rcov[atnoi] + rcov[atnoj])
                if r < r_model:
                    bondlist.append((atomi,atomj))                
    return bondlist

def getgeo(geofilename):
    geopat = re.compile('[new|Input|Symmetrized] geometry')
    file = open(geofilename,'r')
    while 1:
        line = file.readline()
        if not line: break
        if geopat.search(line):
            geo = getonegeo(file)
    file.close()
    return geo

def getonegeo(file):
    geo = []
    line = file.readline()
    line = file.readline()
    while 1:
        line = file.readline()
        if not line:break
        words = string.split(line)
        if len(words) < 4: break
        sym,x,y,z = words
        x = eval(x)
        y = eval(y)
        z = eval(z)
        sym = cleansym(sym)
        atno = sym2no[sym]
        geo.append((atno,x,y,z))
    return geo

def cleansym(sym):
    pat = re.compile('[^a-zA-Z]')
    newsym = pat.split(sym)[0]
    return newsym

######### DATA SECTION########

rint = [
    0.973,
    0.973, 0.944,
    0.698, 0.675, 0.000, 0.184, 0.000, 1.000, 0.933, 0.850,
    0.941, 0.133, 0.184, 0.933, 1.000, 0.678, 0.000, 0.005,
    0.941, 0.133, 0.184, 0.933, 1.000, 0.678, 0.000, 0.005,
    0.941, 0.133, 0.184, 0.933, 1.000, 0.678, 0.000, 0.005,
    0.941, 0.133, 0.184, 0.933, 1.000, 0.678, 0.000, 0.005,
    0.941, 0.133, 0.184, 0.933, 1.000, 0.678, 0.000, 0.005
    ]

gint = [
    0.844,
    0.973, 0.900,
    0.133, 0.133, 0.392, 0.310, 0.000, 0.000, 0.910, 0.777,
    0.973, 0.545, 0.310, 0.910, 0.647, 1.000, 0.392, 0.500,
    0.973, 0.545, 0.310, 0.910, 0.647, 1.000, 0.392, 0.500,
    0.973, 0.545, 0.310, 0.910, 0.647, 1.000, 0.392, 0.500,
    0.973, 0.545, 0.310, 0.910, 0.647, 1.000, 0.392, 0.500,
    0.973, 0.545, 0.310, 0.910, 0.647, 1.000, 0.392, 0.500
    ]

bint = [
    1.0,
    1.0, 1.0,
    0.000, 0.000, 0.000, 0.310, 1.000, 0.000, 0.667, 0.600,
    1.000, 0.133, 0.310, 0.310, 0.667, 0.000, 0.184, 0.200,
    1.000, 0.133, 0.310, 0.310, 0.667, 0.000, 0.184, 0.200,
    1.000, 0.133, 0.310, 0.310, 0.667, 0.000, 0.184, 0.200,
    1.000, 0.133, 0.310, 0.310, 0.667, 0.000, 0.184, 0.200,
    1.000, 0.133, 0.310, 0.310, 0.667, 0.000, 0.184, 0.200
    ]

rcov = [
    1.25,
    1.25, 1.25,
    1.86, 1.86, 1.86, 1.65, 1.50, 1.35, 1.40, 1.40,
    2.2, 2.2, 2.2, 2.2, 2.2, 2.2, 2.2, 2.2,
    2.2, 2.2, 2.2, 2.2, 2.2, 2.2, 2.2, 2.2,
    2.2, 2.2, 2.2, 2.2, 3.1, 3.1, 3.1, 3.1,
    3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1,
    3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1,
    3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1,
    3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1,
    3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1,
    3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1,
    3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1, 3.1
    ]

rvdw = [
    0.24,
    0.24, 0.24,
    0.93, 0.93, 0.62, 0.58, 0.56, 0.55, 0.54, 0.54,
    1.16, 1.02, 0.89, 0.87, 0.80, 0.77, 0.74, 0.74,
    1.16, 1.02, 0.89, 0.87, 0.80, 0.77, 0.74, 0.74,
    1.16, 1.02, 0.89, 0.87, 0.80, 0.77, 0.74, 0.74,
    1.16, 1.02, 0.89, 0.87, 0.80, 0.77, 0.74, 0.74,
    1.16, 1.02, 0.89, 0.87, 0.80, 0.77, 0.74, 0.74
    ]

sym2no = {
    'X' : 0,
    '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,
    'Fe' : 26,
    'Ni' : 28,
    'Ga' : 31,
    'As' : 33,
    'Mo' : 42,
    'Ru' : 44,
    'Ag' : 47,
    'In' : 49,
    'x' : 0, 
    '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,
    'fe' : 26,
    'ni' : 28,
    'ga' : 31,
    'as' : 33,
    'mo' : 42,
    'ru' : 44,
    'ag' : 47, 
    'in' : 49}
    
##############################
from math import sqrt
import getopt

if len(sys.argv) < 2:
    help()

opts,args = getopt.getopt(sys.argv[1:],'ho:c:b:GCLV')

# Read the file
name = args[0]

opacity = 0.7
contourvalue = 0.05
background = 1.0
spheres = 1
cylinders=1
cylinder_color = 0
lines = 0
line_color = 0
orbs = 1
vrml_output = 0
for opt in opts:
    key,value = opt
    if key == '-h':
        help()
    elif key == '-o':
        opacity = eval(value)
    elif key == '-c':
        contourvalue = eval(value)
    elif key == '-b':
        background = eval(value)
    elif key == '-G':
        orbs = 0
    elif key == '-C':
        spheres = 0
        cylinder_color = 1
    elif key == '-L':
        lines = 1
        cylinders = 0
        spheres = 0
    elif key == '-V':
        vrml_output = 1

background = check_param(background,0,1)
opacity = check_param(opacity,0.2,1)
contourvalue = check_param(contourvalue,0.0001,0.1)

words = string.split(name,'.')

root, ext = namesplit(name)
    
geoname = root + '.out'
alt_geoname = altname(root)

if vrml_output == 1:
    vrmlname = root + '.wrl'

if os.path.exists(geoname):
    geo = getgeo(geoname)
elif os.path.exists(alt_geoname):
    geo = getgeo(alt_geoname)
else:
    print 'Can\'t find input file ',geoname,' or ',alt_geoname
    sys.exit(1)

bondlist = get_bonds_by_distance(geo)

# Initialize VTK stuff
ren = vtkRenderer()
ren.SetBackground(background,background,background)
renWin = vtkRenderWindow()
renWin.AddRenderer(ren)
renWin.SetSize(500,500)
iren = vtkRenderWindowInteractor()
iren.SetRenderWindow(renWin)

slist = []
mlist = []
alist = []

if spheres:
    for atom in geo:
        atno,x,y,z = atom
        red = rint[atno]
        green = gint[atno]
        blue = bint[atno]
        radius = rvdw[atno]
        sphere = vtkSphereSource()
        sphere.SetRadius(radius)
        sphere.SetThetaResolution(15)
        sphere.SetPhiResolution(15)
        slist.append(sphere)
        mapper = vtkPolyDataMapper()
        mapper.SetInput(sphere.GetOutput())
        mlist.append(mapper)
        actor = vtkActor()
        actor.SetMapper(mapper)
        actor.GetProperty().SetColor(red,green,blue)
        actor.SetPosition(x,y,z)
        ren.AddActor(actor)
        alist.append(actor)

if cylinders:
    for bond in bondlist:
        atomi, atomj = bond
        if cylinder_color:
            r1,g1,b1 = get_color(atomi)
            r2,g2,b2 = get_color(atomj)
            c1,c2 = bond_center_mid(atomi,atomj)
            cx1,cy1,cz1 = c1
            cx2,cy2,cz2 = c2
            al,ax,ay,az = bond_unit_vector(atomi,atomj)
            theta,rx,ry,rz = getrot(ax,ay,az)
            cyl1 = vtkCylinderSource()
            cyl2 = vtkCylinderSource()
            cyl1.SetRadius(0.08)
            cyl2.SetRadius(0.08)
            cyl1.SetHeight(al/2)
            cyl2.SetHeight(al/2)
            slist.append(cyl1)
            slist.append(cyl2)
            mapper1 = vtkPolyDataMapper()
            mapper2 = vtkPolyDataMapper()
            mapper1.SetInput(cyl1.GetOutput())
            mapper2.SetInput(cyl2.GetOutput())
            mlist.append(mapper1)
            mlist.append(mapper2)
            actor1 = vtkActor()
            actor2 = vtkActor()
            actor1.SetMapper(mapper1)
            actor2.SetMapper(mapper2)
            actor1.GetProperty().SetColor(r1,g1,b1)
            actor2.GetProperty().SetColor(r2,g2,b2)
            actor1.RotateWXYZ(theta,rx,ry,rz)
            actor2.RotateWXYZ(theta,rx,ry,rz)
            actor1.SetPosition(cx1,cy1,cz1)
            actor2.SetPosition(cx2,cy2,cz2)
            ren.AddActor(actor1)
            ren.AddActor(actor2)
            alist.append(actor1)
            alist.append(actor2)
        else:
            cx,cy,cz = bond_center(atomi,atomj)
            al,ax,ay,az = bond_unit_vector(atomi,atomj)
            theta,rx,ry,rz = getrot(ax,ay,az)
            cyl = vtkCylinderSource()
            cyl.SetRadius(0.08)
            cyl.SetHeight(al)
            slist.append(cyl)
            mapper = vtkPolyDataMapper()
            mapper.SetInput(cyl.GetOutput())
            mlist.append(mapper)
            actor = vtkActor()
            actor.SetMapper(mapper)
            actor.GetProperty().SetColor(0.8,0.8,0.8)
            actor.RotateWXYZ(theta,rx,ry,rz)
            actor.SetPosition(cx,cy,cz)
            ren.AddActor(actor)
            alist.append(actor)
            

endpat = re.compile('\&end')
origpat = re.compile('origin=')
expat = re.compile('extentx=')
eypat = re.compile('extenty=')
ezpat = re.compile('extentz=')
npat = re.compile('npts=')

file = open(name,'r')
while 1:
    line = file.readline()
    if not line:
        break
    elif endpat.search(line):
        break
    elif origpat.search(line):
        words = string.split(line)
        x = eval(words[1])
        y = eval(words[2])
        z = eval(words[3])
        origin = (0.529*x,0.529*y,0.529*z) #convert from Bohr to Angstroms
    elif expat.search(line):
        words = string.split(line)
        xlength = 0.529*eval(words[1]) #convert from Bohr to Angstroms
    elif eypat.search(line):
        words = string.split(line)
        ylength = 0.529*eval(words[2]) #convert from Bohr to Angstroms
    elif ezpat.search(line):
        words = string.split(line)
        zlength = 0.529*eval(words[3]) #convert from Bohr to Angstroms
    elif npat.search(line):
        words = string.split(line)
        x = eval(words[1])
        y = eval(words[2])
        z = eval(words[3])
        npts = (x,y,z)

data_array = vtkFloatScalars(npts[0]*npts[1]*npts[2])
for i in range(npts[0]):
    for j in range(npts[1]):
        joffset = j*npts[0]
        for k in range(npts[2]):
            koffset = k*npts[0]*npts[1]
            offset = i + joffset + koffset
            line = file.readline()
            value = eval(string.split(line)[0])
            data_array.InsertScalar(offset,value)
file.close()

xspacing = xlength/(npts[0]-1)
yspacing = ylength/(npts[1]-1)
zspacing = zlength/(npts[2]-1)

data = vtkStructuredPoints()
data.SetDimensions(npts[0],npts[1],npts[2])
data.SetOrigin(origin[0],origin[1],origin[2])
data.SetSpacing(xspacing,yspacing,zspacing)
data.GetPointData().SetScalars(data_array)
del data_array

# Add positive contours:
contour = vtkContourFilter()
contour.SetInput(data)
contour.SetValue(0,contourvalue)
slist.append(contour)
cmap = vtkPolyDataMapper()
cmap.SetInput(contour.GetOutput())
cmap.ScalarVisibilityOff()
mlist.append(cmap)
cact = vtkActor()
cact.SetMapper(cmap)
cact.GetProperty().SetColor(1,0,0)
cact.GetProperty().SetOpacity(opacity)
cact.GetProperty().SetInterpolation(100)
cact.GetProperty().SetInterpolationToPhong()
ren.AddActor(cact)
alist.append(cact)

# Add negative contours:
contourvalue = -contourvalue
contour = vtkContourFilter()
contour.SetInput(data)
contour.SetValue(0,contourvalue)
slist.append(contour)
cmap = vtkPolyDataMapper()
cmap.SetInput(contour.GetOutput())
cmap.ScalarVisibilityOff()
mlist.append(cmap)
cact = vtkActor()
cact.SetMapper(cmap)
cact.GetProperty().SetColor(0,1,0)
cact.GetProperty().SetOpacity(opacity)
cact.GetProperty().SetInterpolation(100)
cact.GetProperty().SetInterpolationToPhong()
ren.AddActor(cact)
alist.append(cact)

if vrml_output == 1:
    vrml = vtkVRMLExporter()
    vrml.SetInput(renWin)
    vrml.SetFileName(vrmlname)
    vrml.Write()

iren.Initialize()
iren.Start()