#!/usr/bin/env python
 import vtk
import math
from vtk.util.colors import *
filenames = ["upperarm.stl","forearm.stl"]
dt = 1.0    # degree step in rotation
angle = [0, 0] # shoulder and elbow joint angle
renWin = vtk.vtkRenderWindow()
assembly = vtk.vtkAssembly()
slider_shoulder = vtk.vtkSliderRepresentation2D()
slider_elbow = vtk.vtkSliderRepresentation2D()
actor = list() # the list of links
# Customize vtkInteractorStyleTrackballCamera 
class MyInteractor(vtk.vtkInteractorStyleTrackballCamera):
  def __init__(self,parent=None):
    self.AddObserver("CharEvent",self.OnCharEvent)
    self.AddObserver("KeyPressEvent",self.OnKeyPressEvent)
  # Override the default key operations which currently handle trackball or joystick styles is provided
  # OnChar is triggered when an ASCII key is pressed. Some basic key presses are handled here 
  def OnCharEvent(self,obj,event):
    pass
  def OnKeyPressEvent(self,obj,event):
    global angle
    # Get the compound key strokes for the event
    key = self.GetInteractor().GetKeySym()
    # Output the key that was pressed
    #print "Pressed: " , key
    # Handle an arrow key
    if(key == "Left"):
      actor[1].RotateY(-dt)      
    if(key == "Right"):
      actor[1].RotateY(dt)      
    if(key == "Up"):
      assembly.RotateY(-dt)
      angle[0] += dt
      if angle[0] >= 360.0:
        angle[0] -= 360.0
      slider_shoulder.SetValue(angle[0])  
    if(key == "Down"):
      assembly.RotateY(dt)
      angle[0] -= dt
      if angle[0] < 0.0:
        angle[0] += 360.0 
      slider_shoulder.SetValue(angle[0])
    # Ask each renderer owned by this RenderWindow to render its image and synchronize this process
    renWin.Render()
    return
def LoadSTL(filename):
  reader = vtk.vtkSTLReader()
  reader.SetFileName(filename)
  mapper = vtk.vtkPolyDataMapper() # maps polygonal data to graphics primitives
  mapper.SetInputConnection(reader.GetOutputPort())
  actor = vtk.vtkLODActor() 
  actor.SetMapper(mapper)
  return actor  # represents an entity in a rendered scene
def CreateCoordinates():
  # create coordinate axes in the render window
  axes = vtk.vtkAxesActor() 
  axes.SetTotalLength(100, 100, 100) # Set the total length of the axes in 3 dimensions 
  # Set the type of the shaft to a cylinder:0, line:1, or user defined geometry. 
  axes.SetShaftType(0) 
  axes.SetCylinderRadius(0.02) 
  axes.GetXAxisCaptionActor2D().SetWidth(0.03) 
  axes.GetYAxisCaptionActor2D().SetWidth(0.03) 
  axes.GetZAxisCaptionActor2D().SetWidth(0.03) 
  #axes.SetAxisLabels(0) # Enable:1/disable:0 drawing the axis labels
  #transform = vtk.vtkTransform() 
  #transform.Translate(0.0, 0.0, 0.0)
  #axes.SetUserTransform(transform)
  #axes.GetXAxisCaptionActor2D().GetCaptionTextProperty().SetColor(1,0,0)
  #axes.GetXAxisCaptionActor2D().GetCaptionTextProperty().BoldOff() # disable text bolding
  return axes
def ShoulderSliderCallback(obj,event):
  sliderRepres = obj.GetRepresentation()
  pos = sliderRepres.GetValue() 
  assembly.SetOrientation(0,-pos,0)

  renWin.Render()
def ElbowSliderCallback(obj,event):
  sliderRepres = obj.GetRepresentation()
  pos = sliderRepres.GetValue() 
  actor[1].SetOrientation(0,-pos,0)
  renWin.Render()
def ConfigSlider(sliderRep, TitleText, Yaxes):
  sliderRep.SetMinimumValue(0.0)
  sliderRep.SetMaximumValue(360.0)
  sliderRep.SetValue(0.0) # Specify the current value for the widget
  sliderRep.SetTitleText(TitleText) # Specify the label text for this widget
  sliderRep.GetSliderProperty().SetColor(1,0,0) # Change the color of the knob that slides
  sliderRep.GetSelectedProperty().SetColor(0,0,1) # Change the color of the knob when the mouse is held on it
  sliderRep.GetTubeProperty().SetColor(1,1,0) # Change the color of the bar 
  sliderRep.GetCapProperty().SetColor(0,1,1) # Change the color of the ends of the bar
  #sliderRep.GetTitleProperty().SetColor(1,0,0) # Change the color of the text displaying the value
  # Position the first end point of the slider
  sliderRep.GetPoint1Coordinate().SetCoordinateSystemToDisplay()
  sliderRep.GetPoint1Coordinate().SetValue(50, Yaxes) 
  # Position the second end point of the slider
  sliderRep.GetPoint2Coordinate().SetCoordinateSystemToDisplay()
  sliderRep.GetPoint2Coordinate().SetValue(400, Yaxes) 
  sliderRep.SetSliderLength(0.02) # Specify the length of the slider shape.The slider length by default is 0.05
  sliderRep.SetSliderWidth(0.02) # Set the width of the slider in the directions orthogonal to the slider axis
  sliderRep.SetTubeWidth(0.005)
  sliderRep.SetEndCapWidth(0.03)

  sliderRep.ShowSliderLabelOn() # display the slider text label
  sliderRep.SetLabelFormat("%.1f")

  sliderWidget = vtk.vtkSliderWidget()
  sliderWidget.SetRepresentation(sliderRep)
  sliderWidget.SetAnimationModeToAnimate()

  return sliderWidget
def CreateGround():
  # create plane source
  plane = vtk.vtkPlaneSource()
  plane.SetXResolution(50)
  plane.SetYResolution(50)
  plane.SetCenter(0,0,0)
  plane.SetNormal(0,0,1)  
  # mapper
  mapper = vtk.vtkPolyDataMapper()
  mapper.SetInputConnection(plane.GetOutputPort())
   
  # actor
  actor = vtk.vtkActor()
  actor.SetMapper(mapper)
  actor.GetProperty().SetRepresentationToWireframe()
  #actor.GetProperty().SetOpacity(0.4) # 1.0 is totally opaque and 0.0 is completely transparent
  actor.GetProperty().SetColor(light_grey)
  '''
  # Load in the texture map. A texture is any unsigned char image.
  bmpReader = vtk.vtkBMPReader() 
  bmpReader.SetFileName("ground_texture.bmp") 
  texture = vtk.vtkTexture() 
  texture.SetInputConnection(bmpReader.GetOutputPort()) 
  texture.InterpolateOn() 
  actor.SetTexture(texture)
  '''
  transform = vtk.vtkTransform()
  transform.Scale(2000,2000, 1)
  actor.SetUserTransform(transform)
  return actor  
def CreateScene():
  # Create a rendering window and renderer
  ren = vtk.vtkRenderer()
  #renWin = vtk.vtkRenderWindow()
  renWin.AddRenderer(ren)
  # Create a renderwindowinteractor
  iren = vtk.vtkRenderWindowInteractor()
  iren.SetRenderWindow(renWin)
  style = MyInteractor()
  style.SetDefaultRenderer(ren)
  iren.SetInteractorStyle(style)
  for id, file in enumerate(filenames):
    actor.append(LoadSTL(file))
    #actor[id].GetProperty().SetColor(blue)
    r = vtk.vtkMath.Random(.4, 1.0)
    g = vtk.vtkMath.Random(.4, 1.0)
    b = vtk.vtkMath.Random(.4, 1.0)
    actor[id].GetProperty().SetDiffuseColor(r, g, b)
    actor[id].GetProperty().SetDiffuse(.8)
    actor[id].GetProperty().SetSpecular(.5)
    actor[id].GetProperty().SetSpecularColor(1.0,1.0,1.0)
    actor[id].GetProperty().SetSpecularPower(30.0)
    assembly.AddPart(actor[id])
    # Add the actors to the scene
    #ren.AddActor(actor[id])
  # Also set the origin, position and orientation of assembly in space.
  assembly.SetOrigin(0, 0, 0) # This is the point about which all rotations take place 
  #assembly.AddPosition(0, 0, 0)
  #assembly.RotateX(45)
  actor[1].SetOrigin(274, 0, 0) # initial elbow joint position
  ren.AddActor(assembly)
  # Add coordinates
  axes = CreateCoordinates()
  ren.AddActor(axes)

  # Add ground
  ground = CreateGround()
  ren.AddActor(ground)

  # Add slider to control the robot
  sliderWidget_shoulder = ConfigSlider(slider_shoulder,"Shoulder Joint", 80)
  sliderWidget_shoulder.SetInteractor(iren)
  sliderWidget_shoulder.EnabledOn()
  sliderWidget_shoulder.AddObserver("InteractionEvent", ShoulderSliderCallback)

  sliderWidget_elbow = ConfigSlider(slider_elbow,"Elbow Joint", 160)
  sliderWidget_elbow.SetInteractor(iren)
  sliderWidget_elbow.EnabledOn()
  sliderWidget_elbow.AddObserver("InteractionEvent", ElbowSliderCallback)

  # Set background color
  ren.SetBackground(.2, .2, .2)

  # Set window size
  renWin.SetSize(600, 600)

  # Set up the camera to get a particular view of the scene
  camera = vtk.vtkCamera()
  camera.SetFocalPoint(300, 0, 0)
  camera.SetPosition(300, -400, 350)
  camera.ComputeViewPlaneNormal()
  camera.SetViewUp(0, 1, 0)
  camera.Zoom(0.4)
  ren.SetActiveCamera(camera)
  # Enable user interface interactor
  iren.Initialize()
  iren.Start()
if __name__ == "__main__":
  CreateScene()