#!/usr/bin/python
from pylab import *
import sys

print sys.argv
p = 0.5
q = 0.05
if(len(sys.argv) > 1):
    p = float(sys.argv[1])
if(len(sys.argv) > 2):
    q = float(sys.argv[2])

# This isn't used in the code, it's just here for reference.
# Function describing the initial conditions:
def f(x,p,q):
    if(x <= p):
        return  q/p*x
    elif(x > p):
        return  (q+q*p/(1-p)) - q/(1-p)*x

# This is the definition for the An coefficients (all Bn are zero):
def a(n,p,q):
    if(n==0):
        an = 0
    else:
        m = float(float(q)/float(p))
        b = float(0)
        an1 = -2*b*cos(pi*n*p)/(pi*n) + 2*m*sin(pi*n*p)/(pi**2*n**2) - 2*m*p*cos(pi*n*p)/(pi*n) + 2*b/(pi*n)
        #print m,b,an1
        m = float(-q/(1-p))
        b = float(q+q*p/(1-p))
        an2 = -2*b*cos(pi*n)/(pi*n) - 2*m*cos(pi*n)/(pi*n) - 2*m*sin(pi*n*p)/(pi**2*n**2) + 2*b*cos(pi*n*p)/(pi*n) + 2*m*sin(pi*n)/(pi**2*n**2) + 2*m*p*cos(pi*n*p)/(pi*n)
        #print m,b,an2
        an = an1 + an2
        #print an
    return an

# A function of both time and position describing the displacement from eq
def un(n,t,x,p,q):
    u = 0
    for i in range(n):
        u += float(a(i,p,q))*cos(i*pi*t)*sin(i*pi*x)
    return u

# We want to plot vs time at fixed x, and plot vs x at fixed t
t = arange(0,65.536,0.001)
x = arange(0,1,0.001)
u = un(1000,t,0.75,p,q)
u0 = un(1000,0,x,p,q)

# We also want to look at the FFT
ufft = abs(fft(u).real)


# This will make us some reasonably pretty plots.
subplots_adjust(hspace=0.4)

subplot(311)
plot(x,u0)
title('Initial Conditions')

subplot(312)
plot(t,u)
title('Displacement w/ Time')

subplot(313)
ufft = ufft[0:400]
plot(range(len(ufft)),ufft)
title('FFT of Displacement')

show()