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 = 200

r = .05
L = 2
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 = 1
sliding = 0
air = 0

# Select whether you want a graph shown
graph = 0

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

y = y*start


block = box(length=.5, width=.5, height=.5, pos=(y,0,0), mass=.02, color=color.blue)
ground = box(length=5, width=5, height=.01, pos=(0,-.25,0), color=color.cyan)
support = box(length=.7, width=.5, height=.5, pos=(L,0,0), mass=.02, color=color.green)
spring = cylinder(color=color.red, pos=support.pos, axis=block.pos-support.pos, radius = 0.05)

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, U:red, K+U: magenta')
    uplot=gcurve(color=color.red)
    kplot=gcurve(color=color.blue)
    ukplot=gcurve(color=color.magenta)
    posplot=gcurve(color=color.blue)

while (t < 50):
    rate(800)
    support.pos = (L + amplitude*cos(t*frequency),0,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+(L,0,0))*(-Ks)
    block.p = block.p + (ffriction + fspring)*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.u = (.5*(Ks*((mag(block.pos))**2)))
    spring.uk = spring.u + spring.k
   

## Graph Stuff=======================================================
    if graph == 1:
        uplot.plot(pos=(t , spring.u))
        kplot.plot(pos=(t, spring.k))
        ukplot.plot(pos=(t, spring.uk))