from manim import *

shift_pos = LEFT * config.frame_width / 4
graph_pos = RIGHT * 3 + UP * 2
graph_pos_left = graph_pos + LEFT * 2
text_pos_1 = (1.5, 3, 0)
text_pos_2 = (3.0, 3, 0)
start_size = config.frame_width / 2.5
start_rad = 0.3
l1_color = PURPLE
l2_color = LIGHT_PINK
l3_color = RED
l4_color = YELLOW
l5_color = PINK


class DefaultTemplate(ThreeDScene):
    def construct(self):
        plus = Text("+")
        plus.move_to(text_pos_1).shift(((text_pos_2[0] - text_pos_1[0]) / 2, 0, 0))

        equal = Text("=")
        equal.move_to(text_pos_1).shift(((text_pos_2[0] - text_pos_1[0]) * 1.5, 0, 0))

        top = Square(color=WHITE, side_length=start_size)
        top.shift(shift_pos)

        l1 = Square(
            color=l1_color,
            side_length=start_size,
            grid_xstep=start_size / 2,
            grid_ystep=start_size / 2,
        )
        l1.set_z_index(-1)
        l1.shift(shift_pos)

        def create_node(color, kg, destination, radius=start_rad):
            node = Circle(color=color, radius=radius)
            node.move_to(destination)
            t = Text(str(kg), font_size=(25 * 1 / start_rad) * radius)
            t.move_to(node)
            node = VGroup(node, t)
            return node

        def animate_dot_2_node(dot, node):
            new_node = dot.copy()
            self.play(Transform(new_node, node))
            return new_node

        def add_2_nodes(n1, n2, s, write=True, scale=1, color=GREEN):
            self.play(n1.animate.move_to(text_pos_1))
            if write:
                self.play(Write(plus))
            self.play(n2.animate.move_to(text_pos_2))
            if write:
                self.play(Write(equal))

            result = create_node(color, s, equal)
            result.shift(((text_pos_2[0] - text_pos_1[0]) / 2, 0, 0))
            result.scale(scale)
            self.play(Create(result))
            return result

        def get_sub(root, rad=0.75, start=-2):
            slots = VGroup()
            for i in range(start, 2):
                if start != -2 and i == start:
                    sub = create_node(GRAY_D, "...", root)
                else:
                    sub = create_node(GRAY_D, i + 2, root)
                sub.scale(rad)
                sub.shift(((i + 0.5) * rad, -1, 0))
                slots.add(sub)
            return slots

        def get_square_text(div=1, offset_x=0, offset_y=0):
            slots = VGroup()
            i = 0
            for y in [1, -1]:
                for x in [-1, 1]:
                    text = Text(text=str(i), font_size=int(50 / div))
                    text.move_to(l1)
                    text.shift(
                        (
                            x * start_size / 4 / div + offset_x,
                            y * start_size / 4 / div + offset_y,
                            0,
                        )
                    )
                    slots.add(text)
                    i += 1
            return slots

        def hide_text():
            self.play(Unwrite(plus), Unwrite(equal))

        self.play(Create(top))

        # self.wait(1)

        p1 = self.add_particle((1, 2, 0), BLUE_C)
        p1_node = create_node(BLUE, 5, graph_pos)
        # self.wait(10)
        p1_node = animate_dot_2_node(p1, p1_node)
        self.play(p1.animate.set_color(GRAY))

        p2 = self.add_particle((-2, 1, 0), BLUE_C, 0.5, False)

        p2_node = create_node(BLUE, 1, graph_pos_left)
        p2_node = animate_dot_2_node(p2, p2_node)

        root_1 = add_2_nodes(p1_node, p2_node, 6, color=l1_color)

        self.play(root_1.animate.move_to(graph_pos))
        sub_nodes_r1 = get_sub(root_1)
        sub_text_r1 = get_square_text()

        self.play(Create(l1))

        for i, n in enumerate(sub_nodes_r1):
            arrow = Arrow(start=root_1.get_center(), end=n.get_center(), color=BLUE)
            new = root_1.copy()
            self.play(Write(sub_text_r1[i]), Transform(new, VGroup(n, arrow)))
            sub_nodes_r1[i] = new

        for i, n in enumerate(sub_text_r1):
            self.play(
                n.animate.set_color(GRAY_D),
                l1.animate.set_stroke(opacity=0.5),
                run_time=0.2,
            )

        self.play(p2_node.animate.scale(0.75), p1_node.animate.scale(0.75))
        self.play(
            p1_node.animate.move_to(sub_nodes_r1[1][0]), FadeOut(sub_nodes_r1[1][0])
        )
        self.play(
            p2_node.animate.move_to(sub_nodes_r1[0][0]), FadeOut(sub_nodes_r1[0][0])
        )
        self.play(p2.animate.set_color(GRAY))

        p3 = self.add_particle((1, 1, 0), BLUE_C, 0, False)
        p3_node = create_node(BLUE, 8, graph_pos_left)
        p3_node = animate_dot_2_node(p3, p3_node)
        root_1_1 = add_2_nodes(p3_node, root_1, 14, False, color=l1_color)
        self.play(root_1_1.animate.move_to(graph_pos))
        self.play(FadeOut(root_1))
        root_1 = root_1_1
        root_2 = add_2_nodes(p3_node, p1_node, 13, False, 0.75, l2_color)

        self.play(root_2.animate.move_to(sub_nodes_r1[1][0]))

        l2 = Square(
            color=l2_color,
            side_length=start_size / 2,
            grid_xstep=start_size / 4,
            grid_ystep=start_size / 4,
        )
        l2.shift((start_size / 4, start_size / 4, 0))
        l2.set_z_index(-2)
        l2.shift(shift_pos)
        self.play(Create(l2))

        sub_nodes_r2 = get_sub(root_2, 0.75 * 0.75)
        sub_text_r2 = get_square_text(2, start_size / 4, start_size / 4)
        for i, n in enumerate(sub_nodes_r2):
            arrow = Arrow(start=root_2.get_center(), end=n.get_center(), color=BLUE)
            new = root_2.copy()
            self.play(Write(sub_text_r2[i]), Transform(new, VGroup(n, arrow)))
            sub_nodes_r2[i] = new

        for i, n in enumerate(sub_text_r2):
            self.play(
                n.animate.set_color(GRAY_D),
                l2.animate.set_stroke(opacity=0.5),
                run_time=0.2,
            )

        self.play(p1_node.animate.scale(0.75), p3_node.animate.scale(0.75 * 0.75))
        self.play(
            p1_node.animate.move_to(sub_nodes_r2[0][0]), FadeOut(sub_nodes_r2[0][0])
        )
        self.play(
            p3_node.animate.move_to(sub_nodes_r2[2][0]), FadeOut(sub_nodes_r2[2][0])
        )
        self.play(p3.animate.set_color(GRAY))

        p4 = self.add_particle((2, -1.8, 0), BLUE_C, 0, False)
        p4_node = create_node(BLUE, 42, graph_pos_left)
        p4_node = animate_dot_2_node(p4, p4_node)
        root_1_1 = add_2_nodes(p4_node, root_1, 56, False, color=l1_color)
        self.play(root_1_1.animate.move_to(graph_pos))
        self.play(FadeOut(root_1))
        root_1 = root_1_1
        self.play(p4_node.animate.scale(0.75), p4.animate.set_color(GRAY))
        self.play(
            p4_node.animate.move_to(sub_nodes_r1[3][0]), FadeOut(sub_nodes_r1[3][0])
        )

        # P5
        p5 = self.add_particle((2.5, -2.3, 0), BLUE_C, 0, False)
        p5_node = create_node(BLUE, 2, graph_pos_left)
        p5_node = animate_dot_2_node(p5, p5_node)
        root_1_1 = add_2_nodes(p5_node, root_1, 58, False, color=l1_color)
        self.play(root_1_1.animate.move_to(graph_pos))
        self.play(FadeOut(root_1))
        root_1 = root_1_1

        root_2_4 = add_2_nodes(p5_node, p4_node, 44, False, 0.75, color=l2_color)
        self.play(root_2_4.animate.move_to(sub_nodes_r1[3][0]))

        l3 = Square(
            color=l2_color,
            side_length=start_size / 2,
            grid_xstep=start_size / 4,
            grid_ystep=start_size / 4,
        )
        l3.shift((start_size / 4, -start_size / 4, 0))
        l3.set_z_index(-2)
        l3.shift(shift_pos)
        self.play(Create(l3))
        sub_nodes_r24 = get_sub(root_2_4, 0.75, 0)
        sub_text_r24 = get_square_text(2, start_size / 4, -start_size / 4)

        for n in sub_text_r24:
            self.play(Write(n))

        for i, n in enumerate(sub_nodes_r24):
            arrow = Arrow(start=root_2_4.get_center(), end=n.get_center(), color=BLUE)
            new = root_2_4.copy()
            self.play(Transform(new, VGroup(n, arrow)))
            sub_nodes_r24[i] = new

        self.play(
            p4_node.animate.move_to(sub_nodes_r24[1][0]), FadeOut(sub_nodes_r24[1][0])
        )
        self.play(p4_node.animate.scale(0.75), l3.animate.set_stroke(opacity=0.5))

        # l4
        l4 = Square(
            color=l3_color,
            side_length=start_size / 4,
            grid_xstep=start_size / 8,
            grid_ystep=start_size / 8,
        )
        l4.shift(((start_size / 4) * 1.5, (-start_size / 4) * 1.5, 0))
        l4.set_z_index(-3)
        l4.shift(shift_pos)
        self.play(Create(l4))

        root_3 = add_2_nodes(p5_node, p4_node, 44, False, 0.75**2, l3_color)
        self.play(root_3.animate.move_to(sub_nodes_r24[1][0]))

        sub_nodes_r3 = get_sub(root_3, 0.75**3)

        for i, n in enumerate(sub_nodes_r3):
            arrow = Arrow(start=root_3.get_center(), end=n.get_center(), color=BLUE)
            new = root_3.copy()
            self.play(Transform(new, VGroup(n, arrow)))
            sub_nodes_r3[i] = new

        self.play(
            p5_node.animate.scale(0.75**3), p4_node.animate.scale(0.75), run_time=0.3
        )
        self.play(
            p4_node.animate.move_to(sub_nodes_r3[0][0]), FadeOut(sub_nodes_r3[0][0])
        )
        self.play(
            p5_node.animate.move_to(sub_nodes_r3[3][0]), FadeOut(sub_nodes_r3[3][0])
        )

        hide_text()

        sub_nodes_r1[0][0] = sub_nodes_r1[0][1]
        sub_nodes_r1[1][0] = sub_nodes_r1[1][1]
        sub_nodes_r1[3][0] = sub_nodes_r1[3][1]

        sub_nodes_r2[0][0] = sub_nodes_r2[0][1]
        sub_nodes_r2[2][0] = sub_nodes_r2[2][1]

        sub_nodes_r24[1][0] = sub_nodes_r24[1][1]

        sub_nodes_r3[0][0] = sub_nodes_r3[0][1]
        sub_nodes_r3[3][0] = sub_nodes_r3[3][1]

        graph = VGroup(
            p1_node,
            p2_node,
            p3_node,
            p4_node,
            p5_node,
            root_1,
            root_2,
            root_2_4,
            root_3,
            sub_nodes_r1,
            sub_nodes_r2,
            sub_nodes_r24,
            sub_nodes_r3,
        )
        graph.save_state()
        self.play(
            FadeOut(graph),
            p5.animate.set_color(GRAY),
            FadeOut(sub_text_r1),
            FadeOut(sub_text_r2),
            FadeOut(sub_text_r24),
        )

        v = self.add_particles_and_squares()

        self.play(
            l1.animate.set_stroke(opacity=1),
            l3.animate.set_stroke(opacity=1),
            l2.animate.set_stroke(opacity=1),
            run_time=0.3,
        )

        diagram = VGroup(v, top, p1, p2, p3, p4, p5, l1, l2, l3, l4)

        self.begin_ambient_camera_rotation(90 * DEGREES / 3, about="phi")

        self.play(diagram.animate.move_to((0, 0, 0)), run_time=2.5)
        diagram.save_state()

        self.stop_ambient_camera_rotation(about="phi")

        self.play(
            top.animate.shift((0, 0, 2)),
            l4.animate.shift((0, 0, -1)),
            p4.animate.shift((0, 0, -2)),
            p5.animate.shift((0, 0, -1)),
            p2.animate.shift((0, 0, -2)),
            p1.animate.shift((0, 0, -3)),
            *[x.animate.shift((0, 0, 0)) for x in v[0]],
            *[x.animate.shift((0, 0, -1)) for x in v[1]],
            *[x.animate.shift((0, 0, -2)) for x in v[2]],
            *[x.animate.shift((0, 0, -3)) for x in v[3]],
            run_time=2,
        )

        self.play(
            Rotate(
                diagram, angle=360 * DEGREES, rate_func=rate_functions.ease_in_out_cubic
            ),
            run_time=5,
        )

        self.begin_ambient_camera_rotation(-90 * DEGREES / 3, about="phi")
        self.play(Restore(diagram), run_time=2.5)
        self.stop_ambient_camera_rotation(about="phi")
        self.wait(1)
        diagram = diagram[1:]
        self.play(FadeOut(v), run_time=0.5)
        self.play(Restore(graph), diagram.animate.move_to(shift_pos))
        self.wait(3)

        return

    def create_square(self, color, size, shift_pos, z_index, shift_values):
        square = Square(
            color=color, side_length=size, grid_xstep=size / 2, grid_ystep=size / 2
        )
        square.shift(shift_values)
        square.set_z_index(z_index)
        square.shift(shift_pos)
        self.play(Create(square))
        return square

    def add_particles_and_squares(self):
        t = 0.2
        p6 = self.add_particle((-2.7, -2.7, 0), GRAY, t, False)
        p7 = self.add_particle((-2.1, -1.1, 0), GRAY, t, False)
        l7 = self.create_square(
            l2_color,
            start_size / 2,
            shift_pos,
            -2,
            (-start_size / 4, -start_size / 4, 0),
        )

        p8 = self.add_particle((1.1, 2.7, 0), GRAY, t, False)
        l8 = self.create_square(
            l3_color,
            start_size / 4,
            shift_pos,
            -3,
            ((start_size / 4) * 0.5, (start_size / 4) * 1.5, 0),
        )

        p9 = self.add_particle((1.67, 1.24, 0), GRAY, t, False)
        p11 = self.add_particle((-0.92, 2.12, 0), GRAY, t, False)
        l10 = self.create_square(
            l2_color,
            start_size / 2,
            shift_pos,
            -2,
            (-start_size / 4, start_size / 4, 0),
        )

        p12 = self.add_particle((1.13, -1.45, 0), GRAY, t, False)
        p13 = self.add_particle((-2.48, 0.58, 0), GRAY, t, False)
        l12 = self.create_square(
            l3_color,
            start_size / 4,
            shift_pos,
            -3,
            (-(start_size / 4) * 1.5, (start_size / 4) * 0.5, 0),
        )

        p14 = self.add_particle((0.62, -2.14, 0), GRAY, t, False)
        p15 = self.add_particle((2.35, -1.78, 0), GRAY, t, False)
        p16 = self.add_particle((-1.61, 0.91, 0), GRAY, t, False)
        l16 = self.create_square(
            l4_color,
            start_size / 8,
            shift_pos,
            -4,
            (-(start_size / 4) * 1.25, (start_size / 4) * 0.75, 0),
        )

        p17 = self.add_particle((-2.35, -0.98, 0), GRAY, t, False)
        l17 = self.create_square(
            l3_color,
            start_size / 4,
            shift_pos,
            -3,
            (-(start_size / 4) * 1.5, -(start_size / 4) * 0.5, 0),
        )

        p18 = self.add_particle((0.84, 1.95, 0), GRAY, t, False)
        l18 = self.create_square(
            l4_color,
            start_size / 8,
            shift_pos,
            -4,
            ((start_size / 4) * 0.75, (start_size / 4) * 1.25, 0),
        )

        l19 = self.create_square(
            l5_color,
            start_size / 16,
            shift_pos,
            -5,
            ((start_size / 4) * 0.625, (start_size / 4) * 1.375, 0),
        )

        p20 = self.add_particle((0.5, -0.5, 0), GRAY, t, False)
        l20 = self.create_square(
            l3_color,
            start_size / 4,
            shift_pos,
            -3,
            ((start_size / 4) * 0.5, -(start_size / 4) * 0.5, 0),
        )
        p21 = self.add_particle((-1.8, 1.4, 0), GRAY, t, False)
        p22 = self.add_particle((1.5, -2.0, 0), GRAY, t, False)
        l22 = self.create_square(
            l4_color,
            start_size / 8,
            shift_pos,
            -4,
            ((start_size / 4) * 1.25, -(start_size / 4) * 1.25, 0),
        )
        p23 = self.add_particle((-0.5, -1.6, 0), GRAY, t, False)

        p24 = self.add_particle((-2.0, -2.2, 0), GRAY, t, False)
        l24 = self.create_square(
            l3_color,
            start_size / 4,
            shift_pos,
            -3,
            (-(start_size / 4) * 1.5, -(start_size / 4) * 1.5, 0),
        )

        l2_g = VGroup(l7, p9, l10, p11, p14, p23)
        l3_g = VGroup(p13, p7, p8, l8, p15, l12, p17, l17, p20, p12, l20, p24, p6, l24)
        l4_g = VGroup(p16, l18, l16, p22, l22, p21)
        l5_g = VGroup(p18, l19)
        return VGroup(l2_g, l3_g, l4_g, l5_g)

    def add_particle(self, point, color, runtime=1, write=True):
        dot = Dot(color=color)
        dot.move_to((shift_pos))
        lable = Text(f"{point[:2]}", font_size=25, color=color)
        lable.shift(tuple(x + 0.35 for x in point))
        if runtime > 0:
            if runtime <= 0.3:
                dot.shift(point)
                self.play(Create(dot), run_time=runtime)
            else:
                self.play(Create(dot), run_time=runtime)
                self.play(dot.animate.shift(point), run_time=runtime)
        else:
            dot.shift(point)
            self.add(dot)
        if write:
            lable.move_to(dot)
            lable.shift((0.5, 0.2, 0))
            self.play(Write(lable))
            # self.wait(1)
            self.play(FadeOut(lable))
        return dot