3b1b-manim/active_projects/clacks/solution2/position_phase_space.py

from big_ol_pile_of_manim_imports import *
from active_projects.clacks.question import Block
from active_projects.clacks.question import Wall
from active_projects.clacks.question import ClackFlashes


class PositionPhaseSpaceScene(Scene):
    CONFIG = {
        "rescale_coordinates": True,
        "wall_x": -6,
        "wall_config": {
            "height": 1.6,
            "tick_spacing": 0.35,
            "tick_length": 0.2,
        },
        "wall_height": 1.5,
        "floor_y": -3.5,
        "block1_config": {
            "mass": 10,
            "distance": 9,
            "velocity": 1,
            "width": 1.6,
        },
        "block2_config": {
            "mass": 1,
            "distance": 4,
        },
        "axes_config": {
            "x_min": -0.5,
            "x_max": 31,
            "y_min": -0.5,
            "y_max": 10.5,
            "x_axis_config": {
                "unit_size": 0.4,
                "tick_frequency": 2,
            },
            "y_axis_config": {
                "unit_size": 0.4,
                "tick_frequency": 2,
            },
        },
        "axes_center": 5 * LEFT + 0.65 * DOWN,
        "ps_dot_config": {
            "fill_color": RED,
            "background_stroke_width": 1,
            "background_stroke_color": BLACK,
            "radius": 0.05,
        },
        "ps_d2_label_vect": RIGHT,
        "clack_sound": "clack",
        "mirror_line_class": Line,
        "mirror_line_style": {
            "stroke_color": WHITE,
            "stroke_width": 1,
        },
        "d1_eq_e2_line_color": GREEN_SCREEN,
        "trajectory_style": {
            "stroke_color": YELLOW,
            "stroke_width": 2,
        }
    }

    def setup(self):
        self.total_sliding_time = 0
        self.all_items = [
            self.get_floor(),
            self.get_wall(),
            self.get_blocks(),
            self.get_axes(),
            self.get_phase_space_point(),
            self.get_phase_space_x_line(),
            self.get_phase_space_y_line(),
            self.get_phase_space_dot(),
            self.get_phase_space_d1_label(),
            self.get_phase_space_d2_label(),
            self.get_d1_brace(),
            self.get_d2_brace(),
            self.get_d1_label(),
            self.get_d2_label(),
            self.get_d1_eq_d2_line(),
            self.get_d1_eq_d2_label(),
            self.get_d2_eq_w2_line(),
            self.get_d2_eq_w2_label(),
        ]

    def get_floor_wall_corner(self):
        return self.wall_x * RIGHT + self.floor_y * UP

    def get_mass_ratio(self):
        return op.truediv(
            self.block1.mass,
            self.block2.mass,
        )

    def d1_to_x(self, d1):
        if self.rescale_coordinates:
            d1 *= np.sqrt(self.block1.mass)
        return d1

    def d2_to_y(self, d2):
        if self.rescale_coordinates:
            d2 *= np.sqrt(self.block2.mass)
        return d2

    def ds_to_point(self, d1, d2):
        return self.axes.coords_to_point(
            self.d1_to_x(d1), self.d2_to_y(d2),
        )

    def point_to_ds(self, point):
        x, y = self.axes.point_to_coords(point)
        if self.rescale_coordinates:
            x /= np.sqrt(self.block1.mass)
            y /= np.sqrt(self.block2.mass)
        return (x, y)

    def get_d1(self):
        return self.get_ds()[0]

    def get_d2(self):
        return self.get_ds()[1]

    def get_ds(self):
        return self.point_to_ds(self.ps_point.get_location())

    # Relevant for sliding
    def tie_blocks_to_ps_point(self):
        def update_blocks(blocks):
            d1, d2 = self.point_to_ds(self.ps_point.get_location())
            b1, b2 = blocks
            corner = self.get_floor_wall_corner()
            b1.move_to(corner + d1 * RIGHT, DL)
            b2.move_to(corner + d2 * RIGHT, DR)
        self.blocks.add_updater(update_blocks)

    def time_to_ds(self, time):
        # Deals in its own phase space, different
        # from the one displayed
        m1 = self.block1.mass
        m2 = self.block2.mass
        v1 = self.block1.velocity
        start_d1 = self.block1_config["distance"]
        start_d2 = self.block2_config["distance"]
        w2 = self.block2.width
        start_d2 += w2
        ps_speed = np.sqrt(m1) * abs(v1)
        theta = np.arctan(np.sqrt(m2 / m1))

        def ds_to_ps_point(d1, d2):
            return np.array([
                d1 * np.sqrt(m1),
                d2 * np.sqrt(m2),
                0
            ])

        def ps_point_to_ds(point):
            return (
                point[0] / np.sqrt(m1),
                point[1] / np.sqrt(m2),
            )

        ps_point = ds_to_ps_point(start_d1, start_d2)
        wedge_corner = ds_to_ps_point(w2, w2)
        ps_point -= wedge_corner
        # Pass into the mirror worlds
        ps_point += time * ps_speed * LEFT
        # Reflect back to the real world
        angle = angle_of_vector(ps_point)
        n = int(angle / theta)
        if n % 2 == 0:
            ps_point = rotate_vector(ps_point, -n * theta)
        else:
            ps_point = rotate_vector(
                ps_point,
                -(n + 1) * theta,
            )
            ps_point[1] = abs(ps_point[1])
        ps_point += wedge_corner
        return ps_point_to_ds(ps_point)

    def get_clack_data(self):
        # Copying from time_to_ds.  Not great, but
        # maybe I'll factor things out properly later.
        m1 = self.block1.mass
        m2 = self.block2.mass
        v1 = self.block1.velocity
        w2 = self.block2.get_width()
        h2 = self.block2.get_height()
        ps_speed = np.sqrt(m1) * abs(v1)
        theta = np.arctan(np.sqrt(m2 / m1))

        def ds_to_ps_point(d1, d2):
            return np.array([
                d1 * np.sqrt(m1),
                d2 * np.sqrt(m2),
                0
            ])

        def ps_point_to_ds(point):
            return (
                point[0] / np.sqrt(m1),
                point[1] / np.sqrt(m2),
            )

        ps_point = ds_to_ps_point(*self.get_ds())
        wedge_corner = ds_to_ps_point(w2, w2)
        ps_point -= wedge_corner
        y = ps_point[1]

        clack_data = []
        for k in range(1, int(PI / theta) + 1):
            clack_ps_point = np.array([
                y / np.tan(k * theta), y, 0
            ])
            time = get_norm(ps_point - clack_ps_point) / ps_speed
            reflected_point = rotate_vector(
                clack_ps_point,
                -2 * np.ceil((k - 1) / 2) * theta
            )
            d1, d2 = ps_point_to_ds(reflected_point + wedge_corner)
            loc1 = self.get_floor_wall_corner() + h2 * UP / 2 + d2 * RIGHT
            if k % 2 == 0:
                loc1 += w2 * LEFT
            loc2 = self.ds_to_point(d1, d2)
            clack_data.append((time, loc1, loc2))
        return clack_data

    def get_clack_flashes(self):
        pass  # TODO

    def tie_ps_point_to_time_tracker(self):
        time_tracker = self.get_time_tracker(
            time=self.total_sliding_time
        )

        def update_ps_point(p):
            time = time_tracker.get_value()
            ds = self.time_to_ds(time)
            p.move_to(self.ds_to_point(*ds))
        self.ps_point.add_updater(update_ps_point)
        self.add(time_tracker, self.ps_point)

    def add_clack_flashes(self):
        clack_data = self.get_clack_data()
        self.clack_times = [
            time for (time, loc1, loc2) in clack_data
        ]
        self.block_flashes = ClackFlashes([
            (loc1, time)
            for (time, loc1, loc2) in clack_data
        ])
        self.ps_flashes = ClackFlashes([
            (loc2, time)
            for (time, loc1, loc2) in clack_data
        ])
        self.add(
            self.block_flashes,
            self.ps_flashes,
        )

    def get_continually_building_trajectory(self):
        trajectory = VMobject()
        self.trajectory = trajectory
        trajectory.set_style(**self.trajectory_style)

        def get_point():
            return np.array(self.ps_point.get_location())

        points = [get_point(), get_point()]
        trajectory.set_points_as_corners(points)
        epsilon = 0.001

        def update_trajectory(trajectory):
            new_point = get_point()
            p1, p2 = trajectory.get_anchors()[-2:]
            angle = angle_between_vectors(
                p2 - p1,
                new_point - p2,
            )
            if angle > epsilon:
                points.append(new_point)
            else:
                points[-1] = new_point
            trajectory.set_points_as_corners(points)

        trajectory.add_updater(update_trajectory)
        return trajectory

    def begin_sliding(self, show_trajectory=True):
        self.tie_ps_point_to_time_tracker()
        self.add_clack_flashes()
        if show_trajectory:
            if hasattr(self, "trajectory"):
                self.trajectory.resume_updating()
            else:
                self.add(self.get_continually_building_trajectory())

    def end_sliding(self):
        self.ps_point.clear_updaters()
        self.remove(self.time_tracker)
        to_remove = ["block_flashes", "ps_flashes"]
        for attr in to_remove:
            if hasattr(self, attr):
                self.remove(getattr(self, attr))
        if hasattr(self, "trajectory"):
            self.trajectory.suspend_updating()
        total_time = self.time_tracker.get_value()
        self.total_sliding_time += total_time
        for time in self.clack_times:
            if time < total_time:
                offset = total_time - time
                self.add_sound(
                    "clack",
                    time_offset=-offset,
                )

    def slide(self, time, stop_condition=None):
        self.begin_sliding()
        self.wait(time, stop_condition)
        self.end_sliding()

    def slide_until(self, stop_condition, max_time=60):
        self.slide(max_time, stop_condition=stop_condition)

    def get_ps_point_change_anim(self, d1, d2, **added_kwargs):
        b1 = self.block1
        ps_speed = np.sqrt(b1.mass) * abs(b1.velocity)
        curr_d1, curr_d2 = self.get_ds()
        distance = get_norm([curr_d1 - d1, curr_d2 - d2])

        # Default
        kwargs = {
            "run_time": (distance / ps_speed),
            "rate_func": None,
        }
        kwargs.update(added_kwargs)
        return ApplyMethod(
            self.ps_point.move_to,
            self.ds_to_point(d1, d2),
            **kwargs
        )

    # Mobject getters
    def get_floor(self):
        floor = self.floor = Line(
            self.wall_x * RIGHT,
            FRAME_WIDTH * RIGHT / 2,
            stroke_color=WHITE,
            stroke_width=3,
        )
        floor.move_to(self.get_floor_wall_corner(), LEFT)
        return floor

    def get_wall(self):
        wall = self.wall = Wall(**self.wall_config)
        wall.move_to(self.get_floor_wall_corner(), DR)
        return wall

    def get_blocks(self):
        blocks = self.blocks = VGroup()
        for n in [1, 2]:
            config = getattr(self, "block{}_config".format(n))
            block = Block(**config)
            block.move_to(self.get_floor_wall_corner(), DL)
            block.shift(config["distance"] * RIGHT)
            block.label.move_to(block)
            block.label.set_fill(BLACK)
            block.label.set_stroke(WHITE, 1, background=True)
            self.blocks.add(block)
        self.block1, self.block2 = blocks
        return blocks

    def get_axes(self):
        axes = self.axes = Axes(**self.axes_config)
        axes.set_stroke(LIGHT_GREY, 2)
        axes.shift(
            self.axes_center - axes.coords_to_point(0, 0)
        )
        axes.labels = self.get_axes_labels(axes)
        axes.add(axes.labels)
        axes.added_lines = self.get_added_axes_lines(axes)
        axes.add(axes.added_lines)
        return axes

    def get_added_axes_lines(self, axes):
        c2p = axes.coords_to_point
        x_mult = y_mult = 1
        if self.rescale_coordinates:
            x_mult = np.sqrt(self.block1.mass)
            y_mult = np.sqrt(self.block2.mass)
        y_lines = VGroup(*[
            Line(
                c2p(0, 0), c2p(0, axes.y_max * y_mult + 1),
            ).move_to(c2p(x, 0), DOWN)
            for x in np.arange(0, axes.x_max) * x_mult
        ])
        x_lines = VGroup(*[
            Line(
                c2p(0, 0), c2p(axes.x_max * x_mult, 0),
            ).move_to(c2p(0, y), LEFT)
            for y in np.arange(0, axes.y_max) * y_mult
        ])
        line_groups = VGroup(x_lines, y_lines)
        for lines in line_groups:
            lines.set_stroke(BLUE, 1, 0.5)
            lines[1::2].set_stroke(width=0.5, opacity=0.25)
        return line_groups

    def get_axes_labels(self, axes, with_sqrts=None):
        if with_sqrts is None:
            with_sqrts = self.rescale_coordinates
        x_label = TexMobject("x = ", "d_1")
        y_label = TexMobject("y = ", "d_2")
        labels = VGroup(x_label, y_label)
        if with_sqrts:
            additions = map(TexMobject, [
                "\\sqrt{m_1}", "\\sqrt{m_2}"
            ])
            for label, addition in zip(labels, additions):
                addition.move_to(label[1], DL)
                label[1].next_to(
                    addition, RIGHT, SMALL_BUFF,
                    aligned_edge=DOWN
                )
                addition[2:].set_color(BLUE)
                label.add(addition)
        x_label.next_to(axes.x_axis.get_right(), DL, MED_SMALL_BUFF)
        y_label.next_to(axes.y_axis.get_top(), DR, MED_SMALL_BUFF)
        for label in labels:
            label.shift_onto_screen()
        return labels

    def get_phase_space_point(self):
        ps_point = self.ps_point = VectorizedPoint()
        ps_point.move_to(self.ds_to_point(
            self.block1.distance,
            self.block2.distance + self.block2.width
        ))
        self.tie_blocks_to_ps_point()
        return ps_point

    def get_phase_space_x_line(self):
        def get_x_line():
            origin = self.axes.coords_to_point(0, 0)
            point = self.ps_point.get_location()
            y_axis_point = np.array(origin)
            y_axis_point[1] = point[1]
            return DashedLine(
                y_axis_point, point,
                color=GREEN,
                stroke_width=2,
            )
        self.x_line = updating_mobject_from_func(get_x_line)
        return self.x_line

    def get_phase_space_y_line(self):
        def get_y_line():
            origin = self.axes.coords_to_point(0, 0)
            point = self.ps_point.get_location()
            x_axis_point = np.array(origin)
            x_axis_point[0] = point[0]
            return DashedLine(
                x_axis_point, point,
                color=RED,
                stroke_width=2,
            )
        self.y_line = updating_mobject_from_func(get_y_line)
        return self.y_line

    def get_phase_space_dot(self):
        self.ps_dot = ps_dot = Dot(**self.ps_dot_config)
        ps_dot.add_updater(lambda m: m.move_to(self.ps_point))
        return ps_dot

    def get_d_label(self, n, get_d):
        label = VGroup(
            TexMobject("d_{}".format(n), "="),
            DecimalNumber(),
        )
        color = GREEN if n == 1 else RED
        label[0].set_color_by_tex("d_", color)
        label.scale(0.7)
        label.set_stroke(BLACK, 3, background=True)

        def update_value(label):
            lhs, rhs = label
            rhs.set_value(get_d())
            rhs.next_to(
                lhs, RIGHT, SMALL_BUFF,
                aligned_edge=DOWN,
            )
        label.add_updater(update_value)
        return label

    def get_phase_space_d_label(self, n, get_d, line, vect):
        label = self.get_d_label(n, get_d)
        label.add_updater(
            lambda m: m.next_to(line, vect, SMALL_BUFF)
        )
        return label

    def get_phase_space_d1_label(self):
        self.ps_d1_label = self.get_phase_space_d_label(
            1, self.get_d1, self.x_line, UP,
        )
        return self.ps_d1_label

    def get_phase_space_d2_label(self):
        self.ps_d2_label = self.get_phase_space_d_label(
            2, self.get_d2, self.y_line,
            self.ps_d2_label_vect,
        )
        return self.ps_d2_label

    def get_d_brace(self, get_right_point):
        line = Line(LEFT, RIGHT).set_width(6)

        def get_brace():
            right_point = get_right_point()
            left_point = np.array(right_point)
            left_point[0] = self.wall_x
            line.put_start_and_end_on(left_point, right_point)
            return Brace(line, UP, buff=SMALL_BUFF)

        brace = updating_mobject_from_func(get_brace)
        return brace

    def get_d1_brace(self):
        self.d1_brace = self.get_d_brace(
            lambda: self.block1.get_corner(UL)
        )
        return self.d1_brace

    def get_d2_brace(self):
        self.d2_brace = self.get_d_brace(
            lambda: self.block2.get_corner(UR)
        )
        # self.flip_brace_nip()
        return self.d2_brace

    def flip_brace_nip(self, brace):
        nip_index = (len(brace) // 2) - 1
        nip = brace[nip_index:nip_index + 2]
        rect = brace[nip_index - 1]
        center = rect.get_center()
        center[0] = nip.get_center()[0]
        nip.rotate(PI, about_point=center)

    def get_brace_d_label(self, n, get_d, brace, vect, buff):
        label = self.get_d_label(n, get_d)
        label.add_updater(
            lambda m: m.next_to(brace, vect, buff)
        )
        return label

    def get_d1_label(self):
        self.d1_label = self.get_brace_d_label(
            1, self.get_d1, self.d1_brace, UP, SMALL_BUFF,
        )
        return self.d1_label

    def get_d2_label(self):
        self.d2_label = self.get_brace_d_label(
            2, self.get_d2, self.d2_brace, UP, 0
        )
        return self.d2_label

    def get_d1_eq_d2_line(self):
        start = self.ds_to_point(0, 0)
        end = self.ds_to_point(15, 15)
        line = self.d1_eq_d2_line = self.mirror_line_class(start, end)
        line.set_style(**self.mirror_line_style)
        line.set_color(self.d1_eq_e2_line_color)
        return self.d1_eq_d2_line

    def get_d1_eq_d2_label(self):
        label = TexMobject("d1 = d2")
        label.scale(0.75)
        line = self.d1_eq_d2_line
        point = interpolate(
            line.get_start(), line.get_end(),
            0.7,
        )
        label.next_to(point, DR, SMALL_BUFF)
        label.match_color(line)
        label.set_stroke(BLACK, 5, background=True)
        self.d1_eq_d2_label = label
        return label

    def get_d2_eq_w2_line(self):
        w2 = self.block2.width
        start = self.ds_to_point(0, w2)
        end = np.array(start)
        end[0] = FRAME_WIDTH / 2
        self.d2_eq_w2_line = self.mirror_line_class(start, end)
        self.d2_eq_w2_line.set_style(**self.mirror_line_style)
        return self.d2_eq_w2_line

    def get_d2_eq_w2_label(self):
        label = TexMobject("d2 = \\text{block width}")
        label.scale(0.75)
        label.next_to(self.d2_eq_w2_line, UP, SMALL_BUFF)
        label.to_edge(RIGHT, buff=MED_SMALL_BUFF)
        self.d2_eq_w2_label = label
        return label

    def get_time_tracker(self, time=0):
        time_tracker = self.time_tracker = ValueTracker(time)
        time_tracker.add_updater(
            lambda m, dt: m.increment_value(dt)
        )
        return time_tracker


class IntroducePositionPhaseSpace(PositionPhaseSpaceScene):
    CONFIG = {
        "rescale_coordinates": False,
        "block1_config": {
            "velocity": 1.5,
        },
        "slide_wait_time": 30,
    }

    def setup(self):
        super().setup()
        self.add(
            self.floor,
            self.wall,
            self.blocks,
            self.axes,
        )

    def construct(self):
        self.show_coordinates()
        self.show_xy_line()
        self.let_process_play_out()
        self.show_w2_line()

    def show_coordinates(self):
        ps_point = self.ps_point
        axes = self.axes

        self.play(Write(axes.added_lines))
        self.play(FadeInFromLarge(self.ps_dot, scale_factor=10))
        self.play(
            ShowCreation(self.x_line),
            GrowFromPoint(
                self.d1_brace,
                self.d1_brace.get_left(),
            ),
            Indicate(axes.labels[0]),
        )
        self.play(
            FadeInFromDown(self.ps_d1_label),
            FadeInFromDown(self.d1_label),
        )
        self.play(ps_point.shift, 0.5 * LEFT)
        self.play(ps_point.shift, 0.5 * RIGHT)
        self.wait()
        self.play(
            ShowCreation(self.y_line),
            GrowFromPoint(
                self.d2_brace,
                self.d2_brace.get_left(),
            ),
            Indicate(axes.labels[1]),
        )
        self.play(
            FadeInFromDown(self.ps_d2_label),
            FadeInFromDown(self.d2_label),
        )
        self.play(ps_point.shift, 0.5 * UP)
        self.play(ps_point.shift, 0.5 * DOWN)
        self.wait()
        self.play(Rotating(
            ps_point,
            about_point=ps_point.get_location() + 0.5 * RIGHT,
            run_time=3,
            rate_func=smooth,
        ))
        self.wait()

    def show_xy_line(self):
        ps_point = self.ps_point
        ps_point.save_state()
        d1, d2 = self.point_to_ds(ps_point.get_location())

        xy_line = self.d1_eq_d2_line
        xy_label = self.d1_eq_d2_label

        self.play(
            ShowCreation(xy_line),
            Write(xy_label),
        )
        self.play(
            ps_point.move_to, self.ds_to_point(d2, d2),
            run_time=3
        )
        self.wait()
        for d in [3, 7]:
            self.play(
                ps_point.move_to, self.ds_to_point(d, d),
                run_time=2
            )
            self.wait()
        self.play(ps_point.restore)
        self.wait()

    def let_process_play_out(self):
        self.begin_sliding()
        sliding_trajectory = self.get_continually_building_trajectory()
        self.add(sliding_trajectory, self.ps_dot)
        self.wait(self.slide_wait_time)
        self.end_sliding()

    def show_w2_line(self):
        line = self.d2_eq_w2_line
        label = self.d2_eq_w2_label

        self.play(ShowCreation(line))
        self.play(FadeInFromDown(label))
        self.wait()


class SpecialShowPassingFlash(ShowPassingFlash):
    CONFIG = {
        "max_time_width": 0.1,
    }

    def get_bounds(self, alpha):
        tw = self.time_width
        max_tw = self.max_time_width
        upper = interpolate(0, 1 + max_tw, alpha)
        lower = upper - tw
        upper = min(upper, 1)
        lower = max(lower, 0)
        return (lower, upper)


class EqualMassCase(PositionPhaseSpaceScene):
    CONFIG = {
        "block1_config": {
            "mass": 1,
            "width": 1,
            "velocity": 1.5,
        },
        "rescale_coordinates": False,
    }

    def setup(self):
        super().setup()
        self.add(
            self.floor,
            self.wall,
            self.blocks,
            self.axes,
            self.d1_eq_d2_line,
            self.d1_eq_d2_label,
            self.d2_eq_w2_line,
            self.d2_eq_w2_label,
            self.ps_dot,
            self.x_line,
            self.y_line,
            self.ps_d1_label,
            self.ps_d2_label,
        )

    def construct(self):
        self.show_same_mass()
        self.show_first_point()
        self.up_to_first_collision()
        self.up_to_second_collision()
        self.up_to_third_collision()

        self.fade_distance_indicators()
        self.show_beam_bouncing()

    def show_same_mass(self):
        blocks = self.blocks
        self.play(LaggedStart(
            Indicate, blocks,
            lag_ratio=0.8,
            run_time=1,
        ))

    def show_first_point(self):
        ps_dot = self.ps_dot
        ps_point = self.ps_point
        d1, d2 = self.get_ds()

        self.play(FocusOn(ps_dot))
        self.play(ShowCreationThenFadeOut(
            Circle(color=RED).replace(ps_dot).scale(2),
            run_time=1
        ))
        self.wait()
        self.play(
            ps_point.move_to, self.ds_to_point(d1 - 1, d2),
            rate_func=wiggle,
            run_time=3,
        )
        # self.play(ps_point.move_to, self.ds_to_point(d1, d2))
        self.wait()

    def up_to_first_collision(self):
        ps_point = self.ps_point
        d1, d2 = self.get_ds()
        block1 = self.block1
        block2 = self.block2
        xy_line = self.d1_eq_d2_line
        xy_line_label = self.d1_eq_d2_label

        block_arrow = Vector(LEFT, color=RED)
        block_arrow.block = block1
        block_arrow.add_updater(
            lambda m: m.shift(
                m.block.get_center() - m.get_start()
            )
        )
        ps_arrow = Vector(LEFT, color=RED)
        ps_arrow.next_to(ps_point, DL, buff=SMALL_BUFF)

        block_labels = VGroup(block1.label, block2.label)
        block_label_copies = block_labels.copy()

        def update_bl_copies(bl_copies):
            for bc, b in zip(bl_copies, block_labels):
                bc.move_to(b)
        block_label_copies.add_updater(update_bl_copies)

        trajectory = self.get_continually_building_trajectory()

        self.add(block_arrow, ps_arrow, block_label_copies)
        self.play(
            GrowArrow(block_arrow),
            GrowArrow(ps_arrow),
        )
        self.add(trajectory)
        self.play(self.get_ps_point_change_anim(d2, d2))
        block_arrow.block = block2
        ps_arrow.rotate(90 * DEGREES)
        ps_arrow.next_to(ps_point, DR, SMALL_BUFF)
        self.add_sound(self.clack_sound)
        self.play(
            Flash(ps_point),
            Flash(block1.get_left()),
            self.get_ps_point_change_anim(d2, d2 - 1)
        )
        self.play(
            ShowPassingFlash(
                xy_line.copy().set_stroke(YELLOW, 3)
            ),
            Indicate(xy_line_label),
        )

        trajectory.suspend_updating()
        self.wait()

        self.ps_arrow = ps_arrow
        self.block_arrow = block_arrow

    def up_to_second_collision(self):
        trajectory = self.trajectory
        ps_point = self.ps_point
        ps_arrow = self.ps_arrow
        block_arrow = self.block_arrow

        d1, d2 = self.get_ds()
        w2 = self.block2.get_width()

        trajectory.resume_updating()
        self.play(self.get_ps_point_change_anim(d1, w2))
        block_arrow.rotate(PI)
        ps_arrow.rotate(PI)
        ps_arrow.next_to(ps_point, UR, SMALL_BUFF)
        self.add_sound(self.clack_sound)
        self.play(
            Flash(self.block2.get_left()),
            Flash(ps_point),
            self.get_ps_point_change_anim(d1, w2 + 1)
        )

        trajectory.suspend_updating()
        self.wait()

    def up_to_third_collision(self):
        trajectory = self.trajectory
        ps_point = self.ps_point
        ps_arrow = self.ps_arrow
        block_arrow = self.block_arrow
        d1, d2 = self.get_ds()

        trajectory.resume_updating()
        self.play(self.get_ps_point_change_anim(d1, d1))
        block_arrow.block = self.block1
        ps_arrow.rotate(-90 * DEGREES)
        ps_arrow.next_to(ps_point, DR, SMALL_BUFF)
        self.add_sound(self.clack_sound)
        self.play(
            Flash(self.block2.get_left()),
            Flash(ps_point.get_location()),
            self.get_ps_point_change_anim(d1 + 10, d1)
        )
        trajectory.suspend_updating()

    def fade_distance_indicators(self):
        trajectory = self.trajectory
        self.play(
            trajectory.set_stroke, {"width": 1},
            *map(FadeOut, [
                self.ps_arrow,
                self.block_arrow,
                self.x_line,
                self.y_line,
                self.ps_d1_label,
                self.ps_d2_label,
            ])
        )
        trajectory.clear_updaters()

    def show_beam_bouncing(self):
        d1, d2 = self.get_ds()
        d1 = int(d1)
        d2 = int(d2)
        w2 = self.block2.get_width()
        ps_point = self.ps_point

        points = []
        while d1 > d2:
            points.append(self.ds_to_point(d1, d2))
            d1 -= 1
        while d2 >= int(w2):
            points.append(self.ds_to_point(d1, d2))
            d2 -= 1
        points += list(reversed(points))[1:]
        trajectory = VMobject()
        trajectory.set_points_as_corners(points)
        flashes = [
            SpecialShowPassingFlash(
                trajectory.copy().set_stroke(YELLOW, width=6 - n),
                time_width=(0.01 * n),
                max_time_width=0.05,
                remover=True
            )
            for n in np.arange(0, 6, 0.25)
        ]
        flash_mob = flashes[0].mobject  # Lol

        def update_ps_point_from_flas_mob(ps_point):
            if len(flash_mob.points) > 0:
                ps_point.move_to(flash_mob.points[-1])
            else:
                ps_point.move_to(trajectory.points[0])

        # Mirror words
        xy_line = self.d1_eq_d2_line
        w2_line = self.d2_eq_w2_line
        lines = VGroup(xy_line, w2_line)
        for line in lines:
            word = TextMobject("Mirror")
            word.next_to(ORIGIN, UP, SMALL_BUFF)
            word.rotate(line.get_angle(), about_point=ORIGIN)
            word.shift(line.get_center())
            line.word = word

        for line in lines:
            line.set_stroke(LIGHT_GREY)
            line.set_sheen(1, LEFT)
            self.play(
                Write(line.word),
                line.set_sheen, 1, RIGHT,
                line.set_stroke, {"width": 2},
                run_time=1,
            )

        # TODO, clacks?
        for x in range(3):
            self.play(
                UpdateFromFunc(
                    ps_point,
                    update_ps_point_from_flas_mob,
                ),
                *flashes,
                run_time=3,
                rate_func=None,
            )
            self.wait()


class FailedAngleRelation(PositionPhaseSpaceScene):
    CONFIG = {
        "block1_config": {
            "distance": 10,
            "velocity": -1.5,
        },
        "block2_config": {
            "distance": 5,
        },
        "rescale_coordinates": False,
        "trajectory_style": {
            "stroke_width": 2,
        }
    }

    def setup(self):
        super().setup()
        self.add(
            self.floor,
            self.wall,
            self.blocks,
            self.axes,
            self.ps_dot,
            self.x_line,
            self.y_line,
            self.d1_eq_d2_line,
            self.d1_eq_d2_label,
            self.d2_eq_w2_line,
            self.d2_eq_w2_label,
        )

    def construct(self):
        self.show_first_collision()
        self.show_angles()

    def show_first_collision(self):
        self.slide_until(lambda: self.get_ds()[1] < 2)

    def show_angles(self):
        trajectory = self.trajectory
        arcs = self.get_arcs(trajectory)
        equation = self.get_word_equation()
        equation.next_to(
            trajectory.points[0], UR, MED_SMALL_BUFF,
            index_of_submobject_to_align=0,
        )

        for arc in arcs:
            line = Line(ORIGIN, RIGHT)
            line.set_stroke(WHITE, 2)
            line.rotate(arc.start_angle)
            line.shift(arc.arc_center - line.get_start())
            arc.line = line

        arc1, arc2 = arcs
        arc1.arrow = Arrow(
            equation[0].get_left(), arc1.get_right(),
            buff=SMALL_BUFF,
            color=WHITE,
            path_arc=0,
        )
        arc2.arrow = Arrow(
            equation[2].get_corner(DL),
            arc2.get_left(),
            use_rectangular_stem=False,
            path_arc=-120 * DEGREES,
            buff=SMALL_BUFF,
        )
        arc2.arrow.pointwise_become_partial(arc.arrow, 0, 0.95)

        arc1.word = equation[0]
        arc2.word = equation[1:]

        for arc in arcs:
            self.play(
                FadeInFrom(arc.word, LEFT),
                GrowArrow(arc.arrow, path_arc=arc.arrow.path_arc),
            )
            self.play(
                ShowCreation(arc),
                arc.line.rotate, arc.angle,
                {"about_point": arc.line.get_start()},
                UpdateFromAlphaFunc(
                    arc.line,
                    lambda m, a: m.set_stroke(
                        opacity=(there_and_back(a)**0.5)
                    )
                ),
            )

    #
    def get_arcs(self, trajectory):
        p0, p1, p2 = trajectory.get_anchors()[1:4]
        arc_config = {
            "stroke_color": WHITE,
            "stroke_width": 2,
            "radius": 0.5,
            "arc_center": p1,
        }
        arc1 = Arc(
            start_angle=0,
            angle=45 * DEGREES,
            **arc_config
        )
        a2_start = angle_of_vector(DL)
        a2_angle = angle_between_vectors((p2 - p1), DL)
        arc2 = Arc(
            start_angle=a2_start,
            angle=a2_angle,
            **arc_config
        )
        return VGroup(arc1, arc2)

    def get_word_equation(self):
        result = VGroup(
            TextMobject("Angle of incidence"),
            TexMobject("\\ne").rotate(90 * DEGREES),
            TextMobject("Angle of refraction")
        )
        result.arrange_submobjects(DOWN)
        result.set_stroke(BLACK, 5, background=True)
        return result


class UnscaledPositionPhaseSpaceMass10(FailedAngleRelation):
    CONFIG = {
        "block1_config": {
            "mass": 10
        },
        "wait_time": 25,
    }

    def construct(self):
        self.slide(self.wait_time)


class UnscaledPositionPhaseSpaceMass100(UnscaledPositionPhaseSpaceMass10):
    CONFIG = {
        "block1_config": {
            "mass": 100
        }
    }


class RescaleCoordinates(PositionPhaseSpaceScene, MovingCameraScene):
    CONFIG = {
        "rescale_coordinates": False,
        "ps_d2_label_vect": LEFT,
        "axes_center": 6 * LEFT + 0.65 * DOWN,
        "block1_config": {"distance": 7},
        "wait_time": 30,
    }

    def setup(self):
        PositionPhaseSpaceScene.setup(self)
        MovingCameraScene.setup(self)
        self.add(
            self.floor,
            self.wall,
            self.blocks,
            self.axes,
            self.d1_eq_d2_line,
            self.d1_eq_d2_label,
            self.d2_eq_w2_line,
            self.ps_dot,
            self.x_line,
            self.y_line,
            self.ps_d1_label,
            self.ps_d2_label,
            self.d1_brace,
            self.d2_brace,
            self.d1_label,
            self.d2_label,
        )

    def construct(self):
        self.show_rescaling()
        self.comment_on_ugliness()
        self.put_into_frame()

    def show_rescaling(self):
        axes = self.axes
        blocks = self.blocks
        to_stretch = VGroup(
            axes.added_lines,
            self.d1_eq_d2_line,
            self.ps_point,
        )
        m1 = self.block1.mass
        new_axes_labels = self.get_axes_labels(axes, with_sqrts=True)

        # Show label
        def show_label(index, block, vect):
            self.play(
                ShowCreationThenFadeAround(axes.labels[index])
            )
            self.play(
                Transform(
                    axes.labels[index],
                    new_axes_labels[index][:2],
                ),
                GrowFromCenter(new_axes_labels[index][2])
            )
            group = VGroup(
                new_axes_labels[index][2][-2:].copy(),
                TexMobject("="),
                block.label.copy(),
            )
            group.generate_target()
            group.target.arrange_submobjects(RIGHT, buff=SMALL_BUFF)
            group.target.next_to(block, vect)
            group[1].scale(0)
            group[1].move_to(group.target[1])
            group.target[2].set_fill(WHITE)
            group.target[2].set_stroke(width=0, background=True)
            self.play(MoveToTarget(
                group,
                rate_func=there_and_back_with_pause,
                run_time=3
            ))
            self.remove(group)
            self.wait()

        show_label(0, self.block1, RIGHT)

        # The stretch
        blocks.suspend_updating()
        self.play(
            ApplyMethod(
                to_stretch.stretch, np.sqrt(m1), 0,
                {"about_point": axes.coords_to_point(0, 0)},
            ),
            self.d1_eq_d2_label.shift, 6 * RIGHT,
            run_time=2,
        )
        self.rescale_coordinates = True
        blocks.resume_updating()
        self.wait()

        # Show wiggle
        d1, d2 = self.get_ds()
        for new_d1 in [d1 - 2, d1]:
            self.play(self.get_ps_point_change_anim(
                new_d1, d2,
                rate_func=smooth,
                run_time=2,
            ))
        self.wait()

        # Change y-coord
        show_label(1, self.block2, LEFT)

        axes.remove(axes.labels)
        self.remove(axes.labels)
        axes.labels = new_axes_labels
        axes.add(axes.labels)
        self.add(axes)

    def comment_on_ugliness(self):
        axes = self.axes

        randy = Randolph(height=1.7)
        randy.flip()
        randy.next_to(self.d2_eq_w2_line, UP, buff=0)
        randy.to_edge(RIGHT)
        randy.change("sassy")
        randy.save_state()
        randy.fade(1)
        randy.change("plain")

        self.play(Restore(randy))
        self.play(
            PiCreatureSays(
                randy, "Hideous!",
                bubble_kwargs={"height": 1.5, "width": 2},
                target_mode="angry",
                look_at_arg=axes.labels[0]
            )
        )
        self.play(randy.look_at, axes.labels[1])
        self.play(Blink(randy))
        self.play(
            RemovePiCreatureBubble(
                randy, target_mode="confused"
            )
        )
        self.play(Blink(randy))
        self.play(randy.look_at, axes.labels[0])
        self.wait()
        self.play(FadeOut(randy))

    def put_into_frame(self):
        rect = ScreenRectangle(height=FRAME_HEIGHT + 10)
        inner_rect = ScreenRectangle(height=FRAME_HEIGHT)
        rect.add_subpath(inner_rect.points[::-1])
        rect.set_fill(DARK_GREY, opacity=1)
        frame = self.camera_frame

        self.begin_sliding()
        self.add(rect)
        self.play(
            frame.scale, 1.5,
            {"about_point": frame.get_bottom() + UP},
            run_time=2,
        )
        self.wait(self.wait_time)
        self.end_sliding()

    #
    def get_ds(self):
        if self.rescale_coordinates:
            return super().get_ds()
        return (
            self.block1_config["distance"],
            self.block2_config["distance"],
        )


class RescaleCoordinatesMass16(RescaleCoordinates):
    CONFIG = {
        "block1_config": {
            "mass": 16,
            "distance": 10,
        },
        "rescale_coordinates": True,
        "wait_time": 20,
    }

    def construct(self):
        self.put_into_frame()


class RescaleCoordinatesMass64(RescaleCoordinatesMass16):
    CONFIG = {
        "block1_config": {
            "mass": 64,
            "distance": 6,
        },
        "block2_config": {"distance": 3},
    }

    def construct(self):
        self.put_into_frame()


class NewSceneName(Scene):
    def construct(self):
        pass