# Based on Bruce Sherwood's Double Pendulum
# Modified to look somewhat more like a robot arm
# By Andrew Lee
# Math 198 Fall 2009

#Choose a location to move to by pressing keys 0-9
#Press enter (return) to execute

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 = (0.80, 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)
frame1.axis = (0, 1, 0)
frame2 = frame(frame=frame1)
frame2.pos = frame1.axis*L1
frame2.axis = (0, 1, 0)
frame3 = frame(frame=frame2)
frame3.pos = (frame2.axis+(1, -1, 0))*L2
frame4 = frame(frame=frame3)
frame4.pos = (frame3.axis*L3display)
frame1.rotate(axis = (0, 0, 1), angle = theta1)
frame2.rotate(axis = (0, 0, 1), angle = theta2)
frame3.rotate(axis = (0, 0, 1), angle = theta3)
frame4.rotate(axis = (0, 0, 1), angle = pi/2)
#************************************************
pedestal = box(frame = frame0, 
               pos = (0,(0.55*hpedestal)+tbase,0),
               height = 1.1*hpedestal,
               length = 2*wpedestal,
               width = wpedestal,
               color = (0.8, 0.8, 0.8))
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)
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)
bar3 = box(frame = frame3,
               pos = (L3display/2.-d/2.,0,0),
               size = (L3display, d, d),
               color = color.green)
joint3 = sphere(frame=frame3,
               pos=(0,0,0),
               color=color.green,
               radius=0.1)
floor = box(frame=frameworld,
               pos = (0,0,0),
               height = 0.001,
               length = 1,
               width = 1.75,
               color = (0.5, 0.5, 0.5))
#**************************************************************************
wallthck = 0.05
wallhght = 0.05
mplhght = 0.2
mplcent = 2*L2*cos(theta1)

framemp = frame(frame=frameworld)
framemp.pos =(mplcent, mplhght,0)
framep0 = frame(frame=frameworld)
framep0.pos = (0, 0.2, 0.6)
#**************************************************************************
platform0 = box(frame = framep0,
               height = 0.2,
               length = 0.5,
               width = 0.5,
               pos = (0, -0.1, 0),
               color = (1.0, 1.0, 1.0))
piece = cylinder(frame = framep0,
               pos = (0, 0, 0),
               axis = (0, 0.06, 0),
               radius = 0.2,
               color = color.blue)
mplatform = box(frame = framemp,
               height = mplhght,
               length = 1.5,
               width = 1.5,
               pos = (0, -0.5*mplhght, 0),
               color = (1.0, 1.0, 1.0))
mplatform1 = box(frame = framemp,
               height = wallhght,
               length = wallthck,
               width = mplatform.width,
               pos = (-mplatform.length/2+(0.5*wallthck), 0.5*wallhght, 0),
               color = (0.9, 0.9, 0.9))
mplatform2 = box(frame = framemp,
               height = wallhght,
               length = wallthck,
               width = mplatform.width,
               pos = (mplatform.length/2-(0.5*wallthck), 0.5*wallhght, 0),
               color = mplatform1.color)
mplatform3 = box(frame = framemp,
               height = wallhght,
               length = mplatform.length,
               width = wallthck,
               pos = (0, 0.5*wallhght, -mplatform.width/2+(0.5*wallthck)),
               color = mplatform1.color)
mplatform4 = box(frame = framemp,
               height = wallhght,
               length = mplatform.length,
               width = wallthck,
               pos = (0, 0.5*wallhght, mplatform.width/2-(0.5*wallthck)),
               color = mplatform1.color)
mplatform5 = box(frame = framemp,
               height = wallhght,
               length = mplatform.length,
               width = wallthck,
               pos = (0, 0.5*wallhght, (mplatform.width/6)),
               color = mplatform1.color)
mplatform6 = box(frame = framemp,
               height = wallhght,
               length = mplatform.length,
               width = wallthck,
               pos = (0, 0.5*wallhght, (-mplatform.width/6)),
               color = mplatform1.color)
mplatform7 = box(frame = framemp,
               height = wallhght,
               length = wallthck,
               width = mplatform.width,
               pos = ((-mplatform.length/6), 0.5*wallhght, 0),
               color = mplatform1.color)
mplatform8 = box(frame = framemp,
               height = wallhght,
               length = wallthck,
               width = mplatform.width,
               pos = ((mplatform.length/6), 0.5*wallhght, 0),
               color = mplatform1.color)
#**************************************************************************
goffset = mplatform.width/3-wallthck/4
frameg1 = frame(frame = framemp)
frameg1.pos = (-goffset, 0, -goffset)
frameg2 = frame(frame = framemp)
frameg2.pos = (0, 0, -goffset)
frameg3 = frame(frame = framemp)
frameg3.pos = (goffset, 0, -goffset)
frameg4 = frame(frame = framemp)
frameg4.pos = (-goffset, 0, 0)
frameg5 = frame(frame = framemp)
frameg5.pos = (0, 0, 0)
frameg6 = frame(frame = framemp)
frameg6.pos = (goffset, 0, 0)
frameg7 = frame(frame = framemp)
frameg7.pos = (-goffset, 0, goffset)
frameg8 = frame(frame = framemp)
frameg8.pos = (0, 0, goffset)
frameg9 = frame(frame = framemp)
frameg9.pos = (goffset, 0, goffset)
#**************************************************************************
scene.autoscale = 0
ri = 2*L2*cos(theta1)
raf = rpf = 2*L2*cos(theta1)
rc = 2*L2*cos(theta1)
theta1i = theta1
theta1f = theta1
theta1c = theta1
theta1pf = 0
theta1pc = 0
theta1pi = 0
alphai = 0
alphaf = 0
alphac = 0
zf = 0
cmdlst = ["u", "p"]
#**************************************************************************
while 1:
    if scene.kb.keys:
        key = scene.kb.getkey()
        if key == "1":
            cmdlst.extend(["d", "cfa", "u", "g1", "d", "cfg1", "u"])
        if key == "2":
            cmdlst.extend(["d", "cfa", "u", "g2", "d", "cfg2", "u"])
        if key == "3":
            cmdlst.extend(["d", "cfa", "u", "g3", "d", "cfg3", "u"])
        if key == "4":
            cmdlst.extend(["d", "cfa", "u", "g4", "d", "cfg4", "u"])        
        if key == "5":
            cmdlst.extend(["d", "cfa", "u", "g5", "d", "cfg5", "u"])
        if key == "6":
            cmdlst.extend(["d", "cfa", "u", "g6", "d", "cfg6", "u"])
        if key == "7":
            cmdlst.extend(["d", "cfa", "u", "g7", "d", "cfg7", "u"])
        if key == "8":
            cmdlst.extend(["d", "cfa", "u", "g8", "d", "cfg8", "u"])
        if key == "9":
            cmdlst.extend(["d", "cfa", "u", "g9", "d", "cfg9", "u"])
        if key == "0":
            cmdlst.extend(["d", "cfa", "u", "p", "d", "cfp", "u"])
        if key == "\n":
            for indx in range(0, len(cmdlst), 1):
                cpos = "yes"
                if cmdlst[indx] == "p":
                    alphaf = -pi/2
                    rpf = 0.6
                if cmdlst[indx] == "u":
                    zf = 0.2
                if cmdlst[indx] == "d":
                    zf = 0
                if cmdlst[indx] == "g1":
                    alphaf = atan(goffset/(mplcent-goffset))
                    rpf = (mplcent-goffset)/cos(alphaf)
                if cmdlst[indx] == "g2":
                    alphaf = atan(goffset/mplcent)
                    rpf = mplcent/cos(alphaf)
                if cmdlst[indx] == "g3":
                    alphaf = atan(goffset/(mplcent+goffset))
                    rpf = (mplcent+goffset)/cos(alphaf)
                if cmdlst[indx] == "g4":
                    alphaf = 0
                    rpf = mplcent - goffset
                if cmdlst[indx] == "g5":
                    alphaf = 0
                    rpf = mplcent
                if cmdlst[indx] == "g6":
                    alphaf = 0
                    rpf = mplcent + goffset                    
                if cmdlst[indx] == "g7":
                    alphaf = -atan(goffset/(mplcent-goffset))
                    rpf = (mplcent-goffset)/cos(alphaf) 
                if cmdlst[indx] == "g8":
                    alphaf = -atan(goffset/mplcent)
                    rpf = mplcent/cos(alphaf)
                if cmdlst[indx] == "g9":
                    alphaf = -atan(goffset/(mplcent+goffset))
                    rpf = (mplcent+goffset)/cos(alphaf)
                if cmdlst[indx] == "cfa":
                    piece.frame = frame(frame=frame4)
                    cpos = "no"
                if cmdlst[indx] == "cfg1":
                    piece.frame = frame(frame=frameg1)
                    cpos = "no"
                if cmdlst[indx] == "cfg2":
                    piece.frame = frame(frame=frameg2)
                    cpos = "no"
                if cmdlst[indx] == "cfg3":
                    piece.frame = frame(frame=frameg3)
                    cpos = "no"
                if cmdlst[indx] == "cfg4":
                    piece.frame = frame(frame=frameg4)
                    cpos = "no"
                if cmdlst[indx] == "cfg5":
                    piece.frame = frame(frame=frameg5)
                    cpos = "no"
                if cmdlst[indx] == "cfg6":
                    piece.frame = frame(frame=frameg6)
                    cpos = "no"
                if cmdlst[indx] == "cfg7":
                    piece.frame = frame(frame=frameg7)
                    cpos = "no"
                if cmdlst[indx] == "cfg8":
                    piece.frame = frame(frame=frameg8)
                    cpos = "no"
                if cmdlst[indx] == "cfg9":
                    piece.frame = frame(frame=frameg9)
                    cpos = "no"
                if cmdlst[indx] == "cfp":
                    piece.frame = frame(frame=framep0)
                    cpos = "no"
                theta1pf = atan(zf/rpf)
                raf = rpf/(cos(theta1pf))
                theta1f = acos((raf/(2*L2)))
                if cpos == "yes":
                    for t in arange(0, 1.01, .01):
                        rate(65)
                        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
                        theta1pn = theta1pi*((cos((t*pi)/2))**2) + theta1pf*((sin((t*pi)/2))**2)
                        dtheta1p = theta1pn - theta1pc
                        frame1.rotate(axis = (0, 0, 1), angle = dtheta1p)
                        frame3.rotate(axis = (0, 0, 1), angle = -dtheta1p)
                        theta1pc = theta1pc + dtheta1p
                theta1i = theta1f
                alphai = alphaf
                theta1pi = theta1pf
            del cmdlst[0:len(cmdlst)]

            
#End Program