Contamination d'un cube

from manim import *
import numpy as np

class VirusCube(Cube):
    def __init__(self, *args, **kwargs):
        super(VirusCube, self).__init__(*args, **kwargs)
        self.cube_contaminated = False
        self.faces = []

    def add_face(self, cube):
        self.faces.append(cube)


    def isContaminated(self):
        return self.cube_contaminated

    def contaminate(self):
        self.cube_contaminated = True
        self.contaminated = []
        for cube in self.faces:
            if(cube.cube_contaminated == False):
                cube.cube_contaminated = True
                self.contaminated.append(cube)
        return self.contaminated 


class Virus(ThreeDScene):
    def construct(self):
        TAILLE_DU_CUBE = 3
        cubes = np.fromfunction(np.vectorize(lambda x, y, z: VirusCube()), (TAILLE_DU_CUBE, TAILLE_DU_CUBE, TAILLE_DU_CUBE), dtype=VirusCube)

        self.set_camera_orientation(phi=75 * DEGREES, theta=30 * DEGREES)
        for x in range(TAILLE_DU_CUBE) :
            for z in range(TAILLE_DU_CUBE) :
                for y in range(TAILLE_DU_CUBE) :
                    cube = cubes[x][z][y]
                    cube.set_color(GREEN)
                    cube.set_height(0.7)
                    cube.set_width(0.7)
                    self.add(cube)

                    if(x > 0): # LEFT
                        cube.add_face(cubes[x-1][z][y])

                    if(x < TAILLE_DU_CUBE-1): # RIGHT
                        cube.add_face(cubes[x+1][z][y])

                    if(y > 0): # DOWN
                        cube.add_face(cubes[x][z][y-1])

                    if(y < TAILLE_DU_CUBE-1): # UP
                        cube.add_face(cubes[x][z][y+1])
                    
                    if(z > 0): # IN
                        cube.add_face(cubes[x][z-1][y])

                    if(z < TAILLE_DU_CUBE-1): # OUT
                        cube.add_face(cubes[x][z+1][y])

                    scale = 1.5
                    cube.move_to([x*scale, z*scale, y*scale-1])
            
            # Smooth helper
            fadinator = np.vectorize(lambda cube: FadeIn(cube))
            self.play(*fadinator(cubes[x].flatten()))

        self.start_contamination(cubes)

    def start_contamination(self, cubes):
        all_cubes = cubes.flatten()
        nb_cubes = len(all_cubes)
        contaminated = 0

        stats = self.print_stats(contaminated, nb_cubes, None)

        random_cube = np.random.choice(all_cubes)
        self.color([random_cube])
        contaminated += 1
        stats = self.print_stats(contaminated, nb_cubes, stats)
        self.wait(1)
        next_contamination = random_cube.contaminate()

        while(len(next_contamination) != 0):
            contaminated += len(next_contamination)
            stats = self.print_stats(contaminated, nb_cubes, stats)
            self.color(next_contamination)
            self.wait(1)
            next_contamination = self.contaminate_one_time(next_contamination)

    def color(self, cubes):
        print("play" + str(cubes))
        for cube in cubes:
            cube.set_color(RED)

    def print_stats(self, contaminated, total, stats):
        new_stats = Text("Taux de contamination : " + str(round(contaminated/total*100, 2)) + "%")
        self.add_fixed_in_frame_mobjects(new_stats)
        new_stats.to_corner(UP+LEFT)
        if(stats == None):
            self.add(new_stats)
        else:
            self.add(new_stats)
            self.remove(stats)
        return new_stats


    def contaminate_one_time(self, to_contaminate):
        next_contamination = []
        for cube in to_contaminate:
            next_contamination.extend(cube.contaminate())
        
        return next_contamination