# Code to complete the various primitive types
# Users should never import this module directly.  All of the public types and
# functions will be explicitly imported by __init__.py

from __future__ import division
import cvisual

from cvisual import vector
import crayola
color = crayola
from math import pi
from numpy import *

NCHORDS = 20.0 # number of chords in one coil of a helix

# Scenegraph management:
# Internally, each display maintains three unordered lists of renderable
# objects. The first is the set of all non-transparent objects in world space.
# The second is the set of all transparent objects in world space. The
# transparent layer is depth-sorted and rendered from back to front each cycle.
# Objects in world space are added to a particular display by calling
# display.add_renderable( renderable). The display will automatically determine
# if the object is translucent or not. All objects are
# removed by calling display.remove_renderable( renderable). The frame class
# also manages a list of renderable objects that lie within that reference
# frame. The frame itself is a member of its display's renderable set. When
# a renderable object is within a frame, it should only be added to the frame's
# list with frame.add_renderable(renderable), and not to the set of objects in
# the parent display.
# The py_renderable class encapsulates this logic.

class py_renderable(object):
   """ This base class adds support for the visible, frame, and
   display attributes."""
   def __init__(self, keywords, _other=None):
    if _other:
        self.__display = keywords.get( 'display', _other.display)
        self.__frame = keywords.get( 'frame', _other.frame)
        self.__visible = keywords.get( 'visible', _other.visible)
    else:
        self.__display = keywords.get( 'display',
            cvisual.display.get_selected())
        self.__visible = keywords.get( 'visible', True)
        self.__frame = keywords.get('frame')
    if keywords.has_key('display'):
        del keywords['display']
    if keywords.has_key('visible'):
        del keywords['visible']
    if keywords.has_key('frame'):
        del keywords['frame']
    if self.__frame:
        if self.__frame.display != self.__display:
            raise ValueError, """Cannot initialize an object with a frame
                on a different display."""
    

   def complete_init(self):
    "For internal use only.  This function completes two-stage "
    "initialization in VPython"
    if self.__visible:
        if self.__frame:
            self.__frame.add_renderable(self)
        elif self.__display:
            self.__display.add_renderable(self)

   def set_display(self, display):
    "For internal use only.  The setter for the display property."
    if display != self.__display:
        # Check that we aren't screwing up a frame.
        if self.__frame:
            raise ValueError, """Cannot change displays when within a
                frame.  Make frame None, first."""
        if self.__display:
            self.__display.remove_renderable(self)
        self.__display.add_renderable(self)
        self.__display = display

   def get_display(self):
    "For internal use only.  The getter for the display property."
    return self.__display

   display = property( get_display, set_display)

   def get_frame(self):
    "For internal use only.  The getter for the frame property."
    return self.__frame

   # Overridden by the frame class below to add extra checks.
   def set_frame(self, frame):
    "For internal use only.  The setter for the frame property."
    if frame != self.__frame:
        if frame.display != self.__display:
            raise ValueError, "Cannot set to a frame on a different display."
        if frame and self.__frame:
            # Simply moving from one frame to another.
            self.__frame.remove_renderable(self)
            frame.add_renderable(self)
        elif frame and not self.__frame:
            # Moving into a reference frame when otherwise not in one.
            if self.__display:
                self.__display.remove_renderable(self)
            frame.add_renderable(self)
        elif not frame and self.__frame:
            # Removing from a reference frame.
            self.__frame.remove_renderable(self)
            if self.__display:
                self.__display.add_renderable(self)
        self.__frame = frame
            
   frame = property( get_frame, set_frame)

   def get_visible(self):
    "For internal use only.  The getter for the visible property."
    return self.__visible

   def set_visible(self, visible):
    "For internal use only.  The setter for the visible property."
    if visible and not self.__visible:
        if self.__frame:
            self.__frame.add_renderable(self)
        elif self.__display:
            self.__display.add_renderable(self)
    if not visible and self.__visible:
        if self.__frame:
            self.__frame.remove_renderable(self)
        elif self.__display:
            self.__display.remove_renderable(self)
    self.__visible = visible

   visible = property( get_visible, set_visible)
   

# The following generic algorithm will initialize any primitive in Visual.
# object is a reference to a Visual primitive object.
# keywords must be a dictionary of keyword:value pairs.
def process_init_args_from_keyword_dictionary( displayobject, keywords):
   """ The guts of a factory function that initializes the primitive 
   objects in Visual.  Client code should never call this function."""
   # Some objects (like curve, points, faces, and convex) need to have color
   # set before pos if color is a single tuple. But if color is a list, must
   # set up pos before assigning colors.
   colors = displayobject.display.foreground
   setcolors = True
   if keywords.has_key('color'):
      colors = keywords['color']
      del keywords['color']
      try:
         colors_length = len(colors) # fail if colors not a sequence
         if colors_length == 3 or colors_length == 4:
            displayobject.color =  colors
         else:
            setcolors = False # if colors is a sequence of multiple colors
      except:
         displayobject.color = colors # single scalar value
       
   # Add special handling for the faces object.
   # if displayobject.__class__ == faces:
   #    if keywords.has_key('pos'):
   #        displayobject.pos = keywords['pos']
   #        del keywords['pos']

   # Assign all other properties
   for key, value in keywords.iteritems():
      displayobject.__setattr__(key, value)

   # Complete the initialization of the object.
   if not setcolors: # sequence of colors
      if len(colors) == len(displayobject.pos):
         for n in range(len(displayobject.pos)):
            if len(displayobject.color[n]) == len(colors[n]):
               displayobject.color[n] = colors[n]
            elif len(displayobject.color[n]) == 4 and len(colors[n]) == 3:
               for i in range(3):
                  displayobject.color[n][i] = colors[n][i]
            else:
               raise RuntimeError, "Mismatch of pos (%i components) and color (%i components)" % \
               (len(displayobject.pos[n]), len(colors[n]))
      else:
         raise RuntimeError, "Length of pos (%i) and color (%i) sequences are not equal" % \
               (len(displayobject.pos), len(colors))
   displayobject.complete_init()

def copy_object( copy, other, keywords):
   # Copy the user-specified dict.
   copy.__dict__ = dict(other.__dict__)
   # Set the differing properties.
   for key, value in keywords.iteritems():
    copy.__setattr__(key, value)
   # Complete initialization and make this object visible.
   copy.complete_init()

################################################################################
# Complete each type.

class arrow (cvisual.arrow, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.arrow.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        copy_object( self, _other, keywords)
    else:
        cvisual.arrow.__init__(self)
        py_renderable.__init__(self, keywords)
        process_init_args_from_keyword_dictionary(self, keywords)
    if (not self.display.uniform):
       raise RuntimeError, "Do not create arrow with nonuniform axes."
   def __copy__( self, **keywords):
    return arrow(_other=self, **keywords)


class cone (cvisual.cone, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.cone.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        copy_object( self, _other, keywords)
    else:
        cvisual.cone.__init__(self)
        py_renderable.__init__(self, keywords)
        keywords["shininess"] = keywords.get( "shininess", 0.5)
        process_init_args_from_keyword_dictionary(self, keywords)
    if (not self.display.uniform):
       raise RuntimeError, "Do not create cone with nonuniform axes."
   def __copy__( self, **keywords):
    return cone(_other=self, **keywords)


class cylinder (cvisual.cylinder, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.cylinder.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        copy_object( self, _other, keywords)
    else:
        cvisual.cylinder.__init__(self)
        py_renderable.__init__(self, keywords)
        keywords["shininess"] = keywords.get( "shininess", 0.5)
        process_init_args_from_keyword_dictionary(self, keywords)
    if (not self.display.uniform):
       raise RuntimeError, "Do not create cylinder with nonuniform axes."
   def __copy__( self, **keywords):
    return cylinder(_other=self, **keywords)


class sphere (cvisual.sphere, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.sphere.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        copy_object( self, _other, keywords)
    else:
        cvisual.sphere.__init__(self)
        py_renderable.__init__(self, keywords)
        keywords["shininess"] = keywords.get( "shininess", 0.5)
        process_init_args_from_keyword_dictionary(self, keywords)
    if (not self.display.uniform):
       raise RuntimeError, "Do not create sphere with nonuniform axes."
   def __copy__( self, **keywords):
    return sphere(_other=self, **keywords)


class ring (cvisual.ring, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.ring.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        copy_object( self, _other, keywords)
    else:
        cvisual.ring.__init__(self)
        py_renderable.__init__(self, keywords)
        keywords["shininess"] = keywords.get( "shininess", 0.5)
        process_init_args_from_keyword_dictionary(self, keywords)
    if (not self.display.uniform):
       raise RuntimeError, "Do not create ring with nonuniform axes."
   def __copy__( self, **keywords):
    return ring(_other=self, **keywords)


class box (cvisual.box, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.box.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        # In the case where both axis and length are set, this ensures that
        # length is set after axis.
        if keywords.has_key('axis'):
            self.axis = keywords['axis']
            del keywords['axis']
        copy_object( self, _other, keywords)
    else:
        cvisual.box.__init__(self)
        py_renderable.__init__(self, keywords)
        keywords["shininess"] = keywords.get( "shininess", 0.5)
        if keywords.has_key('axis'):
            self.axis = keywords['axis']
            del keywords['axis']
        process_init_args_from_keyword_dictionary(self, keywords)
    if (not self.display.uniform):
       raise RuntimeError, "Do not create box with nonuniform axes."
   def __copy__( self, **keywords):
    return box(_other=self, **keywords)


class ellipsoid (cvisual.ellipsoid, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.ellipsoid.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        copy_object( self, _other, keywords)
    else:
        cvisual.ellipsoid.__init__(self)
        py_renderable.__init__(self, keywords)
        keywords["shininess"] = keywords.get( "shininess", 0.5)
        process_init_args_from_keyword_dictionary(self, keywords)
    if (not self.display.uniform):
       raise RuntimeError, "Do not create ellipsoid with nonuniform axes."
   def __copy__( self, **keywords):
    return ellipsoid(_other=self, **keywords)

class pyramid (cvisual.pyramid, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.pyramid.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        copy_object( self, _other, keywords)
    else:
        cvisual.pyramid.__init__(self)
        py_renderable.__init__(self, keywords)
        process_init_args_from_keyword_dictionary(self, keywords)
    if (not self.display.uniform):
       raise RuntimeError, "Do not create pyramid with nonuniform axes."
   def __copy__( self, **keywords):
    return pyramid(_other=self, **keywords)


class label (cvisual.label, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.label.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        copy_object( self, _other, keywords)
    else:
        cvisual.label.__init__(self)
        py_renderable.__init__(self, keywords)
        if keywords.has_key('linecolor'):
            self.linecolor = keywords['linecolor']
            del keywords['linecolor']
        else:
            if keywords.has_key('display'):
                self.linecolor = keywords['display'].foreground
            else:
                self.linecolor = cvisual.display.get_selected().foreground
        process_init_args_from_keyword_dictionary(self, keywords)
   def __copy__( self, **keywords):
    return label(_other=self, **keywords)


class frame (cvisual.frame, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.frame.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        copy_object( self, _other, keywords)
    else:
        cvisual.frame.__init__(self)
        py_renderable.__init__(self, keywords)
        process_init_args_from_keyword_dictionary(self, keywords)
    if (not self.display.uniform):
       raise RuntimeError, "Do not create frame with nonuniform axes."
   def __copy__( self, **keywords):
    return frame(_other=self, **keywords)

   def set_frame(self, frame):
    #Check to ensure that we are not establishing a cycle of reference frames.
    frame_iterator = frame.frame
    while frame_iterator:
        if frame_iterator.frame == self:
            raise ValueError, "Attempted to create a cycle of reference frames."
        frame_iterator = frame_iterator.frame
    py_renderable.set_frame( self, frame)

class curve( cvisual.curve, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.curve.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        copy_object( self, _other, keywords)
    else:
        cvisual.curve.__init__(self)
        py_renderable.__init__(self, keywords)
        process_init_args_from_keyword_dictionary(self, keywords)
   def __copy__( self, **keywords):
    return curve(_other=self, **keywords)
   def get_red( self):
    return self.color[:,0]
   def get_green(self):
    return self.color[:,1]
   def get_blue(self):
    return self.color[:,2]
   def get_x(self):
    return self.pos[:,0]
   def get_y(self):
    return self.pos[:,1]
   def get_z(self):
    return self.pos[:,2]
   
   pos = property( cvisual.curve.get_pos, cvisual.curve.set_pos, None)
   color = property( cvisual.curve.get_color, cvisual.curve.set_color, None)
   x = property( get_x, cvisual.curve.set_x, None)
   y = property( get_y, cvisual.curve.set_y, None)
   z = property( get_z, cvisual.curve.set_z, None)
   red = property( get_red, cvisual.curve.set_red, None)
   green = property( get_green, cvisual.curve.set_green, None)
   blue = property( get_blue, cvisual.curve.set_blue, None)

class points( cvisual.points, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.points.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        copy_object( self, _other, keywords)
    else:
        cvisual.points.__init__(self)
        py_renderable.__init__(self, keywords)
        process_init_args_from_keyword_dictionary(self, keywords)
   def __copy__( self, **keywords):
    return points(_other=self, **keywords)
   def get_red( self):
    return self.color[:,0]
   def get_green(self):
    return self.color[:,1]
   def get_blue(self):
    return self.color[:,2]
   def get_opacity(self):
    return self.color[:,3]
   def get_x(self):
    return self.pos[:,0]
   def get_y(self):
    return self.pos[:,1]
   def get_z(self):
    return self.pos[:,2]
   
   pos = property( cvisual.points.get_pos, cvisual.points.set_pos, None)
   color = property( cvisual.points.get_color, cvisual.points.set_color, None)
   x = property( get_x, cvisual.points.set_x, None)
   y = property( get_y, cvisual.points.set_y, None)
   z = property( get_z, cvisual.points.set_z, None)
   red = property( get_red, cvisual.points.set_red, None)
   green = property( get_green, cvisual.points.set_green, None)
   blue = property( get_blue, cvisual.points.set_blue, None)
   opacity = property( get_opacity, cvisual.points.set_opacity, None)
   # alpha has been deprecated for some time, so eliminate it:
   # alpha = property( get_opacity, cvisual.points.set_opacity, None)

class convex( cvisual.convex, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.convex.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        copy_object( self, _other, keywords)
    else:
        cvisual.convex.__init__(self)
        py_renderable.__init__(self, keywords)
        process_init_args_from_keyword_dictionary(self, keywords)
    if (not self.display.uniform):
       raise RuntimeError, "Do not create convex with nonuniform axes."
   def __copy__( self, **keywords):
    return convex(_other=self, **keywords)
   pos = property( cvisual.convex.get_pos, cvisual.convex.set_pos, None)


class faces( cvisual.faces, py_renderable):
   def __init__( self, _other=None, **keywords):
    if _other:
        cvisual.faces.__init__(self, _other)
        py_renderable.__init__(self, keywords, _other)
        copy_object( self, _other, keywords)
    else:
        cvisual.faces.__init__(self)
        py_renderable.__init__(self, keywords)
        process_init_args_from_keyword_dictionary(self, keywords)
   def __copy__( self, **keywords):
    return faces(_other=self, **keywords)
   pos = property( cvisual.faces.get_pos, cvisual.faces.set_pos, None)
   normal = property( cvisual.faces.get_normal, cvisual.faces.set_normal, None)
   color = property( cvisual.faces.get_color, cvisual.faces.set_color, None)

class helix(object):
   def __init__( self, _other=None, pos=vector(),
        x=None, y=None, z=None, red=None, green=None, blue=None,
        axis=vector(1,0,0), radius=1.0, length=None, up=vector(0,1,0),
        coils=5, thickness=None, color=color.white, **extrakwds):
    if extrakwds.has_key('display'):
        disp = extrakwds['display']
    else:
        disp = cvisual.display.get_selected()
    if (not disp.uniform):
       raise RuntimeError, "Do not create helix with nonuniform axes."
    if extrakwds.has_key('frame'):
        fr = extrakwds['frame']
    else:
        fr = None
    if x is not None:
        pos[0] = x
    if y is not None:
        pos[1] = y
    if z is not None:
        pos[2] = z
    if red is not None:
        color[0] = red
    if green is not None:
        color[1] = green
    if blue is not None:
        color[2] = blue
    axis = vector(axis)
    if length is None:
        length = axis.mag
    self.__length = length
    self.__axis = axis
    self.__radius = radius
    self.__up = up
    self.__coils = coils
    self.__thickness = radius/20.
    if thickness:
        self.__thickness = thickness
    self.__frame = frame(display=disp, frame=fr, pos=pos, axis=axis.norm(), up=up)
    self.helix = curve( frame = self.__frame, radius = self.__thickness/2.,
        color = color, **extrakwds)
    self.create_pos()
      
   def create_pos(self):
    k = self.coils*(2*pi/self.__length)
    dx = (self.length/self.coils)/NCHORDS
    x_col = arange(0, self.__length+dx, dx)
    pos_data = zeros((len(x_col),3), float64)
    pos_data[:,0] = arange(0, self.__length+dx, dx)
    pos_data[:,1] = (self.radius) * sin(k*pos_data[:,0])
    pos_data[:,2] = (self.radius) * cos(k*pos_data[:,0])
    self.helix.pos = pos_data

   def set_pos(self, pos):
    self.__frame.pos = vector(pos)
   def get_pos(self):
    return self.__frame.pos

   def set_x(self, x):
    self.__frame.pos.x = x
   def get_x(self):
    return self.__frame.pos.x

   def set_y(self, y):
    self.__frame.pos.y = y
   def get_y(self):
    return self.__frame.pos.y

   def set_z(self, z):
    self.__frame.pos.z = z
   def get_z(self):
    return self.__frame.pos.z

   def set_color(self, color):
    self.helix.color = color
   def get_color(self):
    return self.helix.color

   def set_red(self, red):
    self.helix.red = red
   def get_red(self):
    return self.helix.red

   def set_green(self, green):
    self.helix.green = green
   def get_green(self):
    return self.helix.green

   def set_blue(self, blue):
    self.helix.blue = blue
   def get_blue(self):
    return self.helix.blue

   def set_radius(self, radius):
    scale = radius/self.__radius
    self.__radius = radius
    self.helix.y *= scale
    self.helix.z *= scale
   def get_radius(self):
    return self.__radius

   def set_axis(self, axis):
    axis = vector(axis)
    self.__axis = axis
    self.__frame.axis = axis.norm()
    self.set_length(axis.mag)
   def get_axis(self):
    return self.__axis

   def set_length(self, length):
    self.helix.x *= (length/self.__length)
    self.__length = length
    self.__frame.axis = self.__axis.norm()
    self.__axis = length*self.__frame.axis
   def get_length(self):
    return self.__length

   def set_coils(self, coils):
       if self.__coils == coils: return
       self.__coils = coils
       self.create_pos()
   def get_coils(self):
    return self.__coils

   def set_thickness(self, thickness):
       if self.__thickness == thickness: return
       self.__thickness = thickness
       self.helix.radius = thickness/2.
   def get_thickness(self):
    return self.__thickness

   def set_display(self, disp):
    self.helix.display = self.frame.display = disp
   def get_display(self):
    return self.helix.display

   def set_frame(self, fr):
    self.__frame.frame = fr
   def get_frame(self):
    return self.__frame.frame

   def set_up(self, up):
    self.__frame.up = up
   def get_up(self):
    return self.__frame.up

   pos = property( get_pos, set_pos, None)
   x = property( get_x, set_x, None)
   y = property( get_y, set_y, None)
   z = property( get_z, set_z, None)
   color = property( get_color, set_color, None)
   red = property( get_red, set_red, None)
   green = property( get_green, set_green, None)
   blue = property( get_blue, set_blue, None)
   axis = property( get_axis, set_axis, None)
   radius = property( get_radius, set_radius, None)
   coils = property( get_coils, set_coils, None)
   thickness = property( get_thickness, set_thickness, None)
   length = property( get_length, set_length, None)
   display = property( get_display, set_display, None)
   frame = property( get_frame, set_frame, None)
   up = property( get_up, set_up, None)

class texture(cvisual.texture):
   def __init__(self, **kwargs):
    cvisual.texture.__init__(self)
    for key, value in kwargs.iteritems():
        self.__setattr__(key, value)

