from manim import *

class ClearProof(Scene):
    def construct(self):
        config.frame_width = 12.8
        config.frame_height = 10.24
        colors = {
            "a": "#4D9DE0",
            "b": "#E15554",
            "c": "#3BB273",
            "bg": "#1E1E1E",
            "text": "#FFFFFF"
        }
        self.camera.background_color = colors["bg"]
        a, b = 3, 4  # 直角边长度
        c = 5        # 斜边长度

        # ========== 第一步：初始构造 ==========
        step1 = Text("构造直角三角形", font_size=36, color=colors["text"]).to_edge(UP)
        tri = Polygon([-4, -2, 0], [2, -2, 0], [-4, 2, 0], 
                    color=colors["text"], stroke_width=4)
        labels = VGroup(
            MathTex("a", color=colors["a"]).next_to(tri.get_bottom(), DOWN, buff=0.3),
            MathTex("b", color=colors["b"]).next_to(tri.get_left(), LEFT, buff=0.3),
            MathTex("c", color=colors["c"]).move_to(tri.get_vertices()[2] + UR*0.5)
        )
        
        self.play(Write(step1), run_time=1.5)
        self.play(Create(tri), run_time=2)
        self.play(LaggedStart(*[Write(l) for l in labels], lag_ratio=0.4))
        self.wait(2)

        # ========== 第二步：构建正方形 ==========
        step2 = Text("构建边长的正方形", font_size=36, color=colors["text"]).to_edge(UP)
        square_a = Square(a, color=colors["a"], fill_opacity=0.3).next_to(tri, RIGHT, buff=2)
        square_b = Square(b, color=colors["b"], fill_opacity=0.3).next_to(tri, LEFT, buff=2)
        square_c = Square(c, color=colors["c"], fill_opacity=0.3).move_to(ORIGIN + DOWN*1.5)
        
        self.play(
            ReplacementTransform(step1, step2),
            DrawBorderThenFill(square_a),
            DrawBorderThenFill(square_b),
            run_time=2
        )
        self.wait(1.5)
        
        # ========== 第三步：分割图形 ==========
        step3 = Text("分割正方形为可重组部件", font_size=36, color=colors["text"]).to_edge(UP)
        # 分割a²正方形
        a_pieces = VGroup(
            tri.copy().set_fill(colors["a"], 0.3),
            tri.copy().rotate(PI/2).set_fill(colors["a"], 0.3).shift(RIGHT*a)
        )
        # 分割b²正方形
        b_pieces = VGroup(
            tri.copy().rotate(-PI/2).set_fill(colors["b"], 0.3),
            tri.copy().rotate(PI).set_fill(colors["b"], 0.3).shift(LEFT*b)
        )
        
        self.play(
            ReplacementTransform(step2, step3),
            square_a.animate.set_fill(opacity=0),
            square_b.animate.set_fill(opacity=0),
            LaggedStart(
                Transform(square_a, a_pieces),
                Transform(square_b, b_pieces),
                lag_ratio=0.5
            ),
            run_time=3
        )
        self.wait(2)

        # ========== 第四步：重组图形 ==========
        step4 = Text("重组部件构成大正方形", font_size=36, color=colors["text"]).to_edge(UP)
        # 计算重组后的位置
        final_positions = [
            [-c/2, -c/2 -1.5, 0], [c/2, -c/2 -1.5, 0],
            [c/2, c/2 -1.5, 0], [-c/2, c/2 -1.5, 0]
        ]
        # 创建重组动画
        animations = []
        for i, piece in enumerate(a_pieces):
            animations.append(piece.animate.move_to(final_positions[i]))
        for i, piece in enumerate(b_pieces):
            animations.append(piece.animate.move_to(final_positions[i+2]))
        
        self.play(
            ReplacementTransform(step3, step4),
            LaggedStart(*animations, lag_ratio=0.3),
            FadeIn(square_c, shift=UP),
            run_time=4
        )
        self.wait(2)

        # ========== 第五步：面积等式 ==========
        step5 = Text("面积守恒证明定理", font_size=36, color=colors["text"]).to_edge(UP)
        equation = MathTex(
            "a^2", "+", "b^2", "=", "c^2",
            substrings_to_isolate=["a^2", "b^2", "c^2"]
        ).scale(1.5).set_color_by_tex("a^2", colors["a"]).set_color_by_tex("b^2", colors["b"]).set_color_by_tex("c^2", colors["c"])
        
        self.play(
            ReplacementTransform(step4, step5),
            FadeOut(VGroup(a_pieces, b_pieces, square_c)),
            Write(equation),
            run_time=2
        )
        self.wait(3)

        # ========== 最终总结 ==========
        conclusion = Text("勾股定理得证！", font_size=48, color=colors["c"])
        self.play(
            FadeOut(equation),
            FadeOut(step5),
            Write(conclusion, run_time=2)
        )
        self.wait(3)