import pygame
from random import *

BREITE = 20
HOEHE = 20
LAENGE = 30

# Wahrscheinlichkeit für Fels:
p_fels = 0.4

def zeichne_quadrat(x, y, farbe):
    rechteck = pygame.Rect(x*LAENGE, HOEHE*LAENGE-y*LAENGE-LAENGE, LAENGE, LAENGE)
    pygame.draw.rect(surface=fenster, color=farbe, rect=rechteck)

pygame.init()
pygame.display.set_caption("Simulation: Versickerung")
uhr = pygame.time.Clock()

fenster = pygame.display.set_mode((BREITE*LAENGE, HOEHE*LAENGE))

# Dictionary: Speichert die farben für die drei Zustände eines Feldes:
farbe = {
    'luft' : [0, 0, 0],
    'fels' : [100, 110, 130],
    'wasser' : [0, 0, 200]
}

# Erstmal alle Felder auf Luft setzen.
zustand = [['luft' for y in range(HOEHE)] for x in range(BREITE)]

# Wasser in oberster Zeile.
for x in range(BREITE):
    zustand[x][HOEHE-1] = 'wasser'

# Zufällig Felsen verteilen.
for x in range(BREITE):
    for y in range(HOEHE):
        if random() < p_fels:
            zustand[x][y] = 'fels'

print('Nächster Schritt per Leertaste. Leertaste gedrückt lassen für mehrere Schritte.')

# Sorgt dafür, dass das Gedrücktlassen der Leertaste nach 400 Millisekunden
# weitere Leertaste-gedrückt-Ereignisse liefert, 
# im Abstand von 100 Millisekunden.
pygame.key.set_repeat(400, 100)

weiter = False # wird kurz True, wenn der Benutzer die Leertaste drückt.
veraenderung = True # wird False, sobald der Endzustand erreicht ist.

while True:
    for ereignis in pygame.event.get():
        if ereignis.type == pygame.QUIT:
            exit()
        elif ereignis.type == pygame.KEYDOWN:
            if ereignis.key == pygame.K_SPACE:
                print('Leertaste gedrückt')
                weiter = True
            elif ereignis.key == pygame.K_q:
                exit()
    if weiter and veraenderung:
        neu = [[zustand[x][y] for y in range(HOEHE)] for x in range(BREITE)]
        veraenderung = False
        for x in range(BREITE):
            for y in range(HOEHE):
                if neu[x][y] == 'luft':
                    if (x-1 >= 0 and zustand[x-1][y] == 'wasser') or (x+1<=BREITE-1 and zustand[x+1][y] == 'wasser') or (y-1 >= 0 and zustand[x][y-1] == 'wasser') or (y+1<=HOEHE-1 and zustand[x][y+1] == 'wasser'):
                        neu[x][y] = 'wasser'
                        veraenderung = True

        zustand = [[neu[x][y] for y in range(HOEHE)] for x in range(BREITE)]
        weiter = False

    for x in range(BREITE):
        for y in range(HOEHE):
            zeichne_quadrat(x, y, farbe[zustand[x][y]])

    pygame.display.flip()
    uhr.tick(100)