from numpy import *
import matplotlib.pyplot as plt
from math import sqrt

#2D Waermesim mit Dirichlet-Randbedingungen
class Waermesimulation:
  def __init__(self, lengthx, lengthy, h, leftT, rightT, upperT, lowerT):
    self.zeilen = int(lengthy/h)
    self.spalten = int(lengthx/h)
    #Gebietsmatrix fuer Temperatur:
    self.T = zeros((self.zeilen+2,self.spalten+2))
    for i in range(1, self.zeilen+1):
       self.T[i, 0] = leftT
       self.T[i, self.spalten+1] = rightT
    for i in range(1, self.spalten+1):
       self.T[0, i] = lowerT
       self.T[self.zeilen+1, i] = upperT

  def jacobi(self, maxres):
    res = maxres+1.
    while res > maxres:
      tempT = copy(self.T)
      res = 0
      for i in range(1, self.zeilen+1):
        for j in range(1, self.spalten+1):
          res += (tempT[i-1,j] + tempT[i+1,j] \
               + tempT[i,j-1] + tempT[i,j+1] \
               - 4 * tempT[i,j])**2
          self.T[i,j] = 0.25*(tempT[i-1,j] + tempT[i+1,j] \
                      + tempT[i,j-1] + tempT[i,j+1])
      res = sqrt(res/(self.zeilen*self.spalten))

  def gaussseidel(self, maxres):
    res = maxres+1.
    while res > maxres:
      #tempT = copy(self.T)
      res = 0
      for i in range(1, self.zeilen+1):
        for j in range(1, self.spalten+1):
          res += (tempT[i-1,j] + tempT[i+1,j] \
               + tempT[i,j-1] + tempT[i,j+1] \
               - 4 * tempT[i,j])**2
          self.T[i,j] = 0.25*(self.T[i-1,j] + self.T[i+1,j] \
                      + self.T[i,j-1] + self.T[i,j+1])
      res = sqrt(res/(self.zeilen*self.spalten))

  def glaetter(self, n, level):
    offset = 2**level
    for iters in range(n):
      for i in range(2**level, self.zeilen+1, 2**level):
        for j in range(2**level, self.spalten+1, 2**level):
          self.T[i,j] = 0.25*(self.T[i-offset,j] + self.T[i+offset,j] \
                      + self.T[i,j-offset] + self.T[i,j+offset])

  def prolongation(self, level):
    offset = 2**level
    for i in range(2**level, self.zeilen+1, 2**(level+1)):
      for j in range(2**level, self.spalten+1, 2**(level+1)):
        self.T[i,j] = 0.25*(self.T[i-offset,j] + self.T[i+offset,j] \
                    + self.T[i,j-offset] + self.T[i,j+offset])

  def getResiduum(self):
    res = 0
    for i in range(1, self.zeilen+1):
      for j in range(1, self.spalten+1):
        res += (self.T[i-1,j] + self.T[i+1,j] \
             + self.T[i,j-1] + self.T[i,j+1] \
             - 4 * self.T[i,j])**2
    return sqrt(res/(self.zeilen*self.spalten))


  def mehrgitter_rekursiv(self, level):
    if(2**level < min(self.zeilen, self.spalten)):
      # Vorglaetten
      self.glaetter(2, level)
      # Restriktion
      # -
      # Loesung auf grobem Gitter
      self.mehrgitter_rekursiv(level+1)
      # Prolongation
      self.prolongation(level)
      # Nachglaetten
      self.glaetter(2, level)
    else:
      # oberster Level
      self.glaetter(2, level)

  def mehrgitter(self, maxres):
    res = self.getResiduum()
    while res > maxres:
      self.mehrgitter_rekursiv(0)
      res = self.getResiduum()

  def plot(self):
    print self.T
    plt.pcolormesh(self.T)
    plt.show()