from __future__ import division
from visual import *
from visual.graph import *

# Comment in the type(s) of friction of interest
scene.background = color.white
scene.x = scene.y = 0
scene.width = 700
scene.height = 700

r = .05
L = 2
Slen = 0.7
Blen = 0.35
y = -.5
G = 9.8
dt = .0005
t = 0
u = 1.0*.002 # 0.002 for resonance experiment.
Ks = 0.7

# Select types of friction applying
viscous = 0
sliding = 0
air = 0

# Select whether you want a graph shown
graph = 0

# Select whether you want force vectors shown
showarrow = 0

# Set up oscillator properties. Frequency 5.9 for resonance
amplitude = 0.0  # 0.03 for resonance
frequency = 5.9  # 2 and 20 bracket the possibilities


Bmass = 0.02 # 0.02

y = -0.805 # Starting position, -0.525 for spring natural length, 0.805 for balance with weight


block = box(length=Blen, width=Blen, height=Blen, pos=(0.,y+Blen/2.0,0), mass=Bmass, color=color.blue)
support = box(length=1.5, width=1.5, height=0.3, pos=(0.,L-Slen/2.0,0), color=color.green)
spring = helix(color=color.red, pos=support.pos, axis=block.pos-support.pos,
               radius = 0.1, thickness=0.04, coils=15)


if showarrow == 1:
   springarrow = arrow(pos=(0.3,0,0), axis=(0,0,0), color=color.red, shaftwidth=0.1, fixedwidth=True)
   gravarrow = arrow(pos=(0.5,0,0), axis=(0,-9.8*Bmass,0), color=color.blue,
                     shaftwidth=0.1, fixedwidth=True)
   fricarrow = arrow(pos=(0.7,0,0), axis=(0,0,0), color=color.yellow,
                     shaftwidth=0.1, fixedwidth=True)


gravforce = vector(0,0,0)
scene.autoscale=0

block.p = vector(0,0,0)
block.vel = vector(0,0,0)
ffriction = vector(0,0,0)

if graph == 1:
    gdisplay(x=0, y=200, height=200, width=700, background=color.white,
         title='K:blue, Uspring:red, Ugrav: magenta')
    usplot=gcurve(color=color.red)
    kplot=gcurve(color=color.blue)
    ugplot=gcurve(color=color.magenta)
    posplot=gcurve(color=color.blue)

while (t < 50):
    rate(800)
    support.pos = (0., L -Slen/2.0 + amplitude*cos(t*frequency),0)
    gravforce = G*block.mass
    ## viscous friction
    if viscous == 1:
       ffriction.x = -u*block.vel.x
    ## end viscous friction

    ## sliding and air resistance
    if (block.vel.x < 0):
        if sliding == 1:
            ffriction.x = u*gravforce      ## sliding
        if air == 1:
            ffriction.x = u*(block.vel.x**2)     ## air resistance
    else:
        if sliding == 1:
            ffriction.x = -u*gravforce     ## sliding
        if air == 1:
            ffriction.x = -u*(block.vel.x**2)   ## air resistance
    ## end sliding and air resistance
        
    fspring = (block.pos-support.pos+(0,L,0))*(-Ks)
    fgrav = (0,-9.8*block.mass,0)
    block.p = block.p + (ffriction + fspring + fgrav)*dt
    block.vel = block.p/block.mass
    block.pos = block.pos + (block.vel*dt)
    y=mag(block.p)
    spring.axis = spring.axis + block.p*dt/block.mass
    t = t + dt

##Energy Stuff=======================================================
    spring.k = (mag(block.p)**2) / (2 * block.mass)
    spring.us = (.5*(Ks*((mag(block.pos))**2)))
    spring.ug = 9.8*block.mass*block.pos[1]
   

## Graph Stuff=======================================================
    if graph == 1:
        usplot.plot(pos=(t , spring.us))
        kplot.plot(pos=(t, spring.k))
        ugplot.plot(pos=(t, spring.ug))

## Arrow stuff
    if showarrow == 1:
        springarrow.axis=fspring
        fricarrow.axis=ffriction