import sys,string
from numarray import *
import numarray.linear_algebra as la
import math
#from scipy import gplt

#define all globals
numelements = 19
h = math.pi / 20.0
r = 1.0 / sqrt(20.0)
k = r * h**2
n = [1,2,4,8,16,80,160,320,640,800]
A = zeros( [numelements,numelements], Float64 )

def setACrankNicolson():
	for i in range(len(A[0])) :
		for j in range(len(A[0])) :
			if( j == i-1 ) :
				A[i][j] = -r/2.0
			elif( j == i+1 ) :
				A[i][j] = -r/2.0
			elif( j == i ) :
				A[i][j] = (1.0 + r)

def updateVectorCrankNicolson(v):
	temp = zeros( len(v), Float64 )
	for i in range( len(v) ):
		if( i == 0 ):
			temp[i] = (r/2.0)*v[i+1] + (1-r)*v[i] 
		elif( i == len(v) - 1 ):
			temp[i] = (r/2.0)*v[i-1] + (1-r)*v[i] 
		else :
			temp[i] = (r/2.0)*v[i-1] + (r/2)*v[i+1] + (1-r)*v[i] 
	return temp

def setADouglas() :
	for i in range(len(A[0])) :
		for j in range(len(A[0])) :
			if(( j == i-1 )or(j == i+1)) :
				A[i][j] = -(1.0/2.0) * (r - (1.0/6.0))
			elif( j == i ) :
				A[i][j] = r + (5.0/6.0)

def updateVectorDouglas(v):
	temp = zeros( len(v), Float64 )
	for i in range( len(v) ):
		if( i == 0 ):
			temp[i] = (0.5 * (r + (1.0/6.0)))*v[i+1] + ((5.0/6.0) - r)*v[i] 
		elif( i == len(v) - 1 ):
			temp[i] = (0.5 * (r + (1.0/6.0)))*v[i-1] + ((5.0/6.0) - r)*v[i] 
		else :
			temp[i] = (0.5 * (r + (1.0/6.0)))*v[i+1] + (0.5 * (r + (1.0/6.0)))*v[i-1] + ((5.0/6.0) - r)*v[i] 
	return temp


def initVector(v) :
	for i in range(len(v)) :
		#print i+1,
		v[i] = math.sin((i+1)*h)

def Solver(N,type):
	Xn = zeros( numelements, Float64 )
	initVector(Xn)
	if(type == "Douglas"):
		setADouglas()
	else :
		setACrankNicolson()

	for i in range(N) :
		if(type == "Douglas") :
			Xn = updateVectorDouglas(Xn)
		else :
			Xn = updateVectorCrankNicolson(Xn)
		Xn = la.solve_linear_equations(A,Xn)

	return Xn[9]

def trueSolution(x,t):
	return math.exp(-t) * math.sin(x)

def main() :
	print "\\begin{center}"
	print "\\begin{tabular}{|l|l|l|l|} \\hline"
	print "\\multirow{2}{*}{N} & \\multirow{2}{*}{True Solution} & Crank-Nicolson & \\multirow{2}{*}{Error}\\\\"
	print " & & Douglas & \\\\ \\hline"
	for i in n :
		crank = Solver(i,"CrankNicolson")
		true = trueSolution(10.0*h,i*k)
		doug = Solver(i,"Douglas")
		print "\\multirow{2}{*}{",i,"}",
		print "&","\\multirow{2}{*}{",true,"} &",crank,"&",abs(crank-true),"\\\\"
		print "&","&",doug,"&",abs(doug-true),"\\\\ \\hline"
	print "\\end{tabular}"
	print "\\end{center}"

main()