# difference_of_squares_final.py
from manim import *
import numpy as np

class SquareDifference(Scene):
    def create_bubble(self, text, direction=UP):
        """创建自定义对话气泡（完全兼容版）"""
        # 基础形状组件
        bubble = VGroup()
        
        # 文本内容（确保中文字体可用）
        text_mob = Text(text, font="Microsoft YaHei", font_size=24)
        text_mob.scale(0.8)
        
        # 背景框
        rect = RoundedRectangle(
            width=text_mob.width + 0.5,
            height=text_mob.height + 0.3,
            corner_radius=0.2,
            fill_color=WHITE,
            fill_opacity=1,
            stroke_color=BLACK,
            stroke_width=2
        )
        
        # 箭头方向处理
        arrow = Triangle().scale(0.3)
        arrow.set_fill(WHITE, 1)
        arrow.set_stroke(BLACK, 2)
        
        # 使用向量判断方向
        if np.array_equal(direction, UP):
            arrow.rotate(-PI/2)
            bubble.add(rect, arrow)
            bubble.arrange(DOWN, buff=0.1)
        elif np.array_equal(direction, DOWN):
            arrow.rotate(PI/2)
            bubble.add(arrow, rect)
            bubble.arrange(DOWN, buff=0.1)
        else:  # 默认上方
            arrow.rotate(-PI/2)
            bubble.add(rect, arrow)
            bubble.arrange(DOWN, buff=0.1)
            
        # 定位文本
        text_mob.move_to(rect)
        bubble.add(text_mob)
        return bubble

    def construct(self):
        # 高清配置
        config.frame_width = 16
        config.frame_height = 9
        config.pixel_height = 1440
        config.pixel_width = 2560
        
        # 参数配置
        a = 4
        b = 1
        colors = {
            "a": "#1F77B4",
            "b": "#FF7F0E",
            "text": "#FFFFFF",
            "bg": "#2E2E2E"
        }
        self.camera.background_color = colors["bg"]
        step_delay = 2.5

        # ========== 标题部分 ==========
        title = Text("东师理想数学课件生成助手", 
                   font_size=48,
                   color=colors["text"],
                   weight=BOLD,
                   stroke_width=8,
                   stroke_color=BLACK)
        
        sub_title = Text("平方差公式可视化演示", 
                       font_size=32, 
                       color=colors["text"],
                       stroke_width=4,
                       stroke_color=BLACK).next_to(title, DOWN)
        
        self.play(
            Write(title, run_time=1.5),
            FadeIn(sub_title, shift=UP),
            run_time=2
        )
        self.wait(step_delay)
        self.play(
            FadeOut(title, shift=UP),
            FadeOut(sub_title, shift=DOWN),
            run_time=1.5
        )

        # ========== 公式推导部分 ==========
        formula_steps = VGroup(
            MathTex("(", "a", "+", "b", ")(", "a", "-", "b", ")", font_size=40),
            MathTex("=", "a", "^2", "-", "b", "^2", font_size=40),
            MathTex("=", f"{a**2}", "-", f"{b**2}", font_size=40),
            MathTex("=", f"{(a+b)*(a-b)}", font_size=48, color=GREEN)
        ).arrange(DOWN, aligned_edge=LEFT, buff=0.8)
        
        formula_board = SurroundingRectangle(
            formula_steps, 
            color=WHITE, 
            stroke_width=2,
            fill_color=BLACK,
            fill_opacity=0.7,
            corner_radius=0.2
        )
        
        self.play(
            DrawBorderThenFill(formula_board),
            FadeIn(formula_steps, shift=UP),
            run_time=2
        )
        self.wait(step_delay)

        # ========== 坐标系部分（最终修复） ==========
        # 使用新版背景线直接生成网格
        axes = Axes(
            x_range=[0, a+b+1, 1],
            y_range=[0, a+b+1, 1],
            x_length=8,
            y_length=6,
            axis_config={
                "color": colors["text"],
                "stroke_width": 3,
                "include_numbers": True
            }
        ).shift(DOWN*0.5)
        

        # 大正方形动画
        big_square = Rectangle(
            width=a+b,
            height=a+b,
            color=colors["a"],
            fill_opacity=0.4,
            stroke_width=4
        ).move_to(axes.c2p((a+b)/2, (a+b)/2))
        
        self.play(
            formula_board.animate.scale(0.8).to_edge(UP, buff=0.5),
            Create(axes, run_time=2),
            DrawBorderThenFill(big_square),
            run_time=2
        )
        self.wait(step_delay)


        # ========== 剩余动画部分 ==========
        # 垂直分割线
        v_line = Line(
            start=axes.c2p(a, 0),
            end=axes.c2p(a, a+b),
            color=colors["text"],
            stroke_width=4
        )
        v_label = MathTex("a", 
                        color=colors["text"],
                        stroke_width=3,
                        stroke_color=BLACK).next_to(v_line, RIGHT, buff=0.2)
        
        # 水平分割线
        h_line = Line(
            start=axes.c2p(0, a),
            end=axes.c2p(a+b, a),
            color=colors["text"],
            stroke_width=4
        )
        h_label = MathTex("a", 
                        color=colors["text"],
                        stroke_width=3,
                        stroke_color=BLACK).next_to(h_line, UP, buff=0.2)
        
        self.play(
            Create(v_line),
            Write(v_label),
            Create(h_line),
            Write(h_label),
            run_time=2
        )
        self.wait(step_delay)

        # ========== 面积计算部分 ==========
        part_a = Rectangle(
            width=a,
            height=a,
            color=colors["a"],
            fill_opacity=0.6
        ).move_to(axes.c2p(a/2, a/2))
        
        part_b = Rectangle(
            width=b,
            height=a+b,
            color=colors["b"],
            fill_opacity=0.6
        ).move_to(axes.c2p(a + b/2, (a+b)/2))
        
        self.play(
            Transform(big_square, part_a),
            FadeIn(part_b),
            run_time=2
        )
        self.wait(step_delay)

        # 移动并计算面积差
        moved_part = part_b.copy().shift(DOWN*b)
        area_diff = Rectangle(
            width=b,
            height=2*b,
            color=RED,
            fill_opacity=0.3
        ).move_to(axes.c2p(a + b/2, a - b))
        
        self.play(
            part_b.animate.shift(DOWN*b),
            FadeIn(area_diff),
            run_time=2
        )
        self.wait(step_delay)

        # ========== 最终联动 ==========
        final_formula = formula_steps[-1].copy()
        final_formula.generate_target()
        final_formula.target.scale(1.5).move_to(ORIGIN)
        
        self.play(
            MoveToTarget(final_formula),
            FadeOut(axes),
            FadeOut(part_a),
            FadeOut(part_b),
            FadeOut(area_diff),
            run_time=2
        )
        self.wait(step_delay*2)

        # 结束动画
        self.play(
            final_formula.animate.scale(0.5).to_edge(DOWN, buff=1),
            run_time=1.5
        )
        self.wait(2)

        # 版权信息
        copyright = Text("© 东师理想 版权所有", 
                      font_size=24, 
                      color=colors["text"],
                      stroke_width=2,
                      stroke_color=BLACK).to_edge(DOWN, buff=0.3)
        self.play(Write(copyright), run_time=1)
        self.wait(3)

        # 创建小熊角色（优化动画版）
        bear = VGroup(
            Circle(radius=0.5, color=GOLD, fill_opacity=1),  # 头
            Circle(radius=0.1, fill_opacity=1).shift(UR*0.3),  # 右耳
            Circle(radius=0.1, fill_opacity=1).shift(UL*0.3),  # 左耳
            Ellipse(width=0.8, height=1.2, fill_opacity=1).shift(DOWN*0.8)  # 身体
        ).set_color([GOLD_E, LIGHT_BROWN])
        bear.scale(0.8).to_corner(DL)
        
        # 场景1：小熊出场
        self.play(DrawBorderThenFill(bear), run_time=1.5)
        self.play(
            bear.animate.shift(UP*0.5).scale(1.2),
            rate_func=there_and_back_with_pause,
            run_time=1.5
        )
        
        # 创建对话气泡
        bubble = self.create_bubble("大家好！今天我们来认识一个神奇的三角形")
        bubble.next_to(bear, UP, buff=0.3)
        self.play(
            LaggedStart(
                DrawBorderThenFill(bubble[0]),
                DrawBorderThenFill(bubble[1]),
                Write(bubble[2]),
                lag_ratio=0.3
            ),
            run_time=2
        )
        self.wait(2)
        self.play(FadeOut(bubble))
        
        # 创建3-4-5直角三角形
        triangle = Polygon(
            [-2, -1, 0], [2, -1, 0], [2, 2, 0],
            color=WHITE, 
            fill_opacity=0.5,
            stroke_width=4
        )
        labels = VGroup(
            MathTex("3").next_to(triangle, LEFT, buff=0.3),
            MathTex("4").next_to(triangle, DOWN, buff=0.3),
            MathTex("5", color=YELLOW).next_to(triangle.get_center(), UR, buff=0.3)
        )
        
        self.play(DrawBorderThenFill(triangle), run_time=2)
        self.play(LaggedStart(*[Write(l) for l in labels], lag_ratio=0.3))
        self.play(
            triangle.animate.shift(LEFT*2),
            bear.animate.look_at(triangle),
            run_time=1.5
        )
        self.wait()
        
        # 场景2：构建正方形
        square3 = Square(side_length=3, color=PINK, fill_opacity=0.3)
        square4 = Square(side_length=4, color=BLUE, fill_opacity=0.3)
        square5 = Square(side_length=5, color=GREEN, fill_opacity=0.3)
        
        square3.next_to(triangle, LEFT, buff=0.5)
        square4.next_to(triangle, DOWN, buff=0.5)
        square5.rotate(np.arctan(3/4)).next_to(triangle, UR, buff=0.5)
        
        self.play(
            LaggedStart(
                DrawBorderThenFill(square3),
                DrawBorderThenFill(square4),
                DrawBorderThenFill(square5),
                lag_ratio=0.5
            ),
            bear.animate.wave(rate_func=there_and_back),
            run_time=3
        )
        
        # 添加面积标签
        area_text = VGroup(
            MathTex("3^2=9", color=PINK),
            MathTex("4^2=16", color=BLUE),
            MathTex("5^2=25", color=GREEN)
        ).arrange(DOWN, buff=0.5, aligned_edge=LEFT).to_edge(UR)
        
        self.play(TransformFromCopy(square3, area_text[0]))
        self.play(Indicate(square3, scale_factor=1.2))
        self.play(TransformFromCopy(square4, area_text[1]))
        self.play(Indicate(square4, scale_factor=1.2))
        self.play(TransformFromCopy(square5, area_text[2]))
        self.play(Indicate(square5, scale_factor=1.2))
        self.wait()
        
        # 场景3：组合动画
        self.play(
            square3.animate.shift(UP*2 + LEFT*1).rotate(PI/3),
            square4.animate.shift(RIGHT*3 + DOWN*1).rotate(-PI/4),
            bear.animate.blink(rate_func=there_and_back),
            run_time=2
        )
        
        magic_text = Text("看！9 + 16 = 25", font="Microsoft YaHei", color=YELLOW).scale(0.8)
        self.play(Write(magic_text))
        self.play(
            square3.animate.move_to(ORIGIN).set_opacity(0.3),
            square4.animate.move_to(ORIGIN).set_opacity(0.3),
            square5.animate.move_to(ORIGIN).set_opacity(0.3),
            run_time=2
        )
        self.wait()
        
        # 场景4：公式推导
        equation = MathTex("3^2 + 4^2 = 5^2", color=YELLOW).scale(1.5)
        equation_box = SurroundingRectangle(equation, color=BLUE, buff=0.5, corner_radius=0.2)
        self.play(
            ReplacementTransform(magic_text, equation),
            Create(equation_box),
            bear.animate.jump(rate_func=there_and_back),
            run_time=2
        )
        self.wait(2)
        
        # 场景5：告别
        bye_bubble = self.create_bubble("下次再见啦！")
        bye_bubble.next_to(bear, UP, buff=0.3)
        self.play(
            FadeOut(VGroup(triangle, square3, square4, square5, equation_box, labels, area_text)),
            LaggedStart(
                DrawBorderThenFill(bye_bubble[0]),
                DrawBorderThenFill(bye_bubble[1]),
                Write(bye_bubble[2]),
                lag_ratio=0.3
            ),
            bear.animate.happy().scale(1.2),
            run_time=2
        )
        self.play(
            bear.animate.rotate(PI/2, about_point=bear.get_center()).shift(DOWN*2),
            run_time=2
        )
        self.wait(2)

if __name__ == "__main__":
    config.background_color = "#F0F8FF"  # 浅蓝背景
    config.frame_width = 16
    config.frame_height = 9
    config.pixel_width = 1280
    config.pixel_height = 720
    config.media_dir = "./output"
    scene = SquareDifference()
    scene.render()