Manimce Best Practices

"""
3D Visualization Patterns for Manim Community

Demonstrates ThreeDScene, 3D axes, surfaces, and camera control.
Adapted from 3b1b patterns for ManimCE.

Run with: manim -pql 3d_visualization.py SceneName
"""

from manim import *
import numpy as np


class Basic3DScene(ThreeDScene):
    """Basic 3D scene with shapes."""

    def construct(self):
        # Set camera orientation
        self.set_camera_orientation(phi=60 * DEGREES, theta=-45 * DEGREES)

        # 3D shapes
        sphere = Sphere(radius=1, color=BLUE)
        cube = Cube(side_length=1.5, color=RED, fill_opacity=0.7)
        cone = Cone(base_radius=0.8, height=1.5, color=GREEN)

        # Position shapes
        sphere.shift(LEFT * 3)
        cone.shift(RIGHT * 3)

        self.play(Create(sphere), Create(cube), Create(cone))
        self.wait()

        # Rotate camera
        self.begin_ambient_camera_rotation(rate=0.3)
        self.wait(4)
        self.stop_ambient_camera_rotation()


class ThreeDAxesExample(ThreeDScene):
    """3D coordinate axes and plotting."""

    def construct(self):
        self.set_camera_orientation(phi=70 * DEGREES, theta=-45 * DEGREES)

        # Create 3D axes
        axes = ThreeDAxes(
            x_range=[-3, 3, 1],
            y_range=[-3, 3, 1],
            z_range=[-2, 2, 1],
            x_length=6,
            y_length=6,
            z_length=4,
        )

        # Axis labels
        x_label = axes.get_x_axis_label(r"x")
        y_label = axes.get_y_axis_label(r"y")
        z_label = axes.get_z_axis_label(r"z")

        self.play(Create(axes))
        self.add_fixed_orientation_mobjects(x_label, y_label, z_label)
        self.wait()

        # Add a point
        point = Dot3D(axes.c2p(2, 1, 1.5), color=RED, radius=0.1)
        self.play(Create(point))

        # Camera rotation
        self.begin_ambient_camera_rotation(rate=0.2)
        self.wait(5)


class ParametricSurfaceExample(ThreeDScene):
    """3D parametric surface visualization."""

    def construct(self):
        self.set_camera_orientation(phi=60 * DEGREES, theta=-60 * DEGREES)

        axes = ThreeDAxes(
            x_range=[-3, 3],
            y_range=[-3, 3],
            z_range=[-2, 2],
        )

        # Saddle surface: z = x^2 - y^2
        surface = Surface(
            lambda u, v: axes.c2p(u, v, u ** 2 - v ** 2),
            u_range=[-2, 2],
            v_range=[-2, 2],
            resolution=(20, 20),
            fill_opacity=0.7,
        )
        surface.set_color_by_gradient(BLUE, GREEN, YELLOW)

        self.play(Create(axes))
        self.play(Create(surface), run_time=2)

        self.begin_ambient_camera_rotation(rate=0.15)
        self.wait(5)


class SphereVisualization(ThreeDScene):
    """Sphere with parametric representation."""

    def construct(self):
        self.set_camera_orientation(phi=70 * DEGREES, theta=30 * DEGREES)

        # Parametric sphere
        sphere = Surface(
            lambda u, v: np.array([
                np.cos(v) * np.sin(u),
                np.sin(v) * np.sin(u),
                np.cos(u)
            ]),
            u_range=[0, PI],
            v_range=[0, 2 * PI],
            resolution=(20, 40),
        )
        sphere.set_color_by_gradient(BLUE_E, BLUE, TEAL)

        self.play(Create(sphere), run_time=2)

        # Animate camera
        self.begin_ambient_camera_rotation(rate=0.2)
        self.wait(5)


class Function3DPlot(ThreeDScene):
    """Plotting z = f(x, y) surfaces."""

    def construct(self):
        self.set_camera_orientation(phi=65 * DEGREES, theta=-45 * DEGREES)

        axes = ThreeDAxes(
            x_range=[-3, 3],
            y_range=[-3, 3],
            z_range=[-1, 1],
        )

        # Sine wave surface
        surface = Surface(
            lambda u, v: axes.c2p(
                u, v,
                np.sin(np.sqrt(u ** 2 + v ** 2))
            ),
            u_range=[-3, 3],
            v_range=[-3, 3],
            resolution=(30, 30),
        )
        surface.set_color_by_gradient(PURPLE, RED, ORANGE)

        self.play(Create(axes))
        self.play(Create(surface), run_time=2)

        self.begin_ambient_camera_rotation(rate=0.1)
        self.wait(6)


class VectorField3D(ThreeDScene):
    """3D vector field visualization."""

    def construct(self):
        self.set_camera_orientation(phi=60 * DEGREES, theta=-45 * DEGREES)

        axes = ThreeDAxes(
            x_range=[-3, 3],
            y_range=[-3, 3],
            z_range=[-3, 3],
        )

        # Create arrows representing a vector field
        arrows = VGroup()
        for x in np.arange(-2, 3, 1):
            for y in np.arange(-2, 3, 1):
                for z in np.arange(-2, 3, 1):
                    # Vector field: F = (-y, x, z)
                    start = axes.c2p(x, y, z)
                    direction = np.array([-y, x, z]) * 0.3
                    end = start + direction

                    arrow = Arrow3D(
                        start=start,
                        end=end,
                        color=interpolate_color(
                            BLUE, RED,
                            (z + 2) / 4
                        ),
                    )
                    arrows.add(arrow)

        self.play(Create(axes))
        self.play(LaggedStart(*[Create(a) for a in arrows], lag_ratio=0.02))

        self.begin_ambient_camera_rotation(rate=0.1)
        self.wait(5)


class CameraMovement3D(ThreeDScene):
    """Demonstrating 3D camera controls."""

    def construct(self):
        # Start with a default view
        self.set_camera_orientation(phi=75 * DEGREES, theta=-45 * DEGREES)

        # Create a 3D object
        torus = Torus(
            major_radius=2,
            minor_radius=0.5,
            color=BLUE,
            fill_opacity=0.8
        )

        self.play(Create(torus))
        self.wait()

        # Move camera to different angles
        self.move_camera(phi=30 * DEGREES, theta=0, run_time=2)
        self.wait()

        self.move_camera(phi=90 * DEGREES, theta=90 * DEGREES, run_time=2)
        self.wait()

        # Zoom by adjusting frame
        self.move_camera(zoom=1.5, run_time=1)
        self.wait()

        self.move_camera(zoom=0.7, run_time=1)
        self.wait()


class Line3DExample(ThreeDScene):
    """3D lines and curves."""

    def construct(self):
        self.set_camera_orientation(phi=70 * DEGREES, theta=-45 * DEGREES)

        axes = ThreeDAxes()

        # 3D helix
        helix = ParametricFunction(
            lambda t: np.array([
                np.cos(t),
                np.sin(t),
                t / 4
            ]),
            t_range=[0, 4 * PI],
            color=YELLOW,
        )

        # Line in 3D
        line = Line3D(
            start=axes.c2p(-2, -2, -1),
            end=axes.c2p(2, 2, 1),
            color=RED,
        )

        self.play(Create(axes))
        self.play(Create(helix), run_time=2)
        self.play(Create(line))

        self.begin_ambient_camera_rotation(rate=0.15)
        self.wait(5)


class TextIn3D(ThreeDScene):
    """Text and math in 3D scenes."""

    def construct(self):
        self.set_camera_orientation(phi=60 * DEGREES, theta=-45 * DEGREES)

        axes = ThreeDAxes()

        # 3D text (stays fixed to camera)
        title = Text("3D Visualization", font_size=48)
        title.to_corner(UL)
        self.add_fixed_in_frame_mobjects(title)

        # Math label fixed to camera
        equation = MathTex(r"z = x^2 + y^2")
        equation.to_corner(UR)
        self.add_fixed_in_frame_mobjects(equation)

        # Surface
        paraboloid = Surface(
            lambda u, v: axes.c2p(u, v, u ** 2 + v ** 2),
            u_range=[-1.5, 1.5],
            v_range=[-1.5, 1.5],
            resolution=(15, 15),
        )
        paraboloid.set_color_by_gradient(BLUE, GREEN)

        self.play(Write(title), Write(equation))
        self.play(Create(axes))
        self.play(Create(paraboloid))

        self.begin_ambient_camera_rotation(rate=0.1)
        self.wait(5)


class AnimatedSurface(ThreeDScene):
    """Surface that changes over time."""

    def construct(self):
        self.set_camera_orientation(phi=65 * DEGREES, theta=-45 * DEGREES)

        axes = ThreeDAxes(
            x_range=[-3, 3],
            y_range=[-3, 3],
            z_range=[-2, 2],
        )

        # Time parameter
        time = ValueTracker(0)

        # Animated wave surface
        surface = always_redraw(
            lambda: Surface(
                lambda u, v: axes.c2p(
                    u, v,
                    np.sin(np.sqrt(u ** 2 + v ** 2) - time.get_value())
                ),
                u_range=[-3, 3],
                v_range=[-3, 3],
                resolution=(25, 25),
            ).set_color_by_gradient(BLUE, TEAL)
        )

        self.add(axes, surface)

        # Animate
        self.play(
            time.animate.set_value(4 * PI),
            run_time=8,
            rate_func=linear
        )


class MultipleObjects3D(ThreeDScene):
    """Combining multiple 3D objects."""

    def construct(self):
        self.set_camera_orientation(phi=60 * DEGREES, theta=-30 * DEGREES)

        # Create various 3D shapes
        sphere = Sphere(radius=0.5, color=RED).shift(LEFT * 2 + UP)
        cube = Cube(side_length=0.8, color=BLUE).shift(RIGHT * 2)
        cylinder = Cylinder(
            radius=0.4,
            height=1.2,
            color=GREEN
        ).shift(DOWN + LEFT)

        # Arrows connecting them
        arrow1 = Arrow3D(
            start=sphere.get_center(),
            end=cube.get_center(),
            color=YELLOW
        )
        arrow2 = Arrow3D(
            start=cube.get_center(),
            end=cylinder.get_center(),
            color=YELLOW
        )

        self.play(
            Create(sphere),
            Create(cube),
            Create(cylinder),
        )
        self.play(Create(arrow1), Create(arrow2))

        self.begin_ambient_camera_rotation(rate=0.2)
        self.wait(5)

# Attention Visualization Package
# Converted from 3b1b ManimGL to ManimCE

"""
Attention Visualization Helpers - Converted from 3b1b ManimGL to ManimCE

Original: videos/_2024/transformers/helpers.py
Contains utility functions and classes for attention visualization.
"""

from manim import *
import numpy as np
import warnings
import random
import itertools as it
from typing import Optional, Tuple


# =============================================================================
# UTILITY FUNCTIONS
# =============================================================================

def softmax(logits, temperature=1.0):
    """Numerically stable softmax function."""
    logits = np.array(logits)
    with warnings.catch_warnings():
        warnings.filterwarnings('ignore')
        logits = logits - np.max(logits)
        exps = np.exp(np.divide(logits, temperature, where=temperature != 0))

    if np.isinf(exps).any() or np.isnan(exps).any() or temperature == 0:
        result = np.zeros_like(logits)
        result[np.argmax(logits)] = 1
        return result
    return exps / np.sum(exps)


def value_to_color(
    value,
    low_positive_color=BLUE_E,
    high_positive_color=BLUE_B,
    low_negative_color=RED_E,
    high_negative_color=RED_B,
    min_value=0.0,
    max_value=10.0
):
    """Map a numeric value to a color based on sign and magnitude."""
    # Clamp alpha between 0 and 1
    alpha = max(0, min(1, abs(value - min_value) / (max_value - min_value))) if max_value != min_value else 0.5

    if value >= 0:
        return interpolate_color(low_positive_color, high_positive_color, alpha)
    else:
        return interpolate_color(low_negative_color, high_negative_color, alpha)


def get_paragraph(words, line_len=40, font_size=48):
    """Handle word wrapping for text display."""
    words = list(map(str.strip, words))
    word_lens = list(map(len, words))
    lines = []
    lh, rh = 0, 0
    while rh < len(words):
        rh += 1
        if sum(word_lens[lh:rh]) > line_len:
            rh -= 1
            lines.append(words[lh:rh])
            lh = rh
    lines.append(words[lh:])
    text = "\n".join([" ".join(line).strip() for line in lines])
    return Text(text, font_size=font_size)


def random_bright_color(hue_range=(0.0, 1.0)):
    """Generate a random bright color within a hue range."""
    import colorsys
    hue = random.uniform(*hue_range)
    rgb = colorsys.hsv_to_rgb(hue, 0.7, 0.9)
    return rgb_to_color(rgb)


# =============================================================================
# CUSTOM MOBJECT CLASSES
# =============================================================================

class NumericEmbedding(VGroup):
    """
    A vertical vector of decimal numbers representing an embedding.
    Displays values with color coding based on magnitude.
    """

    def __init__(
        self,
        values: Optional[np.ndarray] = None,
        length: int = 7,
        num_decimal_places: int = 1,
        value_range: Tuple[float, float] = (-9.9, 9.9),
        show_ellipsis: bool = True,
        ellipsis_row: int = -2,
        dark_color=GREY_C,
        light_color=WHITE,
        bracket_color=GREY_B,
        **kwargs
    ):
        super().__init__(**kwargs)

        self.value_range = value_range
        self.dark_color = dark_color
        self.light_color = light_color

        if values is None:
            values = np.random.uniform(*value_range, size=length)

        self.values = values
        self.length = len(values)

        # Create decimal number entries
        self.elements = VGroup()
        for i, val in enumerate(values):
            if show_ellipsis and i == (ellipsis_row % len(values)):
                entry = MathTex(r"\vdots")
            else:
                entry = DecimalNumber(
                    val,
                    num_decimal_places=num_decimal_places,
                    include_sign=True,
                    font_size=36
                )
                # Color based on value
                alpha = abs(val) / max(abs(value_range[0]), abs(value_range[1]))
                entry.set_color(interpolate_color(dark_color, light_color, alpha))
            self.elements.add(entry)

        self.elements.arrange(DOWN, buff=0.15)

        # Add brackets
        self.left_bracket = MathTex(r"\left[")
        self.right_bracket = MathTex(r"\right]")

        self.left_bracket.stretch_to_fit_height(self.elements.get_height() * 1.1)
        self.right_bracket.stretch_to_fit_height(self.elements.get_height() * 1.1)

        self.left_bracket.next_to(self.elements, LEFT, buff=0.1)
        self.right_bracket.next_to(self.elements, RIGHT, buff=0.1)

        self.left_bracket.set_color(bracket_color)
        self.right_bracket.set_color(bracket_color)

        self.add(self.left_bracket, self.elements, self.right_bracket)

    def get_brackets(self):
        return VGroup(self.left_bracket, self.right_bracket)

    def get_entries(self):
        return self.elements


class WeightMatrix(VGroup):
    """
    A matrix of decimal numbers with color-coded entries.
    Used to represent weight matrices in neural networks.
    """

    def __init__(
        self,
        values: Optional[np.ndarray] = None,
        shape: Tuple[int, int] = (6, 8),
        value_range: Tuple[float, float] = (-9.9, 9.9),
        num_decimal_places: int = 1,
        show_ellipsis: bool = True,
        ellipsis_row: int = -2,
        ellipsis_col: int = -2,
        low_positive_color=BLUE_E,
        high_positive_color=BLUE_B,
        low_negative_color=RED_E,
        high_negative_color=RED_B,
        **kwargs
    ):
        super().__init__(**kwargs)

        self.shape = shape
        self.value_range = value_range
        self.low_positive_color = low_positive_color
        self.high_positive_color = high_positive_color
        self.low_negative_color = low_negative_color
        self.high_negative_color = high_negative_color

        if values is None:
            values = np.random.uniform(*value_range, size=shape)

        self.values = values

        # Create matrix entries
        self.rows = VGroup()
        n_rows, n_cols = shape

        for i in range(n_rows):
            row = VGroup()
            for j in range(n_cols):
                if show_ellipsis and i == (ellipsis_row % n_rows):
                    entry = MathTex(r"\vdots")
                elif show_ellipsis and j == (ellipsis_col % n_cols):
                    entry = MathTex(r"\cdots")
                else:
                    val = values[i, j]
                    entry = DecimalNumber(
                        val,
                        num_decimal_places=num_decimal_places,
                        include_sign=True,
                        font_size=24
                    )
                    entry.set_color(value_to_color(
                        val,
                        low_positive_color,
                        high_positive_color,
                        low_negative_color,
                        high_negative_color,
                        0, max(abs(value_range[0]), abs(value_range[1]))
                    ))
                row.add(entry)
            row.arrange(RIGHT, buff=0.2)
            self.rows.add(row)

        self.rows.arrange(DOWN, buff=0.15)

        # Add brackets
        self.left_bracket = MathTex(r"\left[")
        self.right_bracket = MathTex(r"\right]")

        self.left_bracket.stretch_to_fit_height(self.rows.get_height() * 1.1)
        self.right_bracket.stretch_to_fit_height(self.rows.get_height() * 1.1)

        self.left_bracket.next_to(self.rows, LEFT, buff=0.1)
        self.right_bracket.next_to(self.rows, RIGHT, buff=0.1)

        self.add(self.left_bracket, self.rows, self.right_bracket)

    def get_entries(self):
        entries = VGroup()
        for row in self.rows:
            for entry in row:
                if isinstance(entry, DecimalNumber):
                    entries.add(entry)
        return entries

    def get_rows(self):
        return self.rows


class ContextAnimation(LaggedStart):
    """
    Animation showing context flow from source words to target word.
    Creates arcing lines that flash from sources to target.
    """

    def __init__(
        self,
        target,
        sources,
        direction=UP,
        time_width=2,
        min_stroke_width=1,
        max_stroke_width=5,
        strengths=None,
        run_time=3,
        path_arc=PI / 2,
        **kwargs,
    ):
        arcs = VGroup()
        if strengths is None:
            strengths = np.random.random(len(sources)) ** 2

        for source, strength in zip(sources, strengths):
            sign = direction[1] * (-1) ** int(source.get_x() < target.get_x())
            arc = Line(
                source.get_edge_center(direction),
                target.get_edge_center(direction),
                path_arc=sign * path_arc,
            )
            arc.set_stroke(
                color=random_bright_color(hue_range=(0.1, 0.3)),
                width=interpolate(min_stroke_width, max_stroke_width, strength)
            )
            arcs.add(arc)

        arcs.shuffle()
        lag_ratio = 0.5 / max(len(arcs), 1)

        super().__init__(
            *[
                ShowPassingFlash(arc, time_width=time_width)
                for arc in arcs
            ],
            lag_ratio=lag_ratio,
            run_time=run_time,
            **kwargs,
        )


class NeuralNetwork(VGroup):
    """
    Visual representation of a neural network with layers and connections.
    """

    def __init__(
        self,
        layer_sizes=[6, 12, 6],
        neuron_radius=0.1,
        v_buff=0.3,
        h_buff=1.5,
        max_stroke_width=2.0,
        **kwargs
    ):
        super().__init__(**kwargs)

        self.max_stroke_width = max_stroke_width

        # Create layers
        self.layers = VGroup()
        for n in layer_sizes:
            layer = VGroup(*[
                Circle(radius=neuron_radius, color=WHITE, fill_opacity=random.random())
                for _ in range(n)
            ])
            layer.arrange(DOWN, buff=v_buff)
            self.layers.add(layer)

        self.layers.arrange(RIGHT, buff=h_buff)

        # Create connections
        self.lines = VGroup()
        for l1, l2 in zip(self.layers, self.layers[1:]):
            layer_lines = VGroup()
            for n1 in l1:
                for n2 in l2:
                    line = Line(
                        n1.get_center(),
                        n2.get_center(),
                        buff=neuron_radius
                    )
                    line.set_stroke(
                        color=value_to_color(random.uniform(-10, 10)),
                        width=max_stroke_width * random.random(),
                        opacity=random.random() ** 2
                    )
                    layer_lines.add(line)
            self.lines.add(layer_lines)

        self.add(self.lines, self.layers)


class AttentionPattern(VGroup):
    """
    Visual representation of attention weights between tokens.
    Shows which tokens attend to which with varying line widths.
    """

    def __init__(
        self,
        n_tokens=8,
        token_labels=None,
        attention_weights=None,
        **kwargs
    ):
        super().__init__(**kwargs)

        if token_labels is None:
            token_labels = [f"T{i}" for i in range(n_tokens)]

        if attention_weights is None:
            # Random attention pattern
            attention_weights = softmax(np.random.randn(n_tokens, n_tokens), temperature=0.5)

        # Create token representations
        self.tokens = VGroup()
        for label in token_labels:
            token = VGroup(
                Square(side_length=0.8, color=BLUE, fill_opacity=0.3),
                Text(label, font_size=24)
            )
            token[1].move_to(token[0])
            self.tokens.add(token)

        self.tokens.arrange(RIGHT, buff=0.5)

        # Create attention lines (simplified - just showing strongest connections)
        self.attention_lines = VGroup()
        for i in range(n_tokens):
            for j in range(n_tokens):
                if attention_weights[i, j] > 0.1:  # Threshold
                    line = Line(
                        self.tokens[i].get_bottom(),
                        self.tokens[j].get_bottom(),
                        path_arc=-0.5,
                    )
                    line.set_stroke(
                        color=YELLOW,
                        width=attention_weights[i, j] * 5,
                        opacity=attention_weights[i, j]
                    )
                    self.attention_lines.add(line)

        self.add(self.tokens, self.attention_lines)


# =============================================================================
# ANIMATION HELPERS
# =============================================================================

class RandomizeMatrixEntries(Animation):
    """Animation that smoothly randomizes matrix entries."""

    def __init__(self, matrix, **kwargs):
        self.matrix = matrix
        self.entries = matrix.get_entries()
        self.start_values = [
            entry.get_value() if hasattr(entry, 'get_value') else 0
            for entry in self.entries
        ]
        self.target_values = np.random.uniform(
            matrix.value_range[0],
            matrix.value_range[1],
            len(self.entries)
        )
        super().__init__(matrix, **kwargs)

    def interpolate_mobject(self, alpha: float) -> None:
        for index, entry in enumerate(self.entries):
            if hasattr(entry, 'set_value'):
                start = self.start_values[index]
                target = self.target_values[index]
                entry.set_value(interpolate(start, target, alpha))


def show_attention_flow(scene, source_mobs, target_mob, weights=None, run_time=2):
    """Helper to animate attention flow from multiple sources to a target."""
    if weights is None:
        weights = np.random.random(len(source_mobs))
        weights = weights / weights.sum()

    arrows = VGroup()
    for source, weight in zip(source_mobs, weights):
        arrow = CurvedArrow(
            source.get_top(),
            target_mob.get_top(),
            angle=-TAU/4
        )
        arrow.set_stroke(width=weight * 5, color=YELLOW)
        arrow.set_opacity(weight)
        arrows.add(arrow)

    scene.play(
        LaggedStart(*[Create(a) for a in arrows], lag_ratio=0.2),
        run_time=run_time
    )
    return arrows

"""
Attention Mechanism Visualization - Converted from 3b1b ManimGL to ManimCE

Original: videos/_2024/transformers/attention.py
Demonstrates the attention mechanism used in transformers.

Run with: manim -pql scenes.py SceneName
"""

from manim import *
import numpy as np
import sys
from pathlib import Path

# Add parent directory to path for helpers import
sys.path.insert(0, str(Path(__file__).parent))
from helpers import (
    NumericEmbedding, WeightMatrix, ContextAnimation,
    NeuralNetwork, AttentionPattern, softmax, value_to_color,
    random_bright_color, show_attention_flow
)


class AttentionPatterns(Scene):
    """
    Demonstrates how attention allows words to influence each other.
    Shows adjectives modifying nouns through attention connections.
    """

    def construct(self):
        # Add sentence
        phrase = "a fluffy blue creature roamed the verdant forest"
        phrase_mob = Text(phrase, font_size=36)
        phrase_mob.move_to(2 * UP)

        words = phrase.split()
        word_mobs = VGroup()

        # Create individual word mobjects
        current_x = phrase_mob.get_left()[0]
        for word in words:
            # Find the word in the phrase
            word_mob = Text(word, font_size=36)
            word_mobs.add(word_mob)

        word_mobs.arrange(RIGHT, buff=0.3)
        word_mobs.move_to(2 * UP)

        self.play(
            LaggedStart(*[FadeIn(w, shift=0.5 * UP) for w in word_mobs], lag_ratio=0.15)
        )
        self.wait()

        # Create word rectangles
        word_rects = VGroup()
        for word_mob in word_mobs:
            rect = SurroundingRectangle(word_mob, buff=0.1)
            rect.set_stroke(GREY, 2)
            rect.set_fill(GREY, 0.2)
            word_rects.add(rect)

        # Identify adjectives and nouns
        adj_indices = [1, 2, 6]  # fluffy, blue, verdant
        noun_indices = [3, 7]     # creature, forest

        adj_rects = VGroup(*[word_rects[i] for i in adj_indices])
        noun_rects = VGroup(*[word_rects[i] for i in noun_indices])

        adj_mobs = VGroup(*[word_mobs[i] for i in adj_indices])
        noun_mobs = VGroup(*[word_mobs[i] for i in noun_indices])

        # Color the rectangles
        adj_rects[0].set_fill(BLUE_C, 0.3)
        adj_rects[1].set_fill(BLUE_D, 0.3)
        adj_rects[2].set_fill(GREEN, 0.3)
        noun_rects.set_fill(GREY_BROWN, 0.3)

        self.play(
            LaggedStart(*[DrawBorderThenFill(r) for r in adj_rects], lag_ratio=0.2),
        )
        self.wait()

        # Show arrows from adjectives to nouns
        adj_arrows = VGroup(
            CurvedArrow(adj_mobs[0].get_top(), noun_mobs[0].get_top(), angle=-0.5),
            CurvedArrow(adj_mobs[1].get_top(), noun_mobs[0].get_top(), angle=-0.5),
            CurvedArrow(adj_mobs[2].get_top(), noun_mobs[1].get_top(), angle=-0.5),
        )
        adj_arrows.set_color(GREY_B)

        self.play(
            LaggedStart(*[DrawBorderThenFill(r) for r in noun_rects], lag_ratio=0.2),
            LaggedStart(*[Create(a) for a in adj_arrows], lag_ratio=0.2),
        )
        self.wait()

        # Animate context flow
        self.play(
            ContextAnimation(noun_mobs[0], adj_mobs[:2], strengths=[1, 1]),
            ContextAnimation(noun_mobs[1], adj_mobs[2:], strengths=[1]),
        )
        self.wait()

        # Show embeddings
        all_rects = word_rects.copy()
        embeddings = VGroup()
        for rect in all_rects:
            emb = NumericEmbedding(length=8)
            emb.set_width(0.4)
            emb.next_to(rect, DOWN, buff=1.2)
            embeddings.add(emb)

        emb_arrows = VGroup()
        for rect, emb in zip(all_rects, embeddings):
            arrow = Arrow(rect.get_bottom(), emb.get_top(), buff=0.1)
            emb_arrows.add(arrow)

        self.play(
            FadeIn(word_rects),
            LaggedStart(*[GrowArrow(a) for a in emb_arrows], lag_ratio=0.1),
            LaggedStart(*[FadeIn(e, shift=0.5 * DOWN) for e in embeddings], lag_ratio=0.1),
            FadeOut(adj_arrows)
        )
        self.wait()

        # Show embedding dimension
        brace = Brace(embeddings[0], LEFT, buff=SMALL_BUFF)
        dim_label = Text("12,288", font_size=24, color=YELLOW)
        dim_label.next_to(brace, LEFT)

        self.play(
            GrowFromCenter(brace),
            FadeIn(dim_label)
        )
        self.wait(2)


class QueryKeyValueExplanation(Scene):
    """
    Explains the Query, Key, Value mechanism in attention.
    """

    def construct(self):
        # Title
        title = Text("Query, Key, Value", font_size=48)
        title.to_edge(UP)
        self.play(Write(title))
        self.wait()

        # Create three matrices
        q_matrix = WeightMatrix(shape=(4, 4), value_range=(-5, 5))
        k_matrix = WeightMatrix(shape=(4, 4), value_range=(-5, 5))
        v_matrix = WeightMatrix(shape=(4, 4), value_range=(-5, 5))

        matrices = VGroup(q_matrix, k_matrix, v_matrix)
        matrices.arrange(RIGHT, buff=1)
        matrices.set_height(2)
        matrices.next_to(title, DOWN, buff=0.8)

        # Labels
        q_label = Text("Query (Q)", font_size=30, color=BLUE)
        k_label = Text("Key (K)", font_size=30, color=GREEN)
        v_label = Text("Value (V)", font_size=30, color=RED)

        q_label.next_to(q_matrix, UP)
        k_label.next_to(k_matrix, UP)
        v_label.next_to(v_matrix, UP)

        self.play(
            FadeIn(q_matrix),
            Write(q_label),
        )
        self.wait()

        self.play(
            FadeIn(k_matrix),
            Write(k_label),
        )
        self.wait()

        self.play(
            FadeIn(v_matrix),
            Write(v_label),
        )
        self.wait()

        # Explain the formula
        formula = MathTex(
            r"\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V",
            font_size=36
        )
        formula.next_to(matrices, DOWN, buff=1)

        self.play(Write(formula))
        self.wait(2)

        # Highlight different parts
        # Q*K^T computes similarity
        explanation1 = Text("Q·K^T → measures similarity between queries and keys", font_size=24)
        explanation1.next_to(formula, DOWN, buff=0.5)

        self.play(Write(explanation1))
        self.wait()

        # Softmax normalizes
        explanation2 = Text("Softmax → converts to attention weights (probabilities)", font_size=24)
        explanation2.next_to(explanation1, DOWN, buff=0.3)

        self.play(Write(explanation2))
        self.wait()

        # Multiply by V
        explanation3 = Text("× V → weighted sum of values", font_size=24)
        explanation3.next_to(explanation2, DOWN, buff=0.3)

        self.play(Write(explanation3))
        self.wait(2)


class AttentionMatrixVisualization(Scene):
    """
    Shows how attention scores form a matrix pattern.
    """

    def construct(self):
        # Create tokens
        tokens = ["The", "cat", "sat", "on", "the", "mat"]
        n = len(tokens)

        # Token labels on left (queries)
        query_labels = VGroup(*[Text(t, font_size=24) for t in tokens])
        query_labels.arrange(DOWN, buff=0.4)
        query_labels.shift(LEFT * 4)

        # Token labels on top (keys)
        key_labels = VGroup(*[Text(t, font_size=24) for t in tokens])
        key_labels.arrange(RIGHT, buff=0.4)
        key_labels.next_to(query_labels, RIGHT, buff=1)
        key_labels.shift(UP * 2)

        # Create attention grid
        attention_scores = softmax(np.random.randn(n, n) * 2, temperature=0.3)

        grid = VGroup()
        for i in range(n):
            row = VGroup()
            for j in range(n):
                score = attention_scores[i, j]
                cell = Square(side_length=0.5)
                cell.set_fill(
                    color=interpolate_color(BLACK, YELLOW, score),
                    opacity=0.8
                )
                cell.set_stroke(WHITE, 0.5)
                row.add(cell)
            row.arrange(RIGHT, buff=0)
            grid.add(row)

        grid.arrange(DOWN, buff=0)
        grid.next_to(query_labels, RIGHT, buff=0.5)
        grid.align_to(query_labels, UP)

        # Adjust key labels position
        key_labels.move_to(grid.get_top() + UP * 0.5)
        key_labels.align_to(grid, LEFT)

        # Title
        title = Text("Attention Matrix", font_size=36)
        title.to_edge(UP)

        self.play(Write(title))
        self.play(
            LaggedStart(*[FadeIn(l) for l in query_labels], lag_ratio=0.1),
            LaggedStart(*[FadeIn(l) for l in key_labels], lag_ratio=0.1),
        )
        self.wait()

        # Animate grid appearing
        all_cells = VGroup(*[cell for row in grid for cell in row])
        self.play(
            LaggedStart(*[FadeIn(c, scale=0.5) for c in all_cells], lag_ratio=0.02)
        )
        self.wait()

        # Highlight a row (how "cat" attends to all words)
        highlight_row = 1  # "cat"
        row_highlight = SurroundingRectangle(grid[highlight_row], color=BLUE, buff=0.05)

        explanation = Text(
            '"cat" attends mostly to itself and "sat"',
            font_size=24
        )
        explanation.next_to(grid, DOWN, buff=0.5)

        self.play(Create(row_highlight), Write(explanation))
        self.wait(2)


class MultiHeadedAttention(ThreeDScene):
    """
    Explains multi-head attention mechanism with 3D visualization.
    Shows multiple attention heads arranged in depth with camera rotation.
    Inspired by 3b1b's transformer visualization.
    """

    def construct(self):
        # Title animation: Single head -> Multi-head
        single_title = Text("Single head of attention", font_size=42)
        multiple_title = Text("Multi-headed attention", font_size=42)

        for title in [single_title, multiple_title]:
            title.to_edge(UP)

        self.play(Write(single_title))
        self.wait(0.5)

        # Flash around "head"
        head_text = single_title[7:11]  # "head"
        self.play(
            Indicate(head_text, color=YELLOW, scale_factor=1.2),
            head_text.animate.set_color(YELLOW),
        )
        self.wait(0.5)

        # Transform title
        self.play(
            TransformMatchingShapes(single_title, multiple_title),
            run_time=1.5
        )
        self.wait()

        # Create attention pattern visualization (grid with dots)
        def create_attention_pattern(n_rows=8, seed=None):
            """Create a grid visualization of attention weights."""
            if seed is not None:
                np.random.seed(seed)

            # Create the base grid
            grid = VGroup()
            cell_size = 0.4
            for i in range(n_rows):
                for j in range(n_rows):
                    cell = Square(side_length=cell_size)
                    cell.set_stroke(WHITE, 0.5, opacity=0.3)
                    cell.move_to(np.array([j * cell_size, -i * cell_size, 0]))
                    grid.add(cell)

            grid.center()

            # Generate causal attention pattern (lower triangular)
            pattern = np.random.normal(0, 1, (n_rows, n_rows))
            for n in range(n_rows):
                pattern[:, n][n + 1:] = -np.inf  # Mask future tokens
                exp_vals = np.exp(pattern[:, n] - np.max(pattern[:, n][pattern[:, n] > -np.inf]))
                pattern[:, n] = exp_vals / np.sum(exp_vals[exp_vals < np.inf])
            pattern = np.nan_to_num(pattern, nan=0.0, posinf=0.0, neginf=0.0)

            # Add dots based on attention weights
            dots = VGroup()
            for i in range(n_rows):
                for j in range(n_rows):
                    value = pattern[i, j]
                    if value > 0.05:  # Threshold for visibility
                        dot = Dot(
                            radius=cell_size * 0.4 * value,
                            color=GREY_B,
                            fill_opacity=0.8
                        )
                        dot.move_to(grid[i * n_rows + j].get_center())
                        dots.add(dot)

            # Create border rectangle
            border = SurroundingRectangle(grid, buff=0.05)
            border.set_stroke(WHITE, 2)
            border.set_fill(BLACK, 0.9)

            pattern_mob = VGroup(border, grid, dots)
            return pattern_mob

        # Create multiple attention heads
        n_heads = 12
        heads = VGroup()
        for i in range(n_heads):
            head = create_attention_pattern(n_rows=6, seed=i * 42)
            head.set_height(2.5)
            heads.add(head)

        # Arrange in 3D depth (along z-axis)
        for i, head in enumerate(heads):
            head.shift(OUT * i * 0.5)  # Stack in z direction

        heads.center()
        heads.shift(DOWN * 0.5)

        # Show first head (screen rectangle style)
        first_head = heads[-1].copy()
        first_head.move_to(ORIGIN + DOWN * 0.5)
        first_head.shift(IN * (n_heads - 1) * 0.25)  # Reset z position

        self.play(FadeIn(first_head))
        self.wait()

        # Add fixed-in-frame elements
        self.add_fixed_in_frame_mobjects(multiple_title)

        # Rotate camera to reveal depth
        self.move_camera(
            phi=70 * DEGREES,
            theta=-60 * DEGREES,
            run_time=2
        )

        # Fan out the heads from the first one
        self.play(
            LaggedStart(
                *[FadeIn(head, shift=OUT * 0.3) for head in heads[:-1]],
                lag_ratio=0.15
            ),
            FadeOut(first_head),
            run_time=3
        )
        self.add(heads)
        self.wait()

        # Add matrix labels for each head (W_Q, W_K)
        wq_labels = VGroup()
        wk_labels = VGroup()

        colors = [YELLOW, TEAL]
        n_shown = min(5, n_heads)

        for i, head in enumerate(list(heads)[-n_shown:]):
            head_num = n_heads - n_shown + i + 1
            wq = MathTex(f"W_Q^{{({head_num})}}", font_size=28, color=YELLOW)
            wk = MathTex(f"W_K^{{({head_num})}}", font_size=28, color=TEAL)

            # Position above each head
            wq.next_to(head, UP, buff=0.2)
            wq.shift(LEFT * 0.3)
            wk.next_to(head, UP, buff=0.2)
            wk.shift(RIGHT * 0.3)

            # Rotate to face camera
            for label in [wq, wk]:
                label.rotate(70 * DEGREES, axis=RIGHT)
                label.rotate(-60 * DEGREES, axis=OUT)

            wq_labels.add(wq)
            wk_labels.add(wk)

        # Add dots to indicate more heads
        dots_label = MathTex(r"\cdots", font_size=48, color=WHITE)
        dots_label.next_to(heads[0], OUT, buff=0.5)
        dots_label.rotate(70 * DEGREES, axis=RIGHT)
        dots_label.rotate(-60 * DEGREES, axis=OUT)

        self.play(
            LaggedStart(*[FadeIn(wq, shift=UP * 0.2) for wq in wq_labels], lag_ratio=0.2),
            run_time=1.5
        )
        self.play(
            LaggedStart(*[FadeIn(wk, shift=UP * 0.2) for wk in wk_labels], lag_ratio=0.2),
            FadeIn(dots_label),
            run_time=1.5
        )
        self.wait()

        # Add brace showing "96 heads" (scaled down for demonstration)
        brace_text = Text("96 heads", font_size=36, color=WHITE)
        brace_text.rotate(70 * DEGREES, axis=RIGHT)
        brace_text.rotate(-60 * DEGREES, axis=OUT)
        brace_text.next_to(heads, UP, buff=0.8)
        brace_text.shift(LEFT * 2)

        self.play(FadeIn(brace_text, shift=UP * 0.3))
        self.wait()

        # Rotate camera to show different angle
        self.move_camera(
            phi=60 * DEGREES,
            theta=-80 * DEGREES,
            run_time=2
        )
        self.wait()

        # Explanation text (fixed in frame)
        explanation = VGroup(
            Text("Each head learns different patterns:", font_size=24),
            Text("• Syntactic relationships", font_size=20, color=BLUE),
            Text("• Semantic connections", font_size=20, color=GREEN),
            Text("• Positional patterns", font_size=20, color=YELLOW),
        )
        explanation.arrange(DOWN, aligned_edge=LEFT, buff=0.15)
        explanation.to_corner(DL, buff=0.5)

        self.add_fixed_in_frame_mobjects(explanation)
        self.play(
            LaggedStart(*[Write(e) for e in explanation], lag_ratio=0.3)
        )
        self.wait()

        # Return to front view
        self.move_camera(
            phi=0,
            theta=-90 * DEGREES,
            run_time=2
        )
        self.wait()

        # Show concatenation concept
        concat_text = Text("Concatenate outputs from all heads", font_size=28)
        concat_text.to_edge(DOWN, buff=0.5)
        self.add_fixed_in_frame_mobjects(concat_text)
        self.play(Write(concat_text))
        self.wait()

        # Final hold
        self.wait()
        self.wait(2)


class SelfAttentionDemo(Scene):
    """
    Interactive demonstration of self-attention on a simple sentence.
    """

    def construct(self):
        # Title
        title = Text("Self-Attention in Action", font_size=42)
        title.to_edge(UP)
        self.play(Write(title))

        # Create sentence
        sentence = "The quick brown fox"
        words = sentence.split()

        word_boxes = VGroup()
        for word in words:
            box = VGroup(
                RoundedRectangle(
                    width=1.5, height=0.8,
                    corner_radius=0.1,
                    fill_opacity=0.3,
                    fill_color=BLUE,
                    stroke_color=WHITE
                ),
                Text(word, font_size=28)
            )
            box[1].move_to(box[0])
            word_boxes.add(box)

        word_boxes.arrange(RIGHT, buff=0.5)
        word_boxes.next_to(title, DOWN, buff=1)

        self.play(
            LaggedStart(*[FadeIn(b, scale=0.8) for b in word_boxes], lag_ratio=0.2)
        )
        self.wait()

        # Show attention from "fox" to other words
        target_idx = 3  # "fox"
        attention_weights = [0.1, 0.3, 0.4, 0.2]  # Attention weights

        # Highlight target
        target_box = word_boxes[target_idx]
        target_highlight = SurroundingRectangle(target_box, color=YELLOW, buff=0.1)

        self.play(Create(target_highlight))

        # Create attention arrows
        attention_arrows = VGroup()
        weight_labels = VGroup()

        for i, (box, weight) in enumerate(zip(word_boxes, attention_weights)):
            if i != target_idx:
                arrow = CurvedArrow(
                    box.get_bottom() + DOWN * 0.1,
                    target_box.get_bottom() + DOWN * 0.1,
                    angle=0.5 if i < target_idx else -0.5
                )
                arrow.set_stroke(
                    color=interpolate_color(GREY, YELLOW, weight),
                    width=weight * 8
                )
                attention_arrows.add(arrow)

                label = DecimalNumber(weight, num_decimal_places=1, font_size=20)
                label.next_to(arrow.point_from_proportion(0.5), DOWN, buff=0.1)
                weight_labels.add(label)

        self.play(
            LaggedStart(*[Create(a) for a in attention_arrows], lag_ratio=0.2),
            LaggedStart(*[FadeIn(l) for l in weight_labels], lag_ratio=0.2),
        )
        self.wait()

        # Show weighted combination
        result_text = Text(
            '"fox" = 0.1×"The" + 0.3×"quick" + 0.4×"brown" + 0.2×"fox"',
            font_size=24
        )
        result_text.next_to(word_boxes, DOWN, buff=1.5)

        self.play(Write(result_text))
        self.wait(2)


class ScaledDotProductAttention(Scene):
    """
    Step-by-step visualization of scaled dot-product attention.
    """

    def construct(self):
        # Title
        title = Text("Scaled Dot-Product Attention", font_size=40)
        title.to_edge(UP)
        self.play(Write(title))

        # Step 1: Show Q, K, V
        step1 = Text("Step 1: Compute Q, K, V from input", font_size=28)
        step1.next_to(title, DOWN, buff=0.5)

        q_vec = NumericEmbedding(length=4).set_height(1.5)
        k_vec = NumericEmbedding(length=4).set_height(1.5)
        v_vec = NumericEmbedding(length=4).set_height(1.5)

        vectors = VGroup(q_vec, k_vec, v_vec)
        vectors.arrange(RIGHT, buff=1)
        vectors.next_to(step1, DOWN, buff=0.5)

        q_label = Text("Q", color=BLUE, font_size=24).next_to(q_vec, UP)
        k_label = Text("K", color=GREEN, font_size=24).next_to(k_vec, UP)
        v_label = Text("V", color=RED, font_size=24).next_to(v_vec, UP)

        self.play(Write(step1))
        self.play(
            FadeIn(q_vec), FadeIn(k_vec), FadeIn(v_vec),
            Write(q_label), Write(k_label), Write(v_label)
        )
        self.wait()

        # Step 2: Compute Q·K^T
        self.play(
            FadeOut(step1),
            VGroup(vectors, q_label, k_label, v_label).animate.shift(UP)
        )

        step2 = Text("Step 2: Q · K^T (dot product)", font_size=28)
        step2.next_to(title, DOWN, buff=0.5)

        dot_product = MathTex(r"Q \cdot K^T = ", font_size=36)
        score = DecimalNumber(2.5, font_size=36, color=YELLOW)
        dot_result = VGroup(dot_product, score).arrange(RIGHT)
        dot_result.next_to(vectors, DOWN, buff=0.5)

        self.play(Write(step2))
        self.play(Write(dot_product), FadeIn(score))
        self.wait()

        # Step 3: Scale
        self.play(FadeOut(step2))
        step3 = Text("Step 3: Scale by √d_k", font_size=28)
        step3.next_to(title, DOWN, buff=0.5)

        scale_formula = MathTex(r"\frac{Q \cdot K^T}{\sqrt{d_k}} = \frac{2.5}{\sqrt{4}} = 1.25", font_size=32)
        scale_formula.next_to(dot_result, DOWN, buff=0.3)

        self.play(Write(step3))
        self.play(Write(scale_formula))
        self.wait()

        # Step 4: Softmax
        self.play(FadeOut(step3))
        step4 = Text("Step 4: Softmax → attention weights", font_size=28)
        step4.next_to(title, DOWN, buff=0.5)

        softmax_text = MathTex(r"\text{softmax}(1.25) \rightarrow \text{weights}", font_size=32)
        softmax_text.next_to(scale_formula, DOWN, buff=0.3)

        self.play(Write(step4))
        self.play(Write(softmax_text))
        self.wait()

        # Step 5: Multiply by V
        self.play(FadeOut(step4))
        step5 = Text("Step 5: Weighted sum of V", font_size=28)
        step5.next_to(title, DOWN, buff=0.5)

        final = MathTex(r"\text{Output} = \text{weights} \times V", font_size=32)
        final.next_to(softmax_text, DOWN, buff=0.3)

        self.play(Write(step5))
        self.play(Write(final))
        self.wait(2)


class PositionalEncoding(Scene):
    """
    Explains positional encoding in transformers.
    """

    def construct(self):
        title = Text("Positional Encoding", font_size=42)
        title.to_edge(UP)
        self.play(Write(title))

        # Problem statement
        problem = Text(
            "Problem: Attention has no sense of word order!",
            font_size=28, color=RED
        )
        problem.next_to(title, DOWN, buff=0.5)
        self.play(Write(problem))
        self.wait()

        # Show two sentences
        sent1 = Text('"The cat ate the fish"', font_size=24)
        sent2 = Text('"The fish ate the cat"', font_size=24)
        sents = VGroup(sent1, sent2).arrange(DOWN, buff=0.3)
        sents.next_to(problem, DOWN, buff=0.5)

        self.play(Write(sent1), Write(sent2))
        self.wait()

        # Show they have same words
        same = Text("Same words, different meanings!", font_size=24, color=YELLOW)
        same.next_to(sents, DOWN, buff=0.3)
        self.play(Write(same))
        self.wait()

        # Solution
        self.play(FadeOut(problem), FadeOut(sents), FadeOut(same))

        solution = Text(
            "Solution: Add position information to embeddings",
            font_size=28, color=GREEN
        )
        solution.next_to(title, DOWN, buff=0.5)
        self.play(Write(solution))

        # Show formula
        formula = MathTex(
            r"PE_{(pos, 2i)} &= \sin\left(\frac{pos}{10000^{2i/d}}\right) \\",
            r"PE_{(pos, 2i+1)} &= \cos\left(\frac{pos}{10000^{2i/d}}\right)",
            font_size=32
        )
        formula.next_to(solution, DOWN, buff=0.5)
        self.play(Write(formula))
        self.wait()

        # Visual representation
        positions = VGroup()
        for i in range(5):
            pos_vec = VGroup()
            for j in range(8):
                val = np.sin(i / (10000 ** (j / 8))) if j % 2 == 0 else np.cos(i / (10000 ** (j / 8)))
                cell = Square(side_length=0.3)
                cell.set_fill(interpolate_color(BLUE, RED, (val + 1) / 2), opacity=0.8)
                cell.set_stroke(WHITE, 0.5)
                pos_vec.add(cell)
            pos_vec.arrange(DOWN, buff=0)
            positions.add(pos_vec)

        positions.arrange(RIGHT, buff=0.2)
        positions.set_height(2)
        positions.next_to(formula, DOWN, buff=0.5)

        pos_labels = VGroup(*[
            Text(f"pos={i}", font_size=16).next_to(p, DOWN, buff=0.1)
            for i, p in enumerate(positions)
        ])

        self.play(
            LaggedStart(*[FadeIn(p) for p in positions], lag_ratio=0.1),
            LaggedStart(*[FadeIn(l) for l in pos_labels], lag_ratio=0.1),
        )
        self.wait(2)


# Additional simplified scenes for the key concepts

class WhatIsAttention(Scene):
    """Simple introduction to attention."""

    def construct(self):
        title = Text("What is Attention?", font_size=48)
        title.to_edge(UP)
        self.play(Write(title))

        # Key idea
        idea = Text(
            "Attention lets each word look at other words\nto understand context",
            font_size=32, line_spacing=1.5
        )
        idea.next_to(title, DOWN, buff=1)
        self.play(Write(idea))
        self.wait()

        # Example
        example_sentence = Text("The bank was steep", font_size=36)
        example_sentence.next_to(idea, DOWN, buff=1)

        self.play(Write(example_sentence))
        self.wait()

        # Highlight "bank" and "steep"
        bank_box = SurroundingRectangle(
            example_sentence[4:8],  # "bank"
            color=YELLOW, buff=0.05
        )
        steep_box = SurroundingRectangle(
            example_sentence[13:18],  # "steep"
            color=GREEN, buff=0.05
        )

        self.play(Create(bank_box))
        self.wait()

        arrow = CurvedArrow(
            steep_box.get_top(),
            bank_box.get_top(),
            angle=-0.5,
            color=YELLOW
        )

        self.play(Create(steep_box), Create(arrow))

        meaning = Text(
            '"steep" helps us know "bank" means riverbank, not financial bank',
            font_size=24, color=GREY_B
        )
        meaning.next_to(example_sentence, DOWN, buff=0.8)
        self.play(Write(meaning))
        self.wait(2)

"""
Basic Animation Patterns for Manim Community

This file demonstrates fundamental animation techniques adapted from 3b1b patterns.
Run with: manim -pql basic_animations.py SceneName
"""

from manim import *


class ShapeCreation(Scene):
    """Demonstrates various ways to create and animate shapes."""

    def construct(self):
        # Create shapes
        circle = Circle(radius=1, color=BLUE, fill_opacity=0.5)
        square = Square(side_length=2, color=RED)
        triangle = Triangle(color=GREEN, fill_opacity=0.8)

        # Arrange shapes
        shapes = VGroup(circle, square, triangle).arrange(RIGHT, buff=1)

        # Different creation animations
        self.play(Create(circle))  # Draw outline progressively
        self.play(DrawBorderThenFill(square))  # Border first, then fill
        self.play(GrowFromCenter(triangle))  # Grow from center point

        self.wait()

        # Transform between shapes
        self.play(Transform(circle, square.copy().shift(UP * 2)))

        self.wait()


class TextAnimations(Scene):
    """Demonstrates text and LaTeX animations."""

    def construct(self):
        # Plain text
        title = Text("Manim Community", font_size=72, color=BLUE)
        self.play(Write(title))
        self.wait()

        # Move title up
        self.play(title.animate.to_edge(UP))

        # LaTeX math
        equation = MathTex(r"e^{i\pi} + 1 = 0", font_size=64)
        self.play(Write(equation))
        self.wait()

        # Transform equation
        expanded = MathTex(r"e^{i\pi} = -1", font_size=64)
        self.play(TransformMatchingTex(equation, expanded))

        self.wait()


class LaggedAnimations(Scene):
    """Demonstrates staggered animations using LaggedStart patterns."""

    def construct(self):
        # Create a grid of dots
        dots = VGroup(*[
            Dot(radius=0.15, color=interpolate_color(BLUE, RED, i / 24))
            for i in range(25)
        ]).arrange_in_grid(rows=5, cols=5, buff=0.5)

        # Staggered fade in
        self.play(
            LaggedStart(*[FadeIn(dot, scale=0.5) for dot in dots], lag_ratio=0.1)
        )
        self.wait()

        # Staggered transformation using LaggedStart with animate
        self.play(
            LaggedStart(
                *[dot.animate.scale(1.5).set_color(YELLOW) for dot in dots],
                lag_ratio=0.05
            )
        )
        self.wait()

        # Wave effect using AnimationGroup with rate_func
        self.play(
            LaggedStart(
                *[dot.animate(rate_func=there_and_back).shift(UP * 0.5) for dot in dots],
                lag_ratio=0.02,
                run_time=2
            )
        )


class AnimationComposition(Scene):
    """Demonstrates combining multiple animations."""

    def construct(self):
        # Create objects
        circle = Circle(color=BLUE, fill_opacity=0.5)
        label = Text("Circle", font_size=36).next_to(circle, DOWN)

        # Group them
        group = VGroup(circle, label)

        # Animate together
        self.play(
            Create(circle),
            Write(label),
            run_time=2
        )
        self.wait()

        # Sequential animations with Succession
        square = Square(color=RED, fill_opacity=0.5).shift(RIGHT * 3)
        square_label = Text("Square", font_size=36).next_to(square, DOWN)

        self.play(
            Succession(
                group.animate.shift(LEFT * 2),
                Create(square),
                Write(square_label),
                lag_ratio=0.5
            )
        )
        self.wait()


class PathAnimations(Scene):
    """Demonstrates movement along paths."""

    def construct(self):
        # Create a path
        path = VMobject()
        path.set_points_smoothly([
            LEFT * 3,
            LEFT * 2 + UP * 2,
            ORIGIN + UP,
            RIGHT * 2 + UP * 2,
            RIGHT * 3,
        ])
        path.set_color(GREY)

        # Create moving object
        dot = Dot(color=RED, radius=0.2)
        dot.move_to(path.get_start())

        self.add(path)
        self.play(Create(path))

        # Move along path
        self.play(MoveAlongPath(dot, path), run_time=3, rate_func=smooth)

        self.wait()


class ColorTransitions(Scene):
    """Demonstrates color manipulation and gradients."""

    def construct(self):
        # Color gradient on shapes
        squares = VGroup(*[
            Square(side_length=0.8, fill_opacity=0.8)
            for _ in range(7)
        ]).arrange(RIGHT, buff=0.2)

        # Apply gradient colors
        colors = [RED, ORANGE, YELLOW, GREEN, BLUE, PURPLE, PINK]
        for square, color in zip(squares, colors):
            square.set_fill(color)
            square.set_stroke(WHITE, width=2)

        self.play(LaggedStartMap(GrowFromCenter, squares, lag_ratio=0.1))
        self.wait()

        # Animate color change
        self.play(
            *[square.animate.set_fill(interpolate_color(BLUE, RED, i / 6))
              for i, square in enumerate(squares)],
            run_time=2
        )
        self.wait()


class GroupOperations(Scene):
    """Demonstrates VGroup operations and arrangements."""

    def construct(self):
        # Create VGroup
        shapes = VGroup(
            Circle(color=RED),
            Square(color=GREEN),
            Triangle(color=BLUE),
        )

        # Arrange horizontally
        shapes.arrange(RIGHT, buff=1)
        self.play(Create(shapes))
        self.wait()

        # Scale entire group
        self.play(shapes.animate.scale(0.5))
        self.wait()

        # Arrange vertically
        self.play(shapes.animate.arrange(DOWN, buff=0.5))
        self.wait()

        # Apply operation to all
        self.play(shapes.animate.set_fill(YELLOW, opacity=0.5))

        self.wait()