from manim import *


class BinomialExpansion(Scene):

    def construct(self):
        # Define the side lengths
        a = 2
        b = 1
        
        # Create the large square
        square = Square(side_length=a + b, color=BLUE)
        square.move_to(ORIGIN)
        
        # Create the smaller squares and rectangles
        square_a = Square(side_length=a, color=RED)
        square_a.move_to(square.get_corner(LEFT + UP))
        
        square_b = Square(side_length=b, color=GREEN)
        square_b.move_to(square.get_corner(RIGHT + DOWN))
        
        rectangle_ab1 = Rectangle(width=a, height=b, color=YELLOW)
        rectangle_ab1.move_to(square.get_corner(LEFT + DOWN))
        
        rectangle_ab2 = Rectangle(width=b, height=a, color=YELLOW)
        rectangle_ab2.move_to(square.get_corner(RIGHT + UP))
        
        # Display the large square
        self.play(Create(square))
        self.wait(1)
        
        # Display the smaller squares and rectangles
        self.play(Create(square_a), Create(square_b), Create(rectangle_ab1), Create(rectangle_ab2))
        self.wait(1)
        
        # Label the areas
        label_a = MathTex('a^2').next_to(square_a, UP)
        label_b = MathTex('b^2').next_to(square_b, DOWN)
        label_ab1 = MathTex('ab').next_to(rectangle_ab1, DOWN)
        label_ab2 = MathTex('ab').next_to(rectangle_ab2, RIGHT)
        
        self.play(Write(label_a), Write(label_b), Write(label_ab1), Write(label_ab2))
        self.wait(2)
        
        # Display the equation
        equation = MathTex('(a + b)^2 = a^2 + 2ab + b^2').to_edge(DOWN)
        self.play(Write(equation))
        self.wait(2)

# To run the animation, use the following command in the terminal:
# conda activate py310
# cd D:\dsWork\QingLong\AI\Manim>
# manim -pql binomial_expansion.py BinomialExpansion