import numpy as np
from matplotlib import pyplot as plt
from matplotlib import colormaps
import matplotlib.animation as animation
from matplotlib.animation import FuncAnimation

# Grösse des quadratischen Feldes
n = 30
# Anzahl der Iterationen
iterationen = 40

# Dreidimensionale Tabelle (Array), erste Koordinate Iterationsnummer (oft = Zeit), zweite Koordinate Zeile, dritte Koordinate Spalte.
u = np.zeros((iterationen + 1, n, n))

# Temperaturverteilung am Anfang
u[0, 5, 9] = 100
u[0, 5, 10] = 100
u[0, 5, n - 10] = 100
u[0, 5, n - 9] = 100

# Start der Simulation, i für Iterationsnummer, z für Zeilennummer, s für Spaltennummer
for i in range(1, iterationen + 1): 
    for z in range(1, n - 1):
        for s in range(1, n - 1):
            u[i, z, s] = (u[i - 1, z - 1, s - 1] + u[i - 1, z - 1, s + 1]) / 2

def animiere(i):
    print(i)
    plt.clf()
    plt.title(f"Heatmap, Iteration = {i}")
    plt.xlabel("s")
    plt.ylabel("z")
    plt.pcolormesh(u[i, :, :], cmap=plt.cm.jet, vmin=0, vmax=100)
    plt.colorbar()
    # plt.show()

anim = animation.FuncAnimation(plt.figure(), animiere, interval=1, frames=iterationen, repeat=False)
anim.save("ausbreitungsbeispiel.gif")