from manim import *


class PythagoreanProof(Scene):
    def construct(self):
        # 场景1：问题引入
        # 创建直角三角形
        tri = Polygon([-2, -1, 0], [1, -1, 0], [-2, 2, 0], color=WHITE)
        labels = VGroup(
            Tex("a", color=RED).next_to(tri.get_vertices()[1], DOWN),
            Tex("b", color=BLUE).next_to(tri.get_vertices()[2], LEFT),
            Tex("c", color=GREEN).next_to(tri.get_midpoint(), UR, buff=0.3)
        )

        self.play(Create(tri), run_time=2)
        self.play(Write(labels), run_time=1.5)

        # 显示平方表达式
        equation = VGroup(
            Tex("a^2 + b^2 = ?", color=RED_B).shift(UP * 2),
            Tex("c^2 = ?", color=GREEN_B).shift(DOWN * 2)
        )
        arrows = VGroup(
            Arrow(equation[0].get_bottom(), labels[0].get_center() + DOWN * 0.5, color=RED),
            Arrow(equation[0].get_bottom(), labels[1].get_center() + LEFT * 0.5, color=BLUE),
            Arrow(equation[1].get_top(), labels[2].get_center() + UR * 0.3, color=GREEN)
        )

        self.play(Write(equation), run_time=1.5)
        self.play(LaggedStart(*[GrowArrow(a) for a in arrows]), run_time=2)
        self.wait(1)

        # 场景2：构造辅助图形
        # 创建三个正方形
        square_a = Square(side_length=2, color=RED).next_to(tri, LEFT, buff=0)
        square_b = Square(side_length=3, color=BLUE).next_to(tri, DOWN, buff=0)
        square_c = Square(side_length=np.sqrt(13), color=GREEN).rotate(np.arctan(3 / 2)).next_to(tri, UR, buff=0)

        for s, t in zip([square_a, square_b, square_c], ["a^2", "b^2", "c^2"]):
            s.set_fill(opacity=0.2)
            self.play(Create(s), Write(Tex(t).move_to(s.get_center())), run_time=1.5)

        # 镜头聚焦
        self.play(
            self.camera.frame.animate.set_width(square_a.width * 3).move_to(VGroup(square_a, square_b)),
            run_time=2
        )

        # 场景3：图形重组动画
        # 切割正方形为三角形
        triangles = VGroup()
        for square in [square_a, square_b]:
            points = square.get_vertices()
            tri1 = Polygon(points[0], points[1], points[3], color=RED if square == square_a else BLUE)
            tri2 = Polygon(points[1], points[2], points[3], color=RED if square == square_a else BLUE)
            triangles.add(tri1, tri2)

        # 动画重组
        self.play(LaggedStart(*[tri.animate.rotate(PI / 2).shift(UR * 2) for tri in triangles]),
                  run_time=3, rate_func=smooth)

        # 拼合新正方形
        new_square = Square(color=GREEN).surround(triangles)
        self.play(Transform(VGroup(triangles), new_square), run_time=2)

        # 场景4：面积等价证明
        # 并排对比
        original = VGroup(tri.copy(), square_a.copy(), square_b.copy()).scale(0.5).to_edge(LEFT)
        transformed = VGroup(new_square.copy(), square_c.copy()).scale(0.5).to_edge(RIGHT)

        self.play(
            FadeOut(arrows),
            Transform(VGroup(tri, square_a, square_b), original),
            Transform(VGroup(triangles, square_c), transformed),
            run_time=2
        )

        # 显示等式
        final_eq = Tex(r"a^2 + b^2 = c^2", color=YELLOW).scale(2)
        self.play(
            Write(final_eq),
            final_eq.animate.shift(UP * 1.5),
            self.camera.frame.animate.move_to(final_eq),
            run_time=2
        )

        # 粒子特效
        self.add(final_eq)
        self.play(
            final_eq.animate.scale(1.5).set_color_by_gradient(RED, BLUE, GREEN),
            rate_func=there_and_back,
            run_time=3
        )

        # 三维旋转（需启用ThreeDScene）
        self.play(
            final_eq.animate.rotate(30 * DEGREES, axis=UP),
            run_time=2,
            rate_func=smooth
        )

        self.wait(3)
if __name__ == "__main__":
    config = {
        "quality": "high_quality",
        "preview": True,
        "media_dir": "./output",
        "pixel_height": 1080,
        "pixel_width": 1920,
        "background_color": "#333333",  # 深色背景
        "format": "png"  # 解决部分系统渲染问题
    }

    with tempconfig(config):
        scene = PythagoreanProof()
        scene.render()