# 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

from visual import *
from visual.controls import *

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

###########################Arm 1##################################
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, 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 = ((0.5*wbase)+L1,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

##################################Arm 2#######################################

frame20 = frame(frame=frameworld)
frame20.pos = (0,0,0)
frame21 = frame(frame=frame20)
frame21.pos = (0,tbase+0.9*hpedestal,0)

pedestal2 = box(frame = frame20, 
               pos = (0,(0.55*hpedestal)+tbase,0),
               height = 1.1*hpedestal,
               length = 2*wpedestal,
               width = wpedestal,
               color = (0.5, 0.5, 0.5))
base2 = cylinder(frame=frame20,
                pos=(0,0,0),
                axis=(0,tbase,0),
                radius=0.5*wbase,
                color=pedestal.color)
axle2 = cylinder(frame = frame21,
                pos=(0,0,offset/2),
                axis = (0,0, -offset),
                radius = d/4.,
                color = color.blue)
bar21 = box(frame = frame21,
           pos = (0,(L1display/2.-d/2.),-(gap+d)/2.),
           size = (d, L1display, d),
           color = color.blue)
bar21b = box(frame = frame21,
            pos = (0,(L1display/2.-d/2.),(gap+d)/2.),
            size = (d, L1display, d),
            color = color.blue)
frame21.axis = (0, 1, 0)
frame22 = frame(frame=frame21)
frame22.pos = frame21.axis*L1
axle21 = cylinder(frame = frame22,
                 pos = (0, 0, -(gap+d)/2.),
                 axis = (0, 0, gap+d),
                 radius = axle.radius,
                 color = axle.color)
bar22 = box(frame = frame22,
           pos = (L2display/2.-d/2., 0,0),
           size = (L2display, d, d),
           color = color.blue)
frame22.axis = (0, 1, 0)
frame23 = frame(frame=frame22)
frame23.pos = (frame22.axis+(1, -1, 0))*L2

bar23 = box(frame = frame23,
           pos = (L3display/2.-d/2.,0,0),
           size = (L3display, d, d),
           color = color.blue)
frame24 = frame(frame=frame23)
frame24.pos = (frame23.axis*L3)

frame21.rotate(axis = (0, 0, 1), angle = theta1)
frame22.rotate(axis = (0, 0, 1), angle = theta2)
frame23.rotate(axis = (0, 0, 1), angle = theta3)

#********************Frame**Markers******************************
#ballf1 = sphere(frame=frame21, pos=(0,0,0), color=color.red, radius=0.1)
#ballf2 = sphere(frame=frame22, pos=(0,0,0), color=color.red, radius=0.1)
ball2f3 = sphere(frame=frame23, pos=(0,0,0), color=color.blue, radius=0.1)
#ballf4 = sphere(frame=frame24, pos=(0,0,0), color=color.red, radius=0.1)

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


#*****"Polar" Coordinate Stuff********
r2i = 2*L2*cos(theta1)
r2f = 2*L2*cos(theta1)
r2c = 2*L2*cos(theta1)
theta21i = theta1
theta21f = theta1
theta21c = theta1
alpha2i = 0
alpha2f = 0
alpha2c = 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":
                    alphaf = alphaf + (pi/12)
                    alpha2f = alpha2f + (pi/12)
                if kcap[indx] == "d":
                    alphaf = alphaf - (pi/12)
                    alpha2f = alpha2f - (pi/12)
                if kcap[indx] == "w":
                    rf = rf + 0.1
                    if rf > (2*L2):
                        rf = (2*L2)
                    r2f = r2f + 0.1
                    if r2f > (2*L2):
                        r2f = (2*L2)
                if kcap[indx] == "s":
                    rf = rf - 0.1
                    if rf < 0.4:
                        rf = 0.4
                    r2f = r2f - 0.1
                    if r2f < 0.4:
                        r2f = 0.4
            del kcap[0:len(kcap)]
            theta1f = acos((rf/(2*L2)))
            theta21f = acos((rf/(2*L2))) 
            for t in arange(0, 1.01, .01):
                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

                theta21n = theta21i*(1-t) + theta21f*(t)
                dtheta21 = theta21n - theta21c
                frame21.rotate(axis = (0, 0, 1), angle = dtheta21)
                frame22.rotate(axis = (0, 0, 1), angle = 2*dtheta21)
                frame23.rotate(axis = (0, 0, 1), angle = dtheta21)
                theta21c = theta21c + dtheta21
                rate(30)
                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

                alpha2n = alpha2i*(1-t) + alpha2f*(t)
                dalpha2 = alpha2n - alpha2c
                frame20.rotate(axis = (0, 1, 0), angle = dalpha2)
                alpha2c = alpha2c + dalpha2
          
            theta21i = theta21f
            alpha2i = alpha2f
            print "alpha2 =", alpha2c 
            print "radius2 =", r2f
            theta1i = theta1f
            alphai = alphaf
            print "alpha =", alphac 
            print "radius =", rf