from visual import*
from visual.graph import *

graph1 = gdisplay(x=0, y=0, width=600, height=550,
             title='Spring Potential Energy', xtitle='x(m)', ytitle='V(J)', 
             xmax=10., xmin=-10., ymax=500, ymin=-0, 
             foreground=color.white, background=color.black)

VPlot = gcurve(color=color.red)


scene.autoscale = 0
scene.range = 50
userspin = 0
equilibrium = vector(10,0,0)
stretch = mag(equilibrium)
ball = sphere(pos = equilibrium, color = (1,0,0), mass = 1, radius = 2)
ball.p = vector(0,0,0)
spring = helix(pos = (0,0,0), coils = 10, axis = ball.pos, length = 10, radius = 0.5, thickness = 0.75)
k = 20 # N/m
pick = None # no object picked out of the scene yet
time = 0
ts = 0.005

while 1:
    rate(100)
    if scene.mouse.events:

        m1 = scene.mouse.getevent() # obtain drag or drop event

        if m1.drag and m1.pick == ball:

            drag_pos = m1.pickpos

            pick = m1.pick

            scene.cursor.visible = 0 # make cursor invisible

        elif m1.drop:

            pick = None # end dragging

            scene.cursor.visible = 1 # cursor visible
            ball.color = (1,0,1)
            stretch = mag(m1.pos)-mag(equilibrium)
            

    if pick:

        new_pos = scene.mouse.project(normal=(0,0,1))

        if new_pos != drag_pos:

            pick.pos += new_pos - drag_pos
            ball.p = vector(0,0,0)

            drag_pos = new_pos
            spring.axis = new_pos
            ball.color = (1,0,0)
    stretch = ball.pos-equilibrium
    spring.axis = ball.pos
    imp = -k*stretch*ts
    ball.p = ball.p + imp
    ball.pos = ball.pos + (ball.p/ball.mass)*ts
    time = time + ts
    
    VPlot.plot(pos=(stretch.x,0.5*k*mag(stretch)**2))