#! /usr/bin/env python
"""Always comment your scripts!
This is done in three quotes, as shown here.

This script calculates the behavior of a harmonic oscillator.

Use the well-known trick to reduce a second-order ODE to a
system of two first-order ODEs

Date: November 9, 2008

Author: Bob Wimmer

"""

from math import *   #This imports all math functions and constants
from scipy.integrate import odeint
from numpy import *


#Define parameters
m = 1.0   #mass
k = 1.0   #spring constant
g = 0.1   #damping constant
A = 1.0   #amplitude of driving force
w = 10.*pi #pi is imported from math
#Define initial conditions
x0 = 0.0  #initial position
xp0 = 1.0 #initial speed
#Define computational parameters
nt = 101      #number of timesteps for output
t_stop = 30. #time to stop
#Define filenames
outfilename = "oscil_v2.dat"


def func(y,t):
    x,v = hsplit(y,2)
    xp = v
    vp = A/m*sin(w*t) - g/m*v - k/m*x
    #print vp
    z = hstack([xp, vp])
    return z


def write_output(t,x):
    #open output file
    outfile = open(outfilename,'w')
    i = 0
    pos = 0
    while i < nt:
        outfile.write('%g \t %g \n' % (t[i], x[i]))
        i = i + 1
    outfile.close()

  
def main():
    x = x0                        #initial position
    v = xp0                       #initial speed
    y = hstack([x,v])             #array contains initial conditions
    t = linspace(0.,t_stop,nt)    #array contains timesat which to evaluate
    z = odeint(func,y,t)
    x_arr, v_arr = hsplit(z,2)
    write_output(t,x_arr) 


main()