#! /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 *
import Gnuplot, Gnuplot.funcutils

#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 = 111      #number of timesteps for output
t_stop = 48. #time to stop
#Define filenames
outfilename = "oscil_v3.dat"


def func(y,t):
    x,v = hsplit(y,2)
    xp = v
    #vp = - g/m*v - k/m*x
    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 plot_oscil(t,x):
    g = Gnuplot.Gnuplot()
    data = Gnuplot.Data(t,x,using=(1,2)) #this ensures that t is used as x axis
    g('set ylabel "y-axis [arb. units]"')
    g('set xlabel "x-axis [arb. units]"')
    g('set style data lines')
    g.plot(data)

  
def main():
    x = x0                        #initial position
    v = xp0                       #initial speed
    y = hstack([x,v])             #array contains initial conditions
    print x
    print v
    print y
    t = linspace(0.,t_stop,nt)    #array contains timesat which to evaluate
    z = odeint(func,y,t)
    #print z
    #print '------------------'
    x_arr, v_arr = hsplit(z,2)
    write_output(t,x_arr)
    plot_oscil(t,x_arr.ravel())
    raw_input('Please press the return key to continue...\n')


main()