from manim import *
import numpy as np

# Color Palette — Classic 3B1B
BG = "#1C1C1C"
PRIMARY = "#58C4DD"      # BLUE — a side, a²
SECONDARY = "#83C167"    # GREEN — b side, b²
ACCENT = "#FFFF00"       # YELLOW — c side, c², highlights
WHITE_C = "#FFFFFF"
MONO = "Menlo"

# Timing
FAST = 0.8; NORMAL = 1.5; SLOW = 2.5

# Triangle: right angle at origin, legs along axes
# a = 3 (vertical), b = 4 (horizontal), c = 5 (hypotenuse)


class Scene1_Title(Scene):
    def construct(self):
        self.camera.background_color = BG

        title = Text(
            "The Pythagorean Theorem",
            font_size=52,
            color=WHITE_C,
            weight=BOLD,
            font=MONO,
        )
        title.move_to(UP * 0.6)

        eq = MathTex(r"a^2 + b^2 = c^2", font_size=72)
        eq.set_color(ACCENT)
        eq.move_to(DOWN * 0.8)

        tagline = Text(
            "Proved not by numbers — but by area.",
            font_size=22,
            color=WHITE_C,
            font=MONO,
            opacity=0.55
        )
        tagline.move_to(DOWN * 2.5)

        self.add_subcaption("Proved not by numbers, but by area.", duration=3)
        self.play(FadeIn(title, run_time=SLOW), run_time=SLOW)
        self.wait(1.5)
        self.play(Write(eq), run_time=SLOW)
        self.wait(1.0)
        self.play(FadeIn(tagline), run_time=NORMAL)
        self.wait(3.0)
        self.play(FadeOut(Group(*self.mobjects)), run_time=0.5)


class Scene2_Triangle(Scene):
    def construct(self):
        self.camera.background_color = BG

        title = Text("The Right Triangle", font_size=38, color=WHITE_C, font=MONO, weight=BOLD)
        title.to_edge(UP, buff=0.6)
        self.add(title)

        # Right triangle: vertices at (0,0), (4,0), (0,3)
        a_len, b_len = 3, 4
        hyp = np.sqrt(a_len**2 + b_len**2)  # = 5

        tri = Polygon(
            [0, 0, 0], [b_len, 0, 0], [0, a_len, 0],
            stroke_color=WHITE_C,
            stroke_opacity=0.9,
            fill_opacity=0,
            stroke_width=2.5
        )

        # Right angle marker at origin
        ra = RightAngle(
            Line([0, 0, 0], [0.6, 0, 0]),
            Line([0, 0, 0], [0, 0.6, 0]),
            length=0.3,
            color=WHITE_C,
            stroke_width=2
        )

        # Labels
        a_label = MathTex(r"a", font_size=40, color=PRIMARY)
        a_label.next_to(Line([0, 0, 0], [0, a_len, 0]).get_center(), LEFT, buff=0.15)

        b_label = MathTex(r"b", font_size=40, color=SECONDARY)
        b_label.next_to(Line([0, 0, 0], [b_len, 0, 0]).get_center(), DOWN, buff=0.15)

        c_label = MathTex(r"c", font_size=40, color=ACCENT)
        c_pos = Line([b_len, 0, 0], [0, a_len, 0]).get_center()
        c_label.move_to(c_pos + [0.3, 0.3, 0])

        # Numeric labels
        nums = MathTex(r"a = 3,\; b = 4,\; c = 5", font_size=28, color=WHITE_C)
        nums.to_edge(DOWN, buff=1.2)

        self.add_subcaption("A right triangle with sides 3, 4, and 5.", duration=3)
        self.play(Create(tri), run_time=SLOW)
        self.wait(1.0)
        self.play(Create(ra), run_time=FAST)
        self.wait(0.8)
        self.play(Write(a_label), Write(b_label), run_time=NORMAL)
        self.wait(0.5)
        self.play(Write(c_label), run_time=NORMAL)
        self.wait(1.0)
        self.add_subcaption("The hypotenuse is always opposite the right angle.", duration=3)
        self.play(Write(nums), run_time=NORMAL)
        self.wait(3.0)
        self.play(FadeOut(Group(*self.mobjects)), run_time=0.5)


class Scene3_ThreeSquares(Scene):
    def construct(self):
        self.camera.background_color = BG

        title = Text("Three Squares", font_size=38, color=WHITE_C, font=MONO, weight=BOLD)
        title.to_edge(UP, buff=0.6)
        self.add(title)

        # Triangle (smaller, for reference)
        tri = Polygon(
            [0, 0, 0], [4, 0, 0], [0, 3, 0],
            stroke_color=WHITE_C,
            stroke_opacity=0.3,
            fill_opacity=0,
            stroke_width=1.5
        )
        tri.move_to(LEFT * 3.5 + DOWN * 0.5)

        # Square on a (blue, side = 3)
        sq_a = Square(side_length=3, stroke_color=PRIMARY, fill_color=PRIMARY, fill_opacity=0.25, stroke_width=2.5)
        sq_a.move_to(LEFT * 3.5 + DOWN * 0.5)
        sq_a_label = MathTex(r"a^2", font_size=36, color=PRIMARY)
        sq_a_num = MathTex(r"= 9", font_size=28, color=PRIMARY)

        # Square on b (green, side = 4)
        sq_b = Square(side_length=4, stroke_color=SECONDARY, fill_color=SECONDARY, fill_opacity=0.2, stroke_width=2.5)
        sq_b_label = MathTex(r"b^2", font_size=36, color=SECONDARY)
        sq_b_num = MathTex(r"= 16", font_size=28, color=SECONDARY)

        # Square on c (yellow, side = 5)
        sq_c = Square(side_length=5, stroke_color=ACCENT, fill_color=ACCENT, fill_opacity=0.15, stroke_width=2.5)
        sq_c_label = MathTex(r"c^2", font_size=36, color=ACCENT)
        sq_c_num = MathTex(r"= 25", font_size=28, color=ACCENT)

        # Position: triangle left, squares stacked right
        tri.move_to(LEFT * 3.5 + DOWN * 0.5)
        sq_a.next_to(tri, RIGHT, buff=0.8)
        sq_b.next_to(sq_a, RIGHT, buff=0.8)
        sq_c.next_to(sq_b, RIGHT, buff=0.8)

        sq_a_lbl = sq_a_label.copy().move_to(sq_a.get_center())
        sq_b_lbl = sq_b_label.copy().move_to(sq_b.get_center())
        sq_c_lbl = sq_c_label.copy().move_to(sq_c.get_center())

        sq_a_n = sq_a_num.copy().next_to(sq_a_lbl, DOWN, buff=0.2)
        sq_b_n = sq_b_num.copy().next_to(sq_b_lbl, DOWN, buff=0.2)
        sq_c_n = sq_c_num.copy().next_to(sq_c_lbl, DOWN, buff=0.2)

        self.add_subcaption("Each side builds its own square. The areas are no coincidence.", duration=3)
        self.play(Create(tri), run_time=NORMAL)
        self.wait(0.5)
        self.play(FadeIn(sq_a, run_time=SLOW), run_time=SLOW)
        self.play(Write(sq_a_lbl), run_time=FAST)
        self.play(Write(sq_a_n), run_time=FAST)
        self.wait(1.5)

        self.add_subcaption("The square on side a has area 9.", duration=2)
        self.play(FadeIn(sq_b, run_time=SLOW), run_time=SLOW)
        self.play(Write(sq_b_lbl), run_time=FAST)
        self.play(Write(sq_b_n), run_time=FAST)
        self.wait(1.5)

        self.add_subcaption("The square on side b has area 16.", duration=2)
        self.play(FadeIn(sq_c, run_time=SLOW), run_time=SLOW)
        self.play(Write(sq_c_lbl), run_time=FAST)
        self.play(Write(sq_c_n), run_time=FAST)
        self.wait(1.5)

        self.add_subcaption("And the square on the hypotenuse c has area 25.", duration=3)
        self.wait(3.0)
        self.play(FadeOut(Group(*self.mobjects)), run_time=0.5)


class Scene4_Proof(Scene):
    def construct(self):
        self.camera.background_color = BG

        title = Text("The Proof by Area", font_size=38, color=WHITE_C, font=MONO, weight=BOLD)
        title.to_edge(UP, buff=0.6)
        self.add(title)

        # Large yellow square (c² = 25)
        # Arrange 4 right triangles inside a 5×5 square:
        # The remaining area = 25 - 2*(3*4)/2*2 = 25 - 24 = 1? No...
        # Classic proof: 4 triangles + 2 squares (a² and b²) = c²
        # Layout: large square of side (a+b) = 7, with 4 triangles cut out
        # Remaining: (a+b)² - 4*(ab/2) = a² + 2ab + b² - 2ab = a² + b² = c²

        a, b = 3, 4
        c = 5
        sq_side = a + b  # = 7

        # Large container square
        container = Square(
            side_length=sq_side,
            stroke_color=ACCENT,
            stroke_opacity=0.4,
            stroke_width=1.5,
            fill_opacity=0
        )
        container.move_to(ORIGIN)

        # Four triangles (each area = ab/2 = 6)
        tri_kwargs = {"stroke_width": 1.5, "fill_opacity": 0.4}

        t_bl = Polygon([0,0,0], [b,0,0], [b,a,0], [0,a,0], stroke_color=PRIMARY, fill_color=PRIMARY, **tri_kwargs)  # bottom-left
        t_br = Polygon([b,0,0], [b+a,0,0], [b+a,a,0], [b,a,0], stroke_color=SECONDARY, fill_color=SECONDARY, **tri_kwargs)  # bottom-right
        t_tr = Polygon([b,a,0], [b+a,a,0], [b+a,b+a,0], [b,a+b,0], stroke_color=PRIMARY, fill_color=PRIMARY, **tri_kwargs)  # top-right
        t_tl = Polygon([0,a,0], [b,a,0], [b,a+b,0], [0,a+b,0], stroke_color=SECONDARY, fill_color=SECONDARY, **tri_kwargs)  # top-left

        # Two inner squares (a² and b²)
        inner_a = Square(side_length=a, stroke_color=PRIMARY, fill_color=PRIMARY, fill_opacity=0.7, stroke_width=2)
        inner_b = Square(side_length=b, stroke_color=SECONDARY, fill_color=SECONDARY, fill_opacity=0.5, stroke_width=2)

        # Position inner squares in the center gap
        # The gap is a square of side (a-b) = 1... no.
        # With 4 triangles at corners of (a+b) square:
        # Remaining is two squares: one a×a and one b×b, arranged diagonally
        # Actually in the classic proof: 4 triangles form an X leaving two squares
        inner_a.move_to([b + a/2, a/2, 0])
        inner_b.move_to([b/2, a + b/2, 0])

        la = MathTex(r"a^2", font_size=32, color=PRIMARY)
        la.move_to(inner_a.get_center())

        lb = MathTex(r"b^2", font_size=32, color=SECONDARY)
        lb.move_to(inner_b.get_center())

        # Labels for the 4 triangles
        tri_label = Text("4 × (ab/2)", font_size=20, color=WHITE_C, opacity=0.5)

        self.add_subcaption("Four identical right triangles form a frame around two squares.", duration=3)
        self.play(Create(container), run_time=NORMAL)
        self.wait(0.5)
        self.play(FadeIn(t_bl), FadeIn(t_br), FadeIn(t_tr), FadeIn(t_tl), run_time=SLOW)
        self.wait(1.5)

        self.add_subcaption("The four triangles have total area 4 times ab over 2.", duration=3)
        self.play(Write(tri_label), run_time=NORMAL)
        self.wait(1.5)

        self.add_subcaption("What's left? A square of area a² and a square of area b².", duration=3)
        self.play(FadeIn(inner_a), FadeIn(inner_b), run_time=SLOW)
        self.play(Write(la), Write(lb), run_time=NORMAL)
        self.wait(2.0)

        # Alternative visualization: rearrange
        alt_title = Text("Rearranged", font_size=26, color=ACCENT, font=MONO, opacity=0.8)
        alt_title.to_edge(DOWN, buff=0.8)
        self.add_subcaption("Now rearrange. Does the large square still fit perfectly?", duration=3)
        self.play(Write(alt_title), run_time=NORMAL)
        self.wait(3.0)
        self.play(FadeOut(Group(*self.mobjects)), run_time=0.5)


class Scene5_Equation(Scene):
    def construct(self):
        self.camera.background_color = BG

        title = Text("The Algebra", font_size=38, color=WHITE_C, font=MONO, weight=BOLD)
        title.to_edge(UP, buff=0.6)
        self.add(title)

        eq1 = MathTex(r"a^2 + b^2 = c^2", font_size=60)
        eq1.to_edge(UP).shift(DOWN * 0.8)

        eq2 = MathTex(r"3^2 + 4^2 = 5^2", font_size=56)
        eq2.move_to(DOWN * 0.5)

        eq3 = MathTex(r"9 + 16 = 25", font_size=52)
        eq3.move_to(DOWN * 2.0)

        self.add_subcaption("In numbers, it looks like this.", duration=2)
        self.play(Write(eq1), run_time=SLOW)
        self.wait(2.0)
        self.play(ReplacementTransform(eq1.copy(), eq2), run_time=SLOW)
        self.wait(2.0)
        self.add_subcaption("Nine plus sixteen equals twenty-five.", duration=2)
        self.play(ReplacementTransform(eq2.copy(), eq3), run_time=SLOW)
        self.wait(3.0)

        # Highlight result
        highlight = SurroundingRectangle(eq3, color=ACCENT, buff=0.3, stroke_opacity=0.7)
        self.play(Create(highlight), run_time=NORMAL)
        self.wait(2.0)
        self.play(FadeOut(Group(*self.mobjects)), run_time=0.5)


class Scene6_TheSilence(Scene):
    def construct(self):
        self.camera.background_color = BG

        eq = MathTex(r"a^2 + b^2 = c^2", font_size=72)
        eq.set_color_by_tex("a^2", PRIMARY)
        eq.set_color_by_tex("b^2", SECONDARY)
        eq.set_color_by_tex("c^2", ACCENT)

        sub = Text(
            '"A beautiful truth, timeless and self-evident."',
            font_size=22,
            color=WHITE_C,
            font=MONO,
            opacity=0.5
        )
        sub.to_edge(DOWN, buff=1.8)

        year = Text("Proved ~500 BC by Pythagoras", font_size=16, color=WHITE_C, font=MONO, opacity=0.3)
        year.to_edge(DOWN, buff=0.8)

        self.add_subcaption("Proved 2500 years ago. Still perfect.", duration=4)
        self.play(FadeIn(eq, run_time=SLOW * 1.5), run_time=SLOW * 1.5)
        self.wait(3.0)
        self.play(Write(sub), run_time=NORMAL)
        self.wait(1.5)
        self.play(Write(year), run_time=NORMAL)
        self.wait(4.0)
        self.play(FadeOut(Group(*self.mobjects)), run_time=0.5)