from manim import *
import numpy as np
import math

class GeneratedScene(Scene):
    def safe_text(self, text, max_width=8, font_size=28, color="#F5F0E1"):
        t = Text(text, font_size=font_size, color=color)
        if t.width > max_width:
            t.scale(max_width / t.width)
        return t

    def construct(self):
        self.camera.background_color = "#0A1628"

        IVORY = "#F5F0E1"
        LIME  = "#B8E835"
        CYAN  = "#4FD1C5"

        A = 3.0   # leg a
        B = 2.0   # leg b
        S = A + B  # 5
        H = S / 2  # 2.5

        # ===== Persistent frame elements =====
        outer = Square(side_length=S, color=IVORY, stroke_width=2)

        tick_len = 0.18
        ticks = VGroup(
            Line([0.5, H, 0], [0.5, H - tick_len, 0], color=IVORY, stroke_width=1.5),
            Line([H, -0.5, 0], [H - tick_len, -0.5, 0], color=IVORY, stroke_width=1.5),
            Line([-0.5, -H, 0], [-0.5, -H + tick_len, 0], color=IVORY, stroke_width=1.5),
            Line([-H, 0.5, 0], [-H + tick_len, 0.5, 0], color=IVORY, stroke_width=1.5),
        )

        label_a = Text("a", font_size=28, color=IVORY)
        label_a.next_to(outer, UP, buff=0.15).shift(LEFT * 1)
        label_b = Text("b", font_size=28, color=IVORY)
        label_b.next_to(outer, RIGHT, buff=0.15).shift(DOWN * 1.5)

        # ===== Scene 1 — Establish the invariant frame (5 s) =====
        cap1 = self.safe_text("A square of side a + b", max_width=6, font_size=24)
        cap1.to_edge(DOWN, buff=0.35)

        self.play(Create(outer), run_time=1.5)
        self.play(
            LaggedStart(*[FadeIn(t) for t in ticks], lag_ratio=0.25),
            run_time=0.9
        )
        self.play(FadeIn(label_a), FadeIn(label_b), run_time=0.5)
        self.play(FadeIn(cap1), run_time=0.4)
        self.wait(1.2)

        # ===== Scene 2 — Arrangement 1: central c² (6 s) =====
        self.play(FadeOut(cap1), run_time=0.3)

        cap2 = self.safe_text("Four identical right triangles", max_width=6, font_size=24)
        cap2.to_edge(DOWN, buff=0.35)

        # Arrangement 1 triangles — vertices in (right-angle, a-end, b-end) order
        tri1 = Polygon(
            [-H, H, 0], [-H + A, H, 0], [-H, H - B, 0],
            color=IVORY, fill_color=IVORY, fill_opacity=0.12, stroke_width=1.5
        )
        tri2 = Polygon(
            [H, H, 0], [H, H - A, 0], [H - B, H, 0],
            color=IVORY, fill_color=IVORY, fill_opacity=0.12, stroke_width=1.5
        )
        tri3 = Polygon(
            [H, -H, 0], [H - A, -H, 0], [H, -H + B, 0],
            color=IVORY, fill_color=IVORY, fill_opacity=0.12, stroke_width=1.5
        )
        tri4 = Polygon(
            [-H, -H, 0], [-H, -H + A, 0], [-H + B, -H, 0],
            color=IVORY, fill_color=IVORY, fill_opacity=0.12, stroke_width=1.5
        )

        # Tilted central square of side c
        csq_verts = [
            [-H + A, H, 0],   # ( 0.5,  2.5)
            [ H, H - A, 0],    # ( 2.5, -0.5)
            [ H - A, -H, 0],   # (-0.5, -2.5)
            [-H, H - B, 0],    # (-2.5,  0.5)
        ]
        csq = Polygon(
            *csq_verts,
            color=LIME, fill_color=LIME, fill_opacity=0.55, stroke_width=1.5
        )

        self.play(
            LaggedStart(
                FadeIn(tri1, scale=0.7), FadeIn(tri2, scale=0.7),
                FadeIn(tri3, scale=0.7), FadeIn(tri4, scale=0.7),
                lag_ratio=0.08
            ),
            run_time=1.2
        )
        self.play(FadeIn(csq), run_time=0.8)
        self.play(FadeIn(cap2), run_time=0.4)
        self.wait(1.5)

        # ===== Scene 3 — Name the three area quantities (5 s) =====
        self.play(FadeOut(cap2), run_time=0.3)

        csq_label = MathTex("c^2", font_size=36, color=IVORY)
        csq_label.move_to(sum([np.array(v) for v in csq_verts]) / 4)

        dash_a = DashedLine(
            [-H + 0.3, H - 0.3, 0], [-H + A - 0.3, H - 0.3, 0],
            color=IVORY, stroke_width=1.5, dash_length=0.12
        )
        dash_b = DashedLine(
            [-H + 0.3, H - 0.3, 0], [-H + 0.3, H - B + 0.3, 0],
            color=IVORY, stroke_width=1.5, dash_length=0.12
        )
        leg_a_lbl = MathTex("a", font_size=22, color=IVORY)
        leg_a_lbl.next_to(dash_a, DOWN, buff=0.08)
        leg_b_lbl = MathTex("b", font_size=22, color=IVORY)
        leg_b_lbl.next_to(dash_b, RIGHT, buff=0.08)

        eq1 = MathTex(
            r"(a+b)^2 = 4 \cdot \frac{1}{2}ab + c^2",
            font_size=28, color=IVORY
        )
        eq1.to_edge(DOWN, buff=0.35)

        self.play(FadeIn(csq_label), run_time=0.4)
        self.play(
            Create(dash_a), Create(dash_b),
            FadeIn(leg_a_lbl), FadeIn(leg_b_lbl),
            run_time=0.6
        )
        self.play(Write(eq1), run_time=1.2)
        self.wait(1.5)

        # ===== Scene 4 — Rearrange to expose a² and b² (8 s) =====
        self.play(
            FadeOut(eq1), FadeOut(csq_label),
            FadeOut(dash_a), FadeOut(dash_b),
            FadeOut(leg_a_lbl), FadeOut(leg_b_lbl),
            run_time=0.3
        )

        cap4 = self.safe_text("Same pieces. New gaps.", max_width=6, font_size=24)
        cap4.to_edge(DOWN, buff=0.35)

        # Arrangement 2 target polygons — same chirality as sources
        tgt1 = Polygon(
            [-0.5, 2.5, 0], [-0.5, -0.5, 0], [-2.5, 2.5, 0],
            color=IVORY, fill_color=IVORY, fill_opacity=0.12, stroke_width=1.5
        )
        tgt2 = Polygon(
            [2.5, -2.5, 0], [-0.5, -2.5, 0], [2.5, -0.5, 0],
            color=IVORY, fill_color=IVORY, fill_opacity=0.12, stroke_width=1.5
        )
        tgt3 = Polygon(
            [-0.5, -0.5, 0], [2.5, -0.5, 0], [-0.5, -2.5, 0],
            color=IVORY, fill_color=IVORY, fill_opacity=0.12, stroke_width=1.5
        )
        tgt4 = Polygon(
            [-2.5, -0.5, 0], [-2.5, 2.5, 0], [-0.5, -0.5, 0],
            color=IVORY, fill_color=IVORY, fill_opacity=0.12, stroke_width=1.5
        )

        # Gap squares
        asq = Square(
            side_length=A, color=CYAN,
            fill_color=CYAN, fill_opacity=0.55, stroke_width=1.5
        )
        asq.move_to([1.0, 1.0, 0])

        bsq = Square(
            side_length=B, color=LIME,
            fill_color=LIME, fill_opacity=0.55, stroke_width=1.5
        )
        bsq.move_to([-1.5, -1.5, 0])

        self.play(FadeOut(csq), run_time=0.3)
        self.play(
            Transform(tri1, tgt1), Transform(tri3, tgt3),
            run_time=1.6, rate_func=smooth
        )
        self.play(
            Transform(tri2, tgt2), Transform(tri4, tgt4),
            run_time=1.2, rate_func=smooth
        )
        self.play(FadeIn(asq), FadeIn(bsq), run_time=0.8)
        self.play(FadeIn(cap4), run_time=0.4)
        self.wait(1.5)

        # ===== Scene 5 — Confirm the invariant (5 s) =====
        self.play(FadeOut(cap4), run_time=0.3)

        txt5 = VGroup(
            self.safe_text("frame unchanged", max_width=3.5, font_size=20),
            self.safe_text("triangles unchanged", max_width=3.5, font_size=20)
        ).arrange(RIGHT, buff=0.6)
        txt5.to_edge(UP, buff=0.3)

        glow = outer.copy().set_stroke(IVORY, width=5, opacity=0.6)
        self.play(FadeIn(glow), FadeIn(txt5), run_time=0.4)
        self.play(FadeOut(glow), run_time=0.2)

        for tri in [tri1, tri2, tri3, tri4]:
            self.play(tri.animate.set_stroke(width=3.5), run_time=0.1)
            self.play(tri.animate.set_stroke(width=1.5), run_time=0.1)

        eq2 = MathTex(
            r"(a+b)^2 = 4 \cdot \frac{1}{2}ab + (a^2 + b^2)",
            font_size=28, color=IVORY
        )
        eq2.to_edge(DOWN, buff=0.35)
        self.play(Write(eq2), run_time=1.2)
        self.wait(1.5)

        # ===== Scene 6 — Merge to c² = a² + b² (6 s) =====
        self.play(FadeOut(eq2), FadeOut(txt5), run_time=0.3)

        ghost_csq = Polygon(
            *csq_verts,
            color=LIME, fill_color=LIME, fill_opacity=0.3, stroke_width=1
        )
        self.play(FadeIn(ghost_csq), run_time=0.6)

        # Slide b² upward until flush with a²
        self.play(bsq.animate.shift(UP * 2), run_time=1.2, rate_func=smooth)

        # Cross-fade combined region to lime
        self.play(
            asq.animate.set_fill(LIME, opacity=0.55).set_stroke(LIME),
            FadeOut(ghost_csq),
            run_time=0.6
        )

        # Final equation
        final_eq = MathTex(
            r"c^2", r"=", r"a^2", r"+", r"b^2",
            font_size=54
        )
        final_eq[0].set_color(LIME)
        final_eq[1].set_color(IVORY)
        final_eq[2].set_color(CYAN)
        final_eq[3].set_color(IVORY)
        final_eq[4].set_color(LIME)
        final_eq.move_to(ORIGIN)

        self.play(
            FadeOut(tri1), FadeOut(tri2), FadeOut(tri3), FadeOut(tri4),
            FadeOut(asq), FadeOut(bsq),
            FadeOut(outer), FadeOut(ticks),
            FadeOut(label_a), FadeOut(label_b),
            run_time=0.8
        )
        self.play(GrowFromCenter(final_eq), run_time=1.0)
        self.wait(2.0)