from manim import *
import random

class Urn(Square):
    def __init__(self):
        Square.__init__(self, side_length=1.0, fill_color=GREY, fill_opacity=1.0)

class Balls(VGroup):
    colors = [BLUE, GREEN, YELLOW, PURPLE, RED, GRAY, MAROON, ORANGE, TEAL]

    def __init__(self, n):
        cols = Balls.colors
        random.shuffle(cols)
        cols = cols[:n]
        super().__init__(*[LabeledDot(str(l), color=c) for l, c in zip(list(range(n)), cols)])

class Main(Scene):
    def construct(self):

        urn = Urn().shift(3*LEFT)

        n = 9
        k = 6
        balls = Balls(n)

        left = list(range(n))
        random.shuffle(left)
        left = left[:k]
        for d in left:

            self.add(balls[d].move_to(urn), urn)
            self.play(balls[d].animate.shift(DOWN))

            self.play(urn.animate.shift(RIGHT))


        self.wait(2)