percolation.py

import random

import pygame
from pygame.locals import *


class Tree:
    BURNED = 0
    BURNING = 1
    ALIVE = 2
    def __init__(self):
        self.state = self.ALIVE
    def __repr__(self):
        s = {self.BURNED:'BURNED', self.BURNING:'BURNING', self.ALIVE:'ALIVE'}[self.state]
        return '<Tree ' + s + '>'

class World:
    def __init__(self, s, p):
        self.map = [list([Tree() for i in range(s)]) for j in range(s)]


        self.burning = []
        self.new_burning = []
        self.burned = []


        self.there_is_fire = False

        self.nb_alive = s ** 2
        self.nb_burned = 0

        self.surface = pygame.Surface((s,s))
        self.surface.fill((0,255,0))

        self.font = pygame.font.Font(None,24)
        self.txt = None

        self.p = p

    def fire_on(self, pos):
        x,y = pos
        if x<0 or x>=len(self.map) or y<0 or y>=len(self.map):
            return
        if self.map[x][y].state is not Tree.ALIVE:
            return
        self.map[x][y].state = Tree.BURNING
        if not self.there_is_fire:
            self.burning = [(x,y)]
            self.there_is_fire = True
            self.nb_burned = 0


        self.nb_burned += 1
        self.new_burning.append((x,y))

    def on_update(self):
        if not self.there_is_fire:
            self.txt = self.font.render("{} % d'arbres brûlés".format(round(self.nb_burned/self.nb_alive * 100, 2)), \
                                        True, (0,0,0), (50,200,255))
            return
        elif self.txt:
            self.txt = None
        self.there_is_fire = False

        self.new_burning = []

        #self.surface.fill((0,255,0))
        burner = lambda : random.random() < self.p

        for x,y in self.burned:
            self.surface.set_at((x,y),(0,0,0))
        self.burned = []

        for x,y in self.burning:
            self.there_is_fire = True
            self.surface.set_at((x,y), (255,0,0))

            self.map[x][y].state = Tree.BURNED

            up = burner()
            down = burner()
            left = burner()
            right = burner()
            if up:
                self.fire_on((x,y+1))
            if down:
                self.fire_on((x,y-1))
            if left:
                self.fire_on((x-1,y))
            if right:
                self.fire_on((x+1,y))

            self.burned.append((x,y))
        self.burning = self.new_burning[:]

    def on_render(self,dst):
        dst.blit(self.surface, (0,0))
        if self.txt:
            dst.blit(self.txt, (0,0))


class App:
    def __init__(self):
        self.window = pygame.display.set_mode((512, 512))
        self.running = False

        self.world = World(512, 0.55)

        self.font = pygame.font.Font(None, 24)
        self.menu = self.font.render("Appuyez sur 'R' pour ré-initialiser la simulation.", \
                                     True, (0,0,0),(50,200,255))

        self.clock= pygame.time.Clock()

    def on_event(self, e):
        if e.type == QUIT:
            self.running = False
        elif e.type == MOUSEBUTTONDOWN:
            self.world.fire_on(e.pos)
        elif e.type == KEYDOWN:
            if e.key == K_r:
                p = self.world.p
                self.world = World(512,p)
            elif e.key == K_LEFT:
                self.world.p -= 0.01
                if self.world.p < 0:
                    self.world.p = 0
            elif e.key == K_RIGHT:
                self.world.p += 0.01
                if self.world.p > 1:
                    self.world.p = 1

    def on_render(self):
        self.world.on_render(self.window)
        if not self.world.there_is_fire:
            prob_txt = "Probabilité : < {} >".format(str(round(self.world.p,2)))
            prob = self.font.render(prob_txt, True, (0,0,0), (50,200,255))
            x = (self.window.get_width() - self.menu.get_width()) // 2
            y = (self.window.get_height() - self.menu.get_height()) // 2
            self.window.blit(self.menu,(x,y))
            x = (self.window.get_width() - prob.get_width()) // 2
            y += 1.2 * self.menu.get_height()
            self.window.blit(prob, (x,y))

    def on_update(self):
        self.world.on_update()

    def on_mainloop(self):
        self.running = True

        while self.running:
            for e in pygame.event.get():
                self.on_event(e)
            self.on_update()
            self.on_render()
            pygame.display.flip()
            pygame.display.set_caption(str(round(self.clock.get_fps(),2))+" FPS")
            self.clock.tick(70)

if __name__ == "__main__":
    pygame.init()
    a = App()
    a.on_mainloop()
    pygame.quit()