import pygame
from random import *

BREITE = 40
HOEHE = 20
LAENGE = 30

# Wahrscheinlichkeit für nach rechts:
p_rechts = BREITE/(BREITE+HOEHE)

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)

def zeichne():
    for x in range(BREITE):
        for y in range(HOEHE):
            zeichne_quadrat(x, y, farbe[zustand[x][y]])
    pygame.display.flip()
    uhr.tick(200)   

def ende_abfragen():
    for ereignis in pygame.event.get():
        if ereignis.type == pygame.QUIT:
            exit()

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],
    'partikel' : [255, 0, 0],
    'fels' : [100, 110, 130],
}

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

while True:
    if zustand[0][HOEHE-1] == 'luft':
        x, y = 0, HOEHE-1
        while True:
            zustand[x][y] = 'partikel'
            zeichne()
            ende_abfragen()
            if x == BREITE-1 or y == 0 or (x+1<=BREITE-1 and zustand[x+1][y] == 'fels') or (x-1 >= 0 and zustand[x-1][y] == 'fels') or (y+1<=HOEHE-1 and zustand[x][y+1] == 'fels') or (y-1 >= 0 and zustand[x][y-1] == 'fels'):
                zustand[x][y] = 'fels'
                break
            else:
                zustand[x][y] = 'luft'
            if random() < p_rechts:
                x = x+1
            else:
                y = y-1
    zeichne()
    ende_abfragen()