3b1b-manim/topics/counting.py

from constants import *

from mobject.mobject import Mobject
from mobject.tex_mobject import TexMobject
from mobject.tex_mobject import TextMobject
from mobject.vectorized_mobject import VGroup
from mobject.vectorized_mobject import VMobject

from animation.creation import ShowCreation
from animation.creation import FadeIn
from animation.transform import MoveToTarget
from animation.transform import Transform
from topics.geometry import Arrow
from topics.geometry import Circle
from topics.geometry import Dot
# from topics.fractals import *
from scene.scene import Scene


class CountingScene(Scene):
    CONFIG = {
        "digit_place_colors" : [YELLOW, MAROON_B, RED, GREEN, BLUE, PURPLE_D],
        "counting_dot_starting_position" : (FRAME_X_RADIUS-1)*RIGHT + (FRAME_Y_RADIUS-1)*UP,
        "count_dot_starting_radius" : 0.5,
        "dot_configuration_height" : 2,
        "ones_configuration_location" : UP+2*RIGHT,
        "num_scale_factor" : 2,
        "num_start_location" : 2*DOWN,
    }
    def setup(self):
        self.dots = VGroup()
        self.number = 0        
        self.max_place = 0
        self.number_mob = VGroup(TexMobject(str(self.number)))
        self.number_mob.scale(self.num_scale_factor)
        self.number_mob.shift(self.num_start_location)

        self.dot_templates = []
        self.dot_template_iterators = []
        self.curr_configurations = []

        self.arrows = VGroup()

        self.add(self.number_mob)

    def get_template_configuration(self, place):
        #This should probably be replaced for non-base-10 counting scenes
        down_right = (0.5)*RIGHT + (np.sqrt(3)/2)*DOWN
        result = []
        for down_right_steps in range(5):
            for left_steps in range(down_right_steps):
                result.append(
                    down_right_steps*down_right + left_steps*LEFT
                )
        return reversed(result[:self.get_place_max(place)])

    def get_dot_template(self, place):
        #This should be replaced for non-base-10 counting scenes
        down_right = (0.5)*RIGHT + (np.sqrt(3)/2)*DOWN
        dots = VGroup(*[
            Dot(
                point, 
                radius = 0.25,
                fill_opacity = 0,
                stroke_width = 2,
                stroke_color = WHITE,
            )
            for point in self.get_template_configuration(place)
        ])
        dots.scale_to_fit_height(self.dot_configuration_height)
        return dots

    def add_configuration(self):
        new_template = self.get_dot_template(len(self.dot_templates))
        new_template.move_to(self.ones_configuration_location)
        left_vect = (new_template.get_width()+LARGE_BUFF)*LEFT
        new_template.shift(
            left_vect*len(self.dot_templates)
        )
        self.dot_templates.append(new_template)
        self.dot_template_iterators.append(
            it.cycle(new_template)
        )
        self.curr_configurations.append(VGroup())

    def count(self, max_val, run_time_per_anim = 1):
        for x in range(max_val):
            self.increment(run_time_per_anim)

    def increment(self, run_time_per_anim = 1):
        moving_dot = Dot(
            self.counting_dot_starting_position,
            radius = self.count_dot_starting_radius,
            color = self.digit_place_colors[0],
        )
        moving_dot.generate_target()
        moving_dot.set_fill(opacity = 0)
        kwargs = {
            "run_time" : run_time_per_anim
        }

        continue_rolling_over = True
        first_move = True
        place = 0
        while continue_rolling_over:
            added_anims = []                
            if first_move:
                added_anims += self.get_digit_increment_animations()
                first_move = False
            moving_dot.target.replace(
                self.dot_template_iterators[place].next()
            )
            self.play(MoveToTarget(moving_dot), *added_anims, **kwargs)
            self.curr_configurations[place].add(moving_dot)


            if len(self.curr_configurations[place].split()) == self.get_place_max(place):
                full_configuration = self.curr_configurations[place]
                self.curr_configurations[place] = VGroup()
                place += 1
                center = full_configuration.get_center_of_mass()
                radius = 0.6*max(
                    full_configuration.get_width(),
                    full_configuration.get_height(),
                )
                circle = Circle(
                    radius = radius,
                    stroke_width = 0,
                    fill_color = self.digit_place_colors[place],
                    fill_opacity = 0.5,
                )
                circle.move_to(center)
                moving_dot = VGroup(circle, full_configuration)
                moving_dot.generate_target()
                moving_dot[0].set_fill(opacity = 0)
            else:
                continue_rolling_over = False

    def get_digit_increment_animations(self):
        result = []
        self.number += 1
        is_next_digit = self.is_next_digit()
        if is_next_digit: self.max_place += 1
        new_number_mob = self.get_number_mob(self.number)
        new_number_mob.move_to(self.number_mob, RIGHT)
        if is_next_digit:
            self.add_configuration()
            place = len(new_number_mob.split())-1
            result.append(FadeIn(self.dot_templates[place]))
            arrow = Arrow(
                new_number_mob[place].get_top(),
                self.dot_templates[place].get_bottom(),
                color = self.digit_place_colors[place]
            )
            self.arrows.add(arrow)
            result.append(ShowCreation(arrow))
        result.append(Transform(
            self.number_mob, new_number_mob,
            submobject_mode = "lagged_start"
        ))
        return result

    def get_number_mob(self, num):
        result = VGroup()
        place = 0
        max_place = self.max_place
        while place < max_place:
            digit = TexMobject(str(self.get_place_num(num, place)))
            if place >= len(self.digit_place_colors):
                self.digit_place_colors += self.digit_place_colors
            digit.set_color(self.digit_place_colors[place])
            digit.scale(self.num_scale_factor)
            digit.next_to(result, LEFT, buff = SMALL_BUFF, aligned_edge = DOWN)
            result.add(digit)
            place += 1
        return result

    def is_next_digit(self):
        return False
    def get_place_num(self, num, place):
        return 0
    def get_place_max(self, place):
        return 0

class PowerCounter(CountingScene):
    def is_next_digit(self):
        number = self.number
        while number > 1:
            if number%self.base != 0:
                return False
            number /= self.base
        return True
    def get_place_max(self, place):
        return self.base
    def get_place_num(self, num, place):
        return (num / (self.base ** place)) % self.base

class CountInDecimal(PowerCounter):
    CONFIG = {
        "base" : 10,
    }
    def construct(self):
        for x in range(11):
            self.increment()
        for x in range(85):
            self.increment(0.25)
        for x in range(20):
            self.increment()

class CountInTernary(PowerCounter):
    CONFIG = {
        "base" : 3,
        "dot_configuration_height" : 1,
        "ones_configuration_location" : UP+4*RIGHT
    }
    def construct(self):
        self.count(27)

    # def get_template_configuration(self):
    #     return [ORIGIN, UP]

class CountInBinaryTo256(PowerCounter):
    CONFIG = {
        "base" : 2,
        "dot_configuration_height" : 1,
        "ones_configuration_location" : UP+5*RIGHT
    }
    def construct(self):
        self.count(128, 0.3)

    def get_template_configuration(self):
        return [ORIGIN, UP]

class FactorialBase(CountingScene):
    CONFIG = {
        "dot_configuration_height" : 1,
        "ones_configuration_location" : UP+4*RIGHT
    }
    def construct(self):
        self.count(30, 0.4)

    def is_next_digit(self):
        return self.number == self.factorial(self.max_place + 1)

    def get_place_max(self, place):
        return place + 2

    def get_place_num(self, num, place):
        return (num / self.factorial(place + 1)) % self.get_place_max(place)

    def factorial(self, n):
        if (n == 1): return 1
        else: return n * self.factorial(n - 1)