import numpy as np
from matplotlib import pyplot as plt

# Länge des Stabes
n = 100
iterationen = 400

# alpha = Temperaturleitfähigkeit, siehe https://de.wikipedia.org/wiki/Temperaturleitf%C3%A4higkeit 
alpha = 1
delta_x = 1
# Die folgende Wahl des Zeitschritts  delta_t = Δt sorgt dafür, dass die numerische Simulation stabil läuft (ohne Begründung).  
delta_t = (delta_x ** 2)/(4 * alpha)
# Die folgende Konstante gamma hängt von delta_t ab und taucht dann in der Simulation auf.
gamma = (alpha * delta_t) / (delta_x ** 2)

# Zweidimensionale Tabelle (Array), erste Koordinate Iterationsnummer, zweite Koordinate Ort.
u = np.zeros((iterationen + 1, n))

# Temperaturverteilung am Anfang
u[0, 1] = 80
# Aufgabe: Eigene Temperaturverteilung am Anfang definieren!


# Randwerte: Temperaturen an den Enden des Stabes (vorgegeben für alle Zeiten)
for i in range(iterationen):
    u[i, 0] = 0
    u[i, n - 1] = 0

# Start der Simulation
for i in range(1, iterationen + 1):
    for x in range(1, n - 1):
        # Aufgabe: Die folgende Zeile so anpassen, dass die Temperaturentwicklung wie in der Aufgabenstellung erläutert verläuft.
        u[i, x] = u[i - 1, x - 1]

# print(u)
# Energie-/Wärme-Erhaltung (abgesehen von konstanten Randwerten)
# print([sum(z) for z in u])

plt.title(f"Heatmap eines horizontalen Stabes\n Zeitschritt Δt = {delta_t} pro Iteration")
plt.xlabel("x")
plt.ylabel(f"Iterationen")
plt.pcolormesh(u, cmap=plt.cm.jet, vmin=0, vmax=100)
plt.colorbar()
plt.savefig("stab-ort-zeit-diagramm.png")
plt.show()