import numpy
from numpy import *
from CGLutil import vrml

# -----------------------------------------------------------------------------
# Return VRML node for a cylinder with axis specified by end points.
#
def cylinder_node(p1, p2, radius, rgb):
    transform = vrml.Transform()
    axis = p2 - p1
    length = numpy.linalg.norm(axis)
    if length > 0:
        center = .5 * (p1 + p2)
        translate = vrml.Transform(translation=tuple(center))
        transform.addChild(translate)
        rot_axis = (axis[2], 0, -axis[0])
        rot_angle = math.atan2(math.sqrt(axis[0]*axis[0] + axis[2]*axis[2]), axis[1])
        rotate = vrml.Transform(rotation=rot_axis + (rot_angle,))
        translate.addChild(rotate)
        cylinder = vrml.Cylinder(radius=radius, height=length, color=rgb,
			     top = 0, bottom = 1)
        rotate.addChild(cylinder)
    return transform

# -----------------------------------------------------------------------------
# Return VRML string that draws a set of capped tubes.
# Spheres are drawn at the end points to provide tube caps.
# The begin point of the cylinder has a flat bottom cap.
# Cylinders are drawn between consecutive pairs of points stored in the list.
def vr_tube(points, ssTypes):
    # Build dictionaries for colors and radii:
    colorDict = {}
    radiiDict = {}
    colorDict['H'] = (1.000, 0.843, 0.000)
    colorDict['S'] = (0.196, 0.804, 0.196)
    colorDict['T'] = (0.827, 0.827, 0.827)
    radiiDict['H'] = 1.5
    radiiDict['S'] = 1.0
    radiiDict['T'] = 0.0

    tubeColor = (.5, .5, 1)
    wrl = vrml.Transform()
    for k in range(len(points)/2):
        p = points[2*k+1]
        translate = vrml.Transform(translation = tuple(p))
        s = vrml.Sphere(radius = radiiDict[ssTypes[k]], color = colorDict[ssTypes[k]])
        translate.addChild(s)
        wrl.addChild(translate)
        c = cylinder_node(points[2*k], points[2*k+1], radiiDict[ssTypes[k]], colorDict[ssTypes[k]])
        wrl.addChild(c)
    return vrml.vrml(wrl)
#====================================== Code above this line to be used by students.


#TESTING:
tubePointList = []
tubeSSlist = []

tubePointList.append(array([1.0, 2.0, 2.5]))        # Coordinates of Helix start.
tubePointList.append(array([9.0, 3.0, 3.5]))        # Coordinates of Helix end.
tubePointList.append(array([3.0, 12.0, 4.5]))       # Coordinates of Strand start.
tubePointList.append(array([5.0, 7.0, 8.5]))        # Coordinates of Strand end.
tubePointList.append(array([11.0, 12.0, 12.5]))     # Coordinates of Helix start.
tubePointList.append(array([6.0, 12.0, 7.5]))       # Coordinates of Helix end.

tubeSSlist.append("H")                              # Secondary structure code for alpha Helix.
tubeSSlist.append("S")                              # Secondary structure code for beta Strand.
tubeSSlist.append("H")                              # Secondary structure code for alpha Helix.

vvrr = vr_tube(tubePointList, tubeSSlist)
vrlmModel = chimera.openModels.open(vvrr, 'VRML')[0]