from __future__ import division

import numpy as np
import scipy as sp
import matplotlib as mpl
import matplotlib.pyplot as plt

def calc_params(nx, nt, alpha, length, tmax):
    # Compute mesh spacing and time step.
    dx = length / (nx-1)
    dt = tmax / (nt-1)
    return (dx, dt)

def create_arrays(nx, nt):
    # Create arrays to save data in process.
    x = np.linspace(0, length+1e-15, nx)
    t = np.linspace(0, tmax+1e-15, nt)
    u = np.zeros([nx, nt])
    return (x, t, u)

def set_ic(x, length):
    # Set initial and boundary conditions.
    u[:, 0] = np.sin(np.pi*x/length)**2
    #boundaries are implicitly set by this initial condition
    return u

def plot_results(u, x):
    # Plot the results (initial and final conditions).
    mpl.rcParams['figure.figsize']=[15,3]
    plt.plot(x, u[:,0], 'k-', lw=2)
    plt.plot(x, u[:,-1],  'b-', lw=2)
    
    #import matplotlib.cm as cm
    #for i in range(len(u[:,0:-1:10])-1):
    #    c = cm.rainbow(i/len(u[:,0:-1:10]), 1)
    #    plt.plot(x[:], u[:,i], color=c, lw=2)
    
    plt.ylim((0,1))
    plt.xlim((0,1))
    plt.show()

def gen_matrix(nx, alpha, dt, dx):
    r  = alpha * dt / (dx*dx)
    s  = 1 - 2*r
    A = np.zeros((nx, nx), dtype=np.float128)
    i,j = np.indices(A.shape)
    A[i==j]   = s
    A[i==j-1] = r
    A[i==j+1] = r
    return A

# Set default values for simulation parameters.
nt = 10
nx = 20
alpha = 0.1
length = 1.0
tmax = 0.5

# Get command-line arguments.
usage = """
This program can be invoked with the following options:
    -t, --nt=       number of time steps
    -x, --nx=       number of spatial intervals
    -a, --alpha=    value of alpha in heat equation
    -l, --length=   length of beam (m)
    -m, --tmax=     maximum time value calculated

Example usage:
    python heat_eqn_cl.py --nt=10 --nx=20 --alpha=0.1 --length=1.0 --tmax=0.5

"""
import getopt, sys
options, args = getopt.getopt(sys.argv[1:],'ht:x:a:l:m:', \
                              ['help','nt=','nx=','alpha=','length=','tmax='])
for option, value in options:
    if option in ('-h','--help'):
        print(usage)
        sys.exit(0)
    elif option in ('-t', '--nt'):
        nt = int(value)
    elif option in ('-x', '--nx'):
        nx = int(value)
    elif option in ('-a', '--alpha'):
        alpha = float(value)
    elif option in ('-l', '--length'):
        length = float(value)
    elif option in ('-m', '--tmax'):
        tmax = float(value)
    else:
        print(usage)
        sys.exit('Invalid command-line argument detected.')

(dx, dt) = calc_params(nx, nt, alpha, length, tmax)
(x, t, u ) = create_arrays(nx, nt)
u = set_ic(x, length)
A = gen_matrix(nx, alpha, dt, dx)

# Loop through each time step.
for n in range(1, nt):
    u[:,n] = A.dot(u[:,n-1])

plot_results(u, x)