from pymol import cmd
from pymol.cgo import *
import string
import math

print 'Enhanced by PyDeT Delaunay Tessalator'

def __init__(self):
        cmd.extend("pydet",PyDeT)
        #cmd.extend("pydetdextract",PyDeT_DExtract)
        #cmd.extend("pydetqdelaunay",PyDeT_QDelaunay)
        #cmd.extend("pydettssvisualise",PyDeT_TessVisualise)

def PyDeT(selection='all',name=''):
	infile="pydelaunay.input"
	outfile="pydelaunay.output"
	statfile="pydelaunay.stat"

	qdelaunay= '"c:\Program Files\qhull\qdelaunay.exe"' #For windows enviorment
	#qdelaunay= '~/bin/./qdelaunay' #For Linux enviorment

	PyDeT_DExtract(infile,selection)
	PyDeT_QDelaunay(qdelaunay,infile,outfile,statfile)
	PyDeT_TessVisualise(outfile,selection,name)

	print ' Returning statistics '

	fpin = open(statfile)
	print '' 
	print fpin.read()

def PyDeT_DExtract(infile, selection='all'):
	print ' Selection in use: '+selection
	print ' Extracting data for qdelaunay into file: '+infile
	print ' -> Collecting data'
	model = cmd.get_model(selection)
	li='3\n'
	li=li+'%s\n'%cmd.count_atoms(selection)
	
	for a in model.atom:
	   li=li+'%s %s %s\n'%(a.coord[0],a.coord[1],a.coord[2])
	
	print ' -> Saveing data'
	fpout = open(infile, 'w')
	fpout.write(li)
	fpout.close()
	
	print ' -> Data extracted'
	
def PyDeT_QDelaunay(qdelaunay,infile,outfile,statfile='', options='-Qt'):
	if statfile=='':
	   print ' Starting qdelaunay! Geometry output is: '+outfile
	   os.system(qdelaunay+' -i'+options+' < '+infile+' > '+outfile)
	else:
	   print ' Starting qdelaunay! Geometry output is: '+outfile+' and stat output is: '+statfile
	   os.system(qdelaunay+' -s -i'+options+' < '+infile+' > '+outfile+' 2> '+statfile)
	
def PyDeT_TessVisualise(outfile,selection,name=''):
	print ' Visualise the Delaunay Tessalation (Edges)'
	print ' -> Reading file '+outfile+' and Calculating distances'
	
	model = cmd.get_model(selection)
	fpin = open(outfile)
	print '    Number of regions: '+fpin.readline()+''
	
	edgeblocker=len(model.atom)*[0]
	for i in range(len(edgeblocker)):
	    edgeblocker[i]=len(model.atom)*[0]

	edgelist=[]
	
	a = model.atom[0]
	b = model.atom[1]
	minco=math.sqrt((a.coord[0]-b.coord[0])**2+(a.coord[1]-b.coord[1])**2+(a.coord[2]-b.coord[2])**2)
	maxco=0
	for line in fpin.readlines():

		line=line.strip()
		corner=line.split(' ')

		for i in range(len(corner)):
		   for j in range(len(corner)):
		      ii=eval(corner[i])
		      jj=eval(corner[j])
	
		      if ii<jj:	       
		       if edgeblocker[ii][jj]==0:
			edgeblocker[ii][jj]=1
			edgeblocker[jj][ii]=1
	
			a = model.atom[ii]
			b = model.atom[jj]
			co=math.sqrt((a.coord[0]-b.coord[0])**2+(a.coord[1]-b.coord[1])**2+(a.coord[2]-b.coord[2])**2)
			minco=min(co,minco)
			maxco=max(co,maxco)			
			edgelist.extend([[a.coord[0],a.coord[1],a.coord[2],b.coord[0],b.coord[1],b.coord[2],co]])
	
	fpin.close()
	print ' -> Generating CGO list'
	
	difco=maxco-minco

	obj=[]
	for e in edgelist:
		co=((e[6]-minco)/difco)**(0.75)
	
		co_r=max(min(1-2*co,1),0)
		co_g=max(min(1-abs(2*co-1),1),0)
		co_b=max(min(2*co-1,1),0)
		obj.extend([ CYLINDER, e[0],e[1], e[2], e[3],e[4], e[5],0.05,co_r,co_g,co_b,co_r,co_g,co_b ])
	
	print ' -> Create CGO objects '
	if name=='':
		cmd.load_cgo(obj,'PyDeT')
	else:
		cmd.load_cgo(obj,'PyDeT_'+name)