import chimera
import sys
import re
import math
from math import sin,cos
import Midas
from Midas.midas_text import makeCommand
def printf( format, *args):
  sys.stdout.write( format % args )
opened = chimera.openModels.open( 'green.pdb', type= "PDB" )
mol1 = opened[0]
mol1.display = True
ATOMS1 = mol1.atoms
count = 0;
xc = 0.0;
yc = 0.0;
zc = 0.0;
vertebrae = 0;
xbc = 0.0;
ybc = 0.0;
zbc = 0.0;
MAINCHAIN = re.compile( "^(N|CA|C)$", re.I )

for a in ATOMS1:
  coord = a.xformCoord(); 
 # printf( '%3d %6.1f %6.1f %6.1f\n',  count, coord.x, coord.y, coord.z )
  xc = xc + coord.x
  yc = yc + coord.y
  zc = zc + coord.z
  if MAINCHAIN.match(a.name):
  	xbc = xbc + coord.x
  	ybc = ybc + coord.y
  	zbc = zbc + coord.z
	vertebrae = vertebrae + 1
  count = count + 1

xc = xc / count;
yc = yc / count;
zc = zc / count;
xbc = xbc / vertebrae;
ybc = ybc / vertebrae;
zbc = zbc / vertebrae;

printf( '%d atoms    centroid at ( %8.2f, %8.2f, %8.2f)\n', count, xc, yc, zc );
printf( '%d vertebra centroid at ( %8.2f, %8.2f, %8.2f)\n', vertebrae, xbc, ybc, zbc );

best = 1000000
d2r = 3.1417 / 180
#for theta in range( 180 ):
#  rtheta = d2r * theta
#  for phi in range( 180 ):
#    rphi = d2r * phi
#    xx = sin( rtheta ) * cos( rphi )
#    yy = sin( rtheta ) * sin( rphi )
#    zz = cos( rtheta )
#    cax = xc + 20 * xx
#    cay = yc + 20 * yy
#    caz = zc + 20 * zz
#    dax = xc - 20 * xx
#    day = yc - 20 * yy
#    daz = zc - 20 * zz
#    now = 0.0
##   printf( ' %4d %4d %6.2f %6.2f %6.2f   %6.4f %6.4f %6.4f\n', theta, phi, cax,cay,caz,dax,day,daz )
#    for a in ATOMS1:
#        coord = a.xformCoord()
#        xx = coord.x
#        yy = coord.y
#        zz = coord.z
#        l1 = sqrt( (cax-xx)**2 + (cay-yy)**2 + (caz-zz)**2 )
#        l2 = sqrt( (dax-xx)**2 + (day-yy)**2 + (daz-zz)**2 )
#        l3 = sqrt( (cax-dax)**2 + (cay-day)**2 + (caz-daz)**2 )
#        t = .5 + .5 * ( l1*l1 - l2*l2 ) / (l3*l3)
#   #     printf( ' %6.2f  %6.2f %6.2f %6.2f %8.4f\n', t, l1, l2, l3, l2*l2 - t*t*l3*l3 )
#        l = sqrt( l1*l1 - t*t*l3*l3 )
#        now = now + l;
#    if now < best:
#        best = now
#        printf( ' %4d %4d %10.2f\n', theta, phi, best );
    
theta = 125
phi   = 101
rtheta = d2r * theta
rphi = d2r * phi
xx = sin( rtheta ) * cos( rphi )
yy = sin( rtheta ) * sin( rphi )
zz = cos( rtheta )


makeCommand( 'select all' )
makeCommand( 'chain @ca,c,n' )
makeCommand( 'color byelement' )
makeCommand( '~select' )
f = open( "green.bild", 'w' )
f.write( ".color red\n" )
f.write( ".sphere %6.2f %6.2f %6.2f 1\n" % ( xc, yc, zc ) );
f.write( ".color green\n" )
f.write( ".sphere %6.2f %6.2f %6.2f 1\n" % ( xbc, ybc, zbc ) );
f.write( ".color red\n" )
f.write( ".arrow %6.2f %6.2f %6.2f %6.2f %6.2f %6.2f  .2\n" % ( xc-30*xx,yc-30*yy,zc-30*zz,xc+30*xx,yc+30*yy,zc+30*zz) )
f.write( ".color green\n" )
f.write( ".cylinder %6.2f %6.2f %6.2f %6.2f %6.2f %6.2f  15 open\n" % ( xc-30*xx,yc-30*yy,zc-30*zz,xc-31*xx,yc-31*yy,zc-31*zz) )
f.write( ".cylinder %6.2f %6.2f %6.2f %6.2f %6.2f %6.2f  15 open\n" % ( xc+30*xx,yc+30*yy,zc+30*zz,xc+31*xx,yc+31*yy,zc+31*zz) )
f.close()
makeCommand( 'open green.bild' )