#!/usr/bin/env python
import numpy as np
cimport numpy as np

DTYPE = np.int
ctypedef np.int_t DTYPE_t

def invert_tridiagonal(np.ndarray[DTYPE_t,ndim=1] a,np.ndarray[DTYPE_t,ndim=1] b,np.ndarray[DTYPE_t,ndim=1] c,np.ndarray[DTYPE_t,ndim=1] d):
	'''
	a,b,c are numpy arrays giving three diagonals.  indices are by row, 
	so first entry of a and last entry of c are not used.
	d is numpy array giving rhs.  shape(d)=shape(b)
	'''
	cdef int N = np.shape(c)[-1]
	cdef np.ndarray cnew = np.empty(N), dnew = np.empty(N), u = np.empty(N)
	cnew[0] = c[0]/b[0]; dnew[0] = d[0]/b[0]
	cdef int i
	for i in range(N-1):
		cnew[i+1] = c[i+1]/(b[i+1]-a[i+1]*cnew[i]) #note: this will set cnew[N], but we don't use it.
		dnew[i+1] = (d[i+1]-a[i+1]*dnew[i])/(b[i+1]-a[i+1]*cnew[i])
	u[N-1] = dnew[N-1]
	for i in range(N-1)[::-1]:
		u[i] = dnew[i] - cnew[i]*u[i+1]
	return u

def tridiagonal_system(np.ndarray[DTYPE_t,ndim=1] a,np.ndarray[DTYPE_t,ndim=1] b,np.ndarray[DTYPE_t,ndim=1] c,np.ndarray[DTYPE_t,ndim=2] d):
	'''
	iterate the above routine along first axis of d
	shape(d) = (N,N)
	'''
	cdef int i
	N = np.shape(d)[0]
	u = np.empty((N,N))
	for i in range(N):
		u[i] = invert_tridiagonal(a,b,c,d[i])
	return u

def sor(np.ndarray[DTYPE_t,ndim=2] u, np.ndarray[DTYPE_t,ndim=2] rho, DTYPE_t alpha, DTYPE_t dx):
	'''
	successive over-relaxation for poisson's equation.
	assume u has boundary conditions enforced and don't touch boundaries of unew
	'''
	cdef int i,j
	cdef np.ndarray[DTYPE_t,ndim=2] unew = np.copy(u)
	cdef n = np.shape(u)[0], m = np.shape(u)[1]
	for i in range(1,n-1):
		for j in range(1,m-1):
			unew[i,j] = (1-alpha)*u[i,j] \
			+ .25*alpha*(u[i+1,j]+unew[i-1,j]+u[i,j+1]+unew[i,j-1]) \
			+ .25*alpha*dx*dx*rho[i,j]
	return unew