from manim import *

class PythagoreanTheorem(Scene):
    """
    A Manim scene that visually proves the Pythagorean Theorem.
    """
    def construct(self):
        # Define side lengths for the right triangle
        a = 3.0
        b = 4.0
        
        # Colors for visual clarity
        tri_color = BLUE
        sq_a_color = RED_C
        sq_b_color = YELLOW_C
        sq_c_color = GREEN_C
        
        # Set up the initial layout and title
        self.camera.background_color = WHITE
        Mobject.set_default(color=BLACK)
        
        title = Title("Visual Proof of the Pythagorean Theorem", color=BLACK)
        equation = MathTex("a^2", "+", "b^2", "=", "c^2", color=BLACK).next_to(title, DOWN, buff=0.5)
        self.play(Write(title), Write(equation))
        self.wait(1)

        self.play(equation.animate.to_corner(UR))
        
        p1 = np.array([a, 0, 0])
        p2 = np.array([a + b, a, 0])
        p3 = np.array([b, a + b, 0])
        p4 = np.array([0, b, 0])
        
        c1 = np.array([0, 0, 0])
        c2 = np.array([a + b, 0, 0])
        c3 = np.array([a + b, a + b, 0])
        c4 = np.array([0, a + b, 0])

        T1 = Polygon(c1, p1, p4, color=tri_color, fill_opacity=0.7)
        T2 = Polygon(c2, p2, p1, color=tri_color, fill_opacity=0.7)
        T3 = Polygon(c3, p3, p2, color=tri_color, fill_opacity=0.7)
        T4 = Polygon(c4, p4, p3, color=tri_color, fill_opacity=0.7)
        sq_c = Polygon(p1, p2, p3, p4, color=sq_c_color, fill_opacity=0.7)
        
        arrangement1 = VGroup(T1, T2, T3, T4, sq_c).center()
        self.play(Create(arrangement1))
        
        c_sq_label = MathTex("c^2", color=BLACK).move_to(sq_c.get_center())
        self.play(Write(c_sq_label))
        self.wait(2)
        
        T1_target = Polygon(np.array([b, 0, 0]), np.array([a+b, 0, 0]), np.array([b, b, 0]), color=tri_color, fill_opacity=0.7)
        T2_target = Polygon(np.array([a+b, 0, 0]), np.array([a+b, b, 0]), np.array([b, b, 0]), color=tri_color, fill_opacity=0.7)
        T3_target = Polygon(np.array([0, b, 0]), np.array([b, b, 0]), np.array([0, a+b, 0]), color=tri_color, fill_opacity=0.7)
        T4_target = Polygon(np.array([b, b, 0]), np.array([b, a+b, 0]), np.array([0, a+b, 0]), color=tri_color, fill_opacity=0.7)

        target_group = VGroup(T1_target, T2_target, T3_target, T4_target).move_to(arrangement1.get_center())

        self.play(
            Transform(T1, T1_target),
            Transform(T2, T2_target),
            Transform(T3, T3_target),
            Transform(T4, T4_target),
            FadeOut(sq_c),
            FadeOut(c_sq_label)
        )
        self.wait(1)

        sq_a = Square(side_length=a, color=sq_a_color, fill_opacity=0.7)
        sq_b = Square(side_length=b, color=sq_b_color, fill_opacity=0.7)
        
        squares_group = VGroup(sq_a, sq_b).center().move_to(arrangement1.get_center())
        sq_a.align_to(T4_target, DOWN).align_to(T1_target, RIGHT)
        sq_b.align_to(T1_target, UP).align_to(T4_target, LEFT)

        a_sq_label = MathTex("a^2", color=BLACK).move_to(sq_a.get_center())
        b_sq_label = MathTex("b^2", color=BLACK).move_to(sq_b.get_center())
        
        self.play(Create(sq_a), Create(sq_b), Write(a_sq_label), Write(b_sq_label))
        self.wait(2)
        self.play(Indicate(equation))
        self.wait(3)