from visual import *

# Some basics for my robot


# Set a few things for the screen
floor=box(pos=vector(0,0,0),height=0.1,width=4,length=10,color=color.green)
scene.center=(0,5,0)
scene.autoscale=0


# Create the parts of the robot
parts=[]
struts=[]

# Create the left foot
parts.append(cone(pos=vector(-1,0,0),axis=(0,0.5,0),radius=0.5,color=color.blue))
parts[0].m=2.
parts[0].v=vector(0,0,0)
parts[0].jet=vector(0,0,0)

# Create the right foot
parts.append(cone(pos=vector(1,0,0),axis=(0,0.5,0),radius=0.5,color=color.blue))
parts[1].m=2.
parts[1].v=vector(0,0,0)
parts[1].jet=vector(0,0,0)

# Add the torso
parts.append(box(pos=(0,3,0),color=color.red,width=0.5,length=1.5,height=2))
parts[2].m=20.
parts[2].v=vector(0.1,0,0)
parts[2].jet=vector(0,0,0)

#####################################################################################
# Students, you should put in new parts here!

nparts=3

#
####################################################################################

#Add all of the struts
#the "side1" and "side2" indicate the parts on either side of the strut
struts.append(helix(radius=0.1))
struts[0].side1=0
struts[0].side2=1
struts[0].k=2000

struts.append(helix(radius=0.1))
struts[1].side1=1
struts[1].side2=2
struts[1].k=2000

struts.append(helix(radius=0.1))
struts[2].side1=2
struts[2].side2=0
struts[2].k=1000


#####################################################################################
# Students, you should put in new struts here!

nstruts=3

#
####################################################################################


#Initialize the struts
for j in range (nstruts):
	struts[j].pos=parts[struts[j].side1].pos
	struts[j].axis=parts[struts[j].side2].pos-parts[struts[j].side1].pos
	struts[j].x0=mag(struts[j].axis)

# Now the evolution
g=9.8
t=0
dt=0.01

#Note: if all of the parts are connected via a "chain", you shouldn't need to change 
#this part at all!
while(t < 20):
	rate(100)
	# First, calculate the forces on all of the parts
	for i in range(nparts):
		parts[i].F=parts[i].m*vector(0,-g,0)+parts[i].jet

	for j in range(nstruts):
		# Also, the forces from the springs
		dF=struts[j].k*struts[j].axis*(mag(struts[j].axis)-struts[j].x0)/\
			mag(struts[j].axis)
		parts[struts[j].side1].F=parts[struts[j].side1].F+dF
		parts[struts[j].side2].F=parts[struts[j].side2].F-dF

	# Next, update the positions of the parts

	for i in range(nparts):
		parts[i].v=parts[i].v+parts[i].F/parts[i].m*dt
		# Check to see if anything has gone through the floor
		if ((parts[i].y < 0) and (parts[i].v.y < 0)):
			parts[i].v.y=-parts[i].v.y
			# This part adds a bit of "friction"
			parts[i].v=parts[i].v*0.9
		parts[i].pos=parts[i].pos+parts[i].v*dt

	#Finally, update the struts
	for j in range (nstruts):
		struts[j].pos=parts[struts[j].side1].pos
		struts[j].axis=parts[struts[j].side2].pos-parts[struts[j].side1].pos
	
	t=t+dt