# Based on Bruce Sherwood's Double Pendulum
# Modified to look somewhat more like a robot arm

#Control movement with "a", "s", "d", and "w" keys
#Press "enter" to execute
#Note, z-axis movement (up and down) not working on this version

from visual import *
from visual.controls import *

scene.title = "Robot Arm"
scene.height = scene.width = 800


d = 0.1                 
gap = 1.5*d              
L1 = 1.0
L1display = L1+d
L2 = 1.0
L2display = L2+d
L3 = 0.45
L3display = L3+d

hpedestal = 0.4
wpedestal = 0.1
tbase = 0.4
wbase = 0.6
offset = 2.*gap
top = vector(0,0,0)
scene.center = (L2*cos(pi/4), 0.55*hpedestal+tbase, 0)

theta1 = pi/4
theta2 = pi/2
theta3 = -pi/4

frameworld = frame()
frame0 = frame(frame=frameworld)
frame0.pos = (0,0,0)
frame1 = frame(frame=frame0)
frame1.pos = (0,tbase+0.9*hpedestal,0)

pedestal = box(frame = frame0, 
               pos = (0,(0.55*hpedestal)+tbase,0),
               height = 1.1*hpedestal,
               length = 2*wpedestal,
               width = wpedestal,
               color = (0.5, 0.5, 0.5))
base = cylinder(frame=frame0,
                pos=(0,0,0),
                axis=(0,tbase,0),
                radius=0.5*wbase,
                color=pedestal.color)
axle = cylinder(frame = frame1,
                pos=(0,0,offset/2),
                axis = (0,0, -offset),
                radius = d/4.,
                color = color.yellow)
bar1 = box(frame = frame1,
           pos = (0,(L1display/2.-d/2.),-(gap+d)/2.),
           size = (d, L1display, d),
           color = color.green)
bar1b = box(frame = frame1,
            pos = (0,(L1display/2.-d/2.),(gap+d)/2.),
            size = (d, L1display, d),
            color = color.green)
frame1.axis = (0, 1, 0)
frame2 = frame(frame=frame1)
frame2.pos = frame1.axis*L1
axle1 = cylinder(frame = frame2,
                 pos = (0, 0, -(gap+d)/2.),
                 axis = (0, 0, gap+d),
                 radius = axle.radius,
                 color = axle.color)
bar2 = box(frame = frame2,
           pos = (L2display/2.-d/2., 0,0),
           size = (L2display, d, d),
           color = color.green)
frame2.axis = (0, 1, 0)
frame3 = frame(frame=frame2)
frame3.pos = (frame2.axis+(1, -1, 0))*L2

bar3 = box(frame = frame3,
           pos = (L3display/2.-d/2.,0,0),
           size = (L3display, d, d),
           color = color.green)
frame4 = frame(frame=frame3)
frame4.pos = (frame3.axis*L3)

frame1.rotate(axis = (0, 0, 1), angle = theta1)
frame2.rotate(axis = (0, 0, 1), angle = theta2)
frame3.rotate(axis = (0, 0, 1), angle = theta3)

#********************Frame**Markers******************************
floor = box(frame=frameworld,
           pos = (0,0,0),
           height = 0.001,
           length = 1,
           width = 1,
           color = (1.0, 1.0, 1.0))
#ballf1 = sphere(frame=frame1, pos=(0,0,0), color=color.red, radius=0.1)
#ballf2 = sphere(frame=frame2, pos=(0,0,0), color=color.red, radius=0.1)
ballf3 = sphere(frame=frame3, pos=(0,0,0), color=color.green, radius=0.1)
#ballf4 = sphere(frame=frame4, pos=(0,0,0), color=color.red, radius=0.1)
#*****************************************************************
platform = box(frame = frameworld, 
               height = 0.2,
               length = 2*L1,
               width = 2*L1,
               pos = (2*L2*cos(pi/4),0.1,0),
               color = (1.0, 1.0, 1.0))
#*****************************************************************

scene.autoscale = 0

dt=0.01
rangle = .01
rotatecount = pi/4


#*****"Polar" Coordinate Stuff********
ri = 2*L2*cos(theta1)
rf = 2*L2*cos(theta1)
rc = 2*L2*cos(theta1)
theta1i = theta1
theta1f = theta1
theta1c = theta1
alphai = 0
alphaf = 0
alphac = 0
kcap = []
kct = 0

#******"Rectangular" Coordinate Stuff**
yf = 0
xf = 2*L2*cos(theta1)
dxyz = 0.1
zf = 0

while 1:
    if scene.kb.keys:
        key = scene.kb.getkey()
        if key != "\n":
            kcap.append(key)
        print kcap
        if key == "\n":
            for indx in range(0, len(kcap), 1):
                if kcap[indx] == "a":
                    xf = xf - dxyz
                if kcap[indx] == "d":
                    xf = xf + dxyz
                if kcap[indx] == "w":
                    yf = yf + dxyz
                if kcap[indx] == "s":
                    yf = yf - dxyz
                if kcap[indx] == "q":
                    zf = zf + dxyz
                if kcap[indx] == "z":
                    zf = zf - dxyz
            del kcap[0:len(kcap)]
            alphaf = atan(yf/xf)
            rf = xf/cos(alphaf)   
            theta1f = acos((rf/(2*L2)))
            for t in arange(0, 1.05, .05):
                rate(30)
                theta1n = theta1i*((cos((t*pi)/2))**2) + theta1f*((sin((t*pi)/2))**2)
                dtheta1 = theta1n - theta1c
                frame1.rotate(axis = (0, 0, 1), angle = dtheta1)
                frame2.rotate(axis = (0, 0, 1), angle = 2*dtheta1)
                frame3.rotate(axis = (0, 0, 1), angle = dtheta1)
                theta1c = theta1c + dtheta1
                alphan = alphai*((cos((t*pi)/2))**2) + alphaf*((sin((t*pi)/2))**2)
                dalpha = alphan - alphac
                frame0.rotate(axis = (0, 1, 0), angle = dalpha)
                alphac = alphac + dalpha
            theta1i = theta1f
            alphai = alphaf
            print "x =", xf-(2*L2*cos(pi/4)) 
            print "y =", yf
        if key == "h":
            ballf4.frame = frameworld
            ballf4.pos = (xf, 0.3, -yf)