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
iterationen = 60

# alpha = Temperaturleitfähigkeit, siehe https://de.wikipedia.org/wiki/Temperaturleitf%C3%A4higkeit 
alpha = 2
delta_x = 1
delta_t = (delta_x ** 2)/(4 * alpha)
gamma = (alpha * delta_t) / (delta_x ** 2)

# Dreidimensionale Tabelle (Array), erste Koordinate Iterationsnummer, zweite Koordinate Zeile, dritte Koordinate Spalte.

u = np.zeros((iterationen + 1, n, n))

# Randwerte: Temperaturen am Rand des Quadrates (vorgegeben für alle Zeiten)
for i in range(iterationen):
    for s in range(n):
        u[i, 0, s] = 100
    # Aufgabe: interessantere Randwerte vereinbaren! Etwa Raum mit offenem Fenster und Heizungen an den Wänden. 

# Anfangswerte
# Aufgabe: Wer mag, definiert interessante Anfangswerte.
u[0, n // 2, n // 2] = 100

# Start der Simulation
for i in range(1, iterationen + 1):
    for z in range(1, n - 1):
        for s in range(1, n - 1):
            # Aufgabe: Wärmeausbreitungsformel anpassen!
            u[i, z, s] = u[i - 1, z - 1, s]

def animiere(i):
    print(i)
    plt.clf()
    plt.title(f"Heatmap, Zeit t = {i * delta_t:.3f}")
    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("feld.gif")