#! /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
import sys


def get_params():
    #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_v4.dat"

    #read variables from command line if there are any
    while len(sys.argv) > 1:
        print sys.argv[0], sys.argv[1]
        option = sys.argv[1];        del sys.argv[1]
        if option == '-m':
            m = float(sys.argv[1]);  del sys.argv[1]
        if option == '-k':
            k = float(sys.argv[1]);  del sys.argv[1]
        if option == '-g':
            g = float(sys.argv[1]);  del sys.argv[1]
        if option == '-A':
            A = float(sys.argv[1]);  del sys.argv[1]
        if option == '-w':
            w = float(sys.argv[1]);  del sys.argv[1]
        if option == '-x0':
            x0 = float(sys.argv[1]);  del sys.argv[1]
        if option == '-xp0':
            xp0 = float(sys.argv[1]);  del sys.argv[1]
        if option == '-nt':
            nt = int(sys.argv[1]);   del sys.argv[1]
        if option == '-o':
            outfilename = sys.argv[1];  del sys.argv[1]

    print m,k,g,A,w,x0,xp0,nt,t_stop,outfilename
    return m,k,g,A,w,x0,xp0,nt,t_stop,outfilename


def func(y,t,A,m,w,g,k):
    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,outfilename):
    #open output file
    outfile = open(outfilename,'w')
    i = 0
    pos = 0
    nt = len(t)
    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(x)

  
def main():
    m,k,g,A,w,x0,xp0,nt,t_stop,outfilename = get_params()
    print m,k,g,A,w,x0,xp0,nt,t_stop,outfilename
    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,args=(A,m,w,g,k))
    x_arr, v_arr = hsplit(z,2)
    write_output(t,x_arr,outfilename) 
    plot_oscil(t,x_arr.ravel())
    raw_input('Please press the return key to continue...\n')


main()