from manim import *
import numpy as np


class PythagoreanProof(Scene):
    def update_step(self, current_text, new_content, duration=1.5):
        """带淡入淡出效果的步骤更新"""
        new_text = Text(new_content, font="Microsoft YaHei", font_size=32).to_edge(UP, buff=0.8)
        new_text.set_color_by_gradient(GOLD, ORANGE)

        if current_text:
            self.play(
                FadeOut(current_text, shift=UP * 0.3),
                FadeIn(new_text, shift=DOWN * 0.3),
                run_time=0.8
            )
        else:
            self.play(FadeIn(new_text, shift=DOWN))

        self.wait(duration)
        return new_text

    def create_square(self, start, end, color):
        """创建与边垂直的正方形"""
        # 计算边的向量
        edge_vector = end - start
        # 获取垂直向量（逆时针90度）
        perp_vector = np.array([-edge_vector[1], edge_vector[0], 0])
        # 标准化垂直向量
        perp_unit = perp_vector / np.linalg.norm(perp_vector)

        # 生成正方形四个顶点
        return Polygon(
            start,
            end,
            end + perp_unit * np.linalg.norm(edge_vector),
            start + perp_unit * np.linalg.norm(edge_vector),
            color=color,
            fill_opacity=0.8,
            stroke_width=4
        )

    def construct(self):
        current_step = None
        elements = VGroup()

        # === 步骤1：绘制直角三角形 ===
        current_step = self.update_step(None, "步骤1：绘制直角三角形")

        # 定义三角形顶点（使用NumPy数组）
        A = np.array([-2.0, -1.0, 0.0])
        B = np.array([2.0, -1.0, 0.0])
        C = np.array([0.0, 2.0, 0.0])  # 调整为更明显的直角三角形

        triangle = Polygon(A, B, C, color=BLUE, fill_opacity=0.8, stroke_width=4)
        elements.add(triangle)
        self.play(Create(triangle), run_time=2)
        self.wait()

        # === 步骤2：构建三个正方形 ===
        current_step = self.update_step(current_step, "步骤2：构建三个正方形")

        # 创建三个边的正方形（正确垂直方向）
        square_AB = self.create_square(A, B, ORANGE)
        square_BC = self.create_square(B, C, GREEN)
        square_AC = self.create_square(C, A, YELLOW)

        # 调整正方形位置使其向外扩展
        square_AB.shift((A + B) / 2 - square_AB.get_center() + DOWN * 0.5)
        square_BC.rotate(PI / 2, about_point=B)
        square_AC.rotate(-PI / 2, about_point=C)

        self.play(
            LaggedStart(
                Create(square_AB),
                Create(square_BC),
                Create(square_AC),
                lag_ratio=0.3
            ),
            run_time=4
        )
        elements.add(square_AB, square_BC, square_AC)
        self.wait()

        # === 步骤3：面积变换演示 ===
        current_step = self.update_step(current_step, "步骤3：面积恒等变换")

        # 创建可移动三角形（从AC边开始）
        moving_tri = self.create_square(A, C, PINK)
        self.play(FadeIn(moving_tri), run_time=1)

        # 旋转动画
        self.play(
            Rotate(
                moving_tri,
                angle=PI / 2,
                about_point=A,
                rate_func=smooth
            ),
            run_time=2
        )
        self.wait()

        # === 步骤4：最终结论 ===
        current_step = self.update_step(current_step, "步骤4：得出结论")

        # 数学公式
        formula = MathTex(r"a^2 + b^2 = c^2", color=GOLD).scale(2)
        formula_box = SurroundingRectangle(formula, color=WHITE, buff=0.5)

        self.play(
            FadeOut(elements),
            FadeOut(moving_tri),
            run_time=1
        )
        self.play(
            DrawBorderThenFill(formula_box),
            Write(formula),
            run_time=2
        )
        self.wait(3)


if __name__ == "__main__":
    config = {
        "quality": "high_quality",
        "preview": True,
        "media_dir": "./output",
        "pixel_height": 1080,
        "pixel_width": 1920,
        "background_color": "#333333"
    }

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