import pygame
import random
import numpy as np

maxx = 50
maxy = 40
pixel_pro_einheit = 20

pause_in_millisekunden = 5

# 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)

rot = np.array([255, 0, 0])
gruen = np.array([0, 255, 0])
blau = np.array([0, 0, 255])

breite = maxx * pixel_pro_einheit
hoehe = maxy * pixel_pro_einheit

def zeichne_quadrat(x, y, farbe="white"):
    rechteck = pygame.Rect(x*pixel_pro_einheit, y*pixel_pro_einheit, pixel_pro_einheit, pixel_pro_einheit)
    pygame.draw.rect(fenster, farbe, rechteck)

def linie(x0, y0, x1, y1, farbe="white"):
    pygame.draw.line(fenster, farbe, (x0*pixel_pro_einheit, y0*pixel_pro_einheit), (x1*pixel_pro_einheit, y1*pixel_pro_einheit))

def rgb_aus_temperatur(temp):
    if temp < 50:
        s = temp/50
        rgb = (1-s)*blau + s*gruen
    else:
        s = (temp-50)/50
        rgb = (1-s)*gruen + s * rot
    return rgb.astype(int)

def zeichne_spielfeld(temperatur):
    for x in range(maxx):
        for y in range(maxy):
            farbe = rgb_aus_temperatur(temperatur[x][y])
            zeichne_quadrat(x, y, farbe)
    for x in range(maxx+1):
        linie(x, 0, x, maxy+1, "gray60")
    for y in range(maxy+1):
        linie(0, y, maxx+1, y, "gray60")
    pygame.display.flip()

def neue_temperatur(alt):
    # Kopie der alten Temperaturverteilung:
    neu = [[alt[x][y] for y in range(maxy)] for x in range(maxx)]
    
    # Aufgabe: Berechne die neue Temperaturverteilung (Variable neu) aus der alten (Variable/Parameter alt)

    return neu

print("Simulation der Wärmeleitungsgleichung")
print("Leertaste: einzelner Zeitschritt")
print("Taste r: starte/stoppe Simulation")
print(f'{delta_t=}')

pygame.init()
fenster = pygame.display.set_mode((breite+1, hoehe+1))
pygame.display.set_caption("Simulation der Wärmeleitungsgleichung")

# Erzeugt ein Array (eine Liste von Listen), alle Einträge (alias Temperaturen) sind Null.
temperatur = [[0 for y in range(maxy)] for x in range(maxx)]

# Aufgabe: Definiere die Temperaturverteilung am Anfang

zeit = 0
print(f'{zeit=}')
zeichne_spielfeld(temperatur)

ende = False
stop = True
einzelschritt = True
while not ende:
    if not stop:
        temperatur = neue_temperatur(temperatur)
        zeit += 1
        print(f'{zeit=}')
        zeichne_spielfeld(temperatur)
        if pause_in_millisekunden > 0:
            pygame.time.delay(pause_in_millisekunden)
        if einzelschritt:
            stop = True
    else:
        pygame.display.update()

    for ereignis in pygame.event.get():
        if ereignis.type == pygame.QUIT:
            ende = True
        elif ereignis.type == pygame.KEYDOWN:
            if ereignis.key in {pygame.K_q, pygame.KMOD_SHIFT+pygame.K_q, pygame.K_ESCAPE}:
                ende = True
            elif ereignis.key == pygame.K_r:
                stop = not stop
                einzelschritt = False
            elif ereignis.key == pygame.K_SPACE:
                stop = False
                einzelschritt = True

pygame.quit()