def sor_diffusion_step(data, alpha, deltaX):
       n = len(data)
       m = len(data[0])
       for j in range(1, n-1):
           for k in range(1, n-1):
               nei_avg = (data[j-1][k] + data[j+1][k] + data[j][k-1]+ data[j][k+1])/4
               data[j][k] = (1-alpha)*data[j][k] + alpha*nei_avg

def adi_diffusion_step(data, alpha, mode):  #implicit on rows
   CORNERS = 0
   EDGES = 1
   n = len(data)
   m = len(data[0])

   a = [-alpha] * (n-2) + [0]
   b = [1] + (n-2) * [1+2*alpha] + [1]
   c = [0] + (n-2) * [-alpha]

   new_data = [data[0][:]]
   d = [None for index in range(n)]
   
   if mode == CORNERS:
       for j in range(1, n-1):

           for k in range(m):
               d[k] = data[j][k] + alpha*(data[j+1][k] -2*data[j][k] + data[j-1][k])
           new_stripe = tridiagonal_solver(a,b,c,d)
           new_data.append(new_stripe)
   
   elif mode == EDGES:
       for j in range(1, n-1):
           d[0] = data[j][0]
           d[m-1] = data[j][m-1]
           for k in range(1, m-1):
               d[k] = data[j][k] + alpha*(data[j+1][k] -2*data[j][k] + data[j-1][k])
           new_stripe = tridiagonal_solver(a,b,c,d)
           new_data.append(new_stripe)

   new_data.append(data[n-1][:])
   return new_data


def transpose(mat):
   return map(list, zip(*mat))


def tridiagonal_solver(a, b, c, d):
   c = c[:]
   a = [None]+a
   n = len(b)
   c[0] = c[0]/b[0]
   d[0] = d[0]/b[0]
   for i in range(1,n-1):
       m = 1/(b[i]-c[i-1]*a[i])
       c[i] = m * c[i]
       d[i] = m * (d[i] - d[i-1]*a[i])

   i = n-1
   m = 1/(b[i]-c[i-1]*a[i])
   d[i] = m * (d[i] - d[i-1]*a[i])

   x = [None]*n
   x[n-1] = d[n-1]
   for i in range(n-1)[::-1]:
       x[i] = d[i] - c[i]*x[i+1]

   return x

def finite_difference_diffusion():
  data = [0.0]*10  + [5.0] + [0.0]*10
  next_data= data[:]
  deltaT = 1.0
  deltaX = 1.0
  D = 1.0
  steps = 100

  history = [data[:]]
  for step in xrange(steps):
      for j in xrange(1, len(data)-1):
          next_data[j] = data[j] + D*deltaT*(data[j+1] - 2.0*data[j]+ data[j-1])/(deltaX**2)
      data = next_data[:]
      history.append(data[:])
  return history
      
def implicit_finit_difference_diffusion():
   n = 103
   data = [0.0]*((n-3)/4)  + [10.0] +  [0.0]*((n-3)/4)  + [20.0] +[0.0]*((n-3)/4)  + [5.0]+  [0.0]*((n-3)/4)

   next_data= data[:]
   deltaT = 1.0
   deltaX = 1.0
   D = .030

   steps = n

   history = [data[:]]
   alpha = D*deltaT/deltaX**2

   a = [-alpha]*(n-2) + [0]
   b = [1] + (n-2)*[1+2*alpha] + [1]
   c = [0] + [-alpha]*(n-2)


   for step in xrange(steps):
       data = tridiagonal_solver(a,b,c,data)
       history.append(data[:])
   return history
      

if __name__ == "__main__":

   n = 5
   alpha = 0.5

   a = [-alpha]*(n-2) + [0]
   b = [1] + (n-2)*[1+2*alpha] + [1]
   c = [0] + [-alpha]*(n-2)
   d = [1.0] * n

   print    tridiagonal_solver(a,b,c,d)