#! /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 10, 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, time


class Oscillator:
    """harmonic or not so harmoic oscillator using scipy."""
    def __init__(self,**kwargs):
        """Initialize parameters from keyword arguments if given."""
        self.m = 1.0
        self.k = 1.0
        self.g = 0.1   #damping constant
        self.A = 1.0   #amplitude of driving force
        self.w = 10.*pi #pi is imported from math
        #Define initial conditions
        self.x0 = 0.0  #initial position
        self.xp0 = 1.0 #initial speed
        #Define computational parameters
        self.nt = 501      #number of timesteps for output
        self.t_stop = 30. #time to stop
        #Define filenames
        self.outfilename = "oscil_v5.dat"
        self.scan()


    def scan(self):
        """Read parameters from standard input."""
        while len(sys.argv) > 1:
            #print sys.argv[0], sys.argv[1]
            option = sys.argv[1];        del sys.argv[1]
            if option == '-m':
                self.m = float(sys.argv[1]);  del sys.argv[1]
            if option == '-k':
                self.k = float(sys.argv[1]);  del sys.argv[1]
            if option == '-g':
                self.g = float(sys.argv[1]);  del sys.argv[1]
            if option == '-A':
                self.A = float(sys.argv[1]);  del sys.argv[1]
            if option == '-w':
                self.w = float(sys.argv[1]);  del sys.argv[1]
            if option == '-x0':
                self.x0 = float(sys.argv[1]);  del sys.argv[1]
            if option == '-xp0':
                self.xp0 = float(sys.argv[1]);  del sys.argv[1]
            if option == '-nt':
                self.nt = int(sys.argv[1]);   del sys.argv[1]
            if option == '-o':
                self.outfilename = sys.argv[1];  del sys.argv[1]


    def solve(self):
        """Solve the oscillator ODE system with odeint."""
        self.x = self.x0                        #initial position
        self.v = self.xp0                       #initial speed
        self.y = hstack([self.x,self.v])             #array contains initial conditions
        self.t = linspace(0.,self.t_stop,self.nt)    #array contains timesat which to evaluate
        self.z = odeint(self.func,self.y,self.t)
        self.x_arr, self.v_arr = hsplit(self.z,2)

                  


    def func(self,y,t):
        """Return the right-hand side of the ODE."""
        x,v = hsplit(y,2)
        xp = v
        vp = self.A/self.m*sin(self.w*t) - self.g/self.m*v - self.k/self.m*x
        z = hstack([xp, vp])
        return z


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


    def plot(self):
        """Plot the results with Gnuplot."""
        data = Gnuplot.Data(self.t,self.x_arr,using=(1,2)) #this ensures that t is used as x axis
        g = Gnuplot.Gnuplot(persist=True)
        g('set ylabel "y-axis [arb. units]"')
        g('set xlabel "x-axis [arb. units]"')
        g('set style data lines')
        g.plot(data)

  
def main():
    s = Oscillator(m=1.0)
    t1 = time.clock()
    s.solve()
    t2 = time.clock()
    s.plot()
    t3 = time.clock()
    s.write_output()
    t4 = time.clock()
    print "Timing information:"
    print "integration took \t"+str(t2-t1)+"\t seconds"
    print "plotting took \t"+str(t3-t2)+"\t seconds"
    print "writing to file took \t"+str(t4-t3)+"\t seconds"
    print "and all togther took \t"+str(t4-t1)+"\t seconds"

    print 'CPU time of oscillator:', t2 - t1

    raw_input('Please press the return key to continue...\n')


main()