"""
Manim Code Generator

Generates Manim Python code from visualization plans.
Includes templates for common educational patterns.
"""

import json
from typing import Dict, Any, List, Optional
from pathlib import Path


class ManimTemplate:
    """Base class for Manim templates."""

    def __init__(self, name: str, description: str):
        self.name = name
        self.description = description

    def generate(self, params: Dict[str, Any]) -> str:
        raise NotImplementedError


class BinarySearchTemplate(ManimTemplate):
    """Template for binary search visualization."""

    def __init__(self):
        super().__init__(
            "binary_search",
            "Visualize binary search algorithm step by step"
        )

    def generate(self, params: Dict[str, Any]) -> str:
        array = params.get("array", [1, 3, 5, 7, 9, 11, 13, 15, 17, 19])
        target = params.get("target", 7)

        return f'''"""
Binary Search Visualization
Generated by LearnViz
"""
from manim import *

class BinarySearchScene(Scene):
    def construct(self):
        # Configuration
        array = {array}
        target = {target}

        # Title
        title = Text("Binary Search Algorithm", font_size=48)
        subtitle = Text(f"Finding {{target}} in a sorted array", font_size=24)
        subtitle.next_to(title, DOWN)

        self.play(Write(title))
        self.play(FadeIn(subtitle))
        self.wait(1)
        self.play(FadeOut(title), FadeOut(subtitle))

        # Create array visualization
        boxes = VGroup()
        numbers = VGroup()
        indices = VGroup()

        for i, val in enumerate(array):
            box = Square(side_length=0.8)
            box.set_stroke(WHITE, 2)
            num = Text(str(val), font_size=24)
            idx = Text(str(i), font_size=16, color=GRAY)

            if i > 0:
                box.next_to(boxes[-1], RIGHT, buff=0.1)

            num.move_to(box.get_center())
            idx.next_to(box, DOWN, buff=0.1)

            boxes.add(box)
            numbers.add(num)
            indices.add(idx)

        array_group = VGroup(boxes, numbers, indices)
        array_group.center()

        # Target display
        target_text = Text(f"Target: {{target}}", font_size=32)
        target_text.to_edge(UP)

        self.play(
            Create(boxes),
            Write(numbers),
            Write(indices),
            Write(target_text)
        )
        self.wait(1)

        # Binary search animation
        left, right = 0, len(array) - 1
        step = 1

        # Pointers
        left_arrow = Arrow(start=DOWN, end=UP, color=BLUE).scale(0.5)
        right_arrow = Arrow(start=DOWN, end=UP, color=RED).scale(0.5)
        left_label = Text("L", font_size=20, color=BLUE)
        right_label = Text("R", font_size=20, color=RED)

        left_arrow.next_to(boxes[left], DOWN, buff=0.5)
        right_arrow.next_to(boxes[right], DOWN, buff=0.5)
        left_label.next_to(left_arrow, DOWN, buff=0.1)
        right_label.next_to(right_arrow, DOWN, buff=0.1)

        self.play(
            Create(left_arrow), Write(left_label),
            Create(right_arrow), Write(right_label)
        )

        while left <= right:
            mid = (left + right) // 2

            # Step indicator
            step_text = Text(f"Step {{step}}: Check middle (index {{mid}})", font_size=24)
            step_text.to_edge(DOWN)
            self.play(Write(step_text))

            # Highlight middle
            self.play(boxes[mid].animate.set_fill(YELLOW, opacity=0.5))
            self.wait(0.5)

            if array[mid] == target:
                # Found!
                found_text = Text(f"Found {{target}} at index {{mid}}!", font_size=32, color=GREEN)
                found_text.next_to(step_text, UP)
                self.play(
                    boxes[mid].animate.set_fill(GREEN, opacity=0.7),
                    Write(found_text)
                )
                break
            elif array[mid] < target:
                # Go right
                compare_text = Text(f"{{array[mid]}} < {{target}}, search right half", font_size=20)
                compare_text.next_to(step_text, UP)
                self.play(Write(compare_text))

                # Gray out left half
                for i in range(left, mid + 1):
                    self.play(
                        boxes[i].animate.set_fill(GRAY, opacity=0.3),
                        run_time=0.2
                    )

                left = mid + 1
                left_arrow.generate_target()
                left_arrow.target.next_to(boxes[left] if left < len(array) else boxes[-1], DOWN, buff=0.5)
                left_label.generate_target()
                left_label.target.next_to(left_arrow.target, DOWN, buff=0.1)
                self.play(
                    MoveToTarget(left_arrow),
                    MoveToTarget(left_label)
                )
                self.play(FadeOut(compare_text))
            else:
                # Go left
                compare_text = Text(f"{{array[mid]}} > {{target}}, search left half", font_size=20)
                compare_text.next_to(step_text, UP)
                self.play(Write(compare_text))

                # Gray out right half
                for i in range(mid, right + 1):
                    self.play(
                        boxes[i].animate.set_fill(GRAY, opacity=0.3),
                        run_time=0.2
                    )

                right = mid - 1
                right_arrow.generate_target()
                right_arrow.target.next_to(boxes[right] if right >= 0 else boxes[0], DOWN, buff=0.5)
                right_label.generate_target()
                right_label.target.next_to(right_arrow.target, DOWN, buff=0.1)
                self.play(
                    MoveToTarget(right_arrow),
                    MoveToTarget(right_label)
                )
                self.play(FadeOut(compare_text))

            self.play(boxes[mid].animate.set_fill(WHITE, opacity=0))
            self.play(FadeOut(step_text))
            step += 1

        self.wait(2)


if __name__ == "__main__":
    # Render with: manim -pql binary_search.py BinarySearchScene
    pass
'''


class SortingTemplate(ManimTemplate):
    """Template for sorting algorithm visualization."""

    def __init__(self):
        super().__init__(
            "sorting",
            "Visualize sorting algorithms (bubble, selection, insertion)"
        )

    def generate(self, params: Dict[str, Any]) -> str:
        array = params.get("array", [64, 34, 25, 12, 22, 11, 90])
        algorithm = params.get("algorithm", "bubble")

        return f'''"""
{algorithm.title()} Sort Visualization
Generated by LearnViz
"""
from manim import *

class SortingScene(Scene):
    def construct(self):
        array = {array}
        n = len(array)

        # Title
        title = Text("{algorithm.title()} Sort", font_size=48)
        self.play(Write(title))
        self.wait(1)
        self.play(title.animate.to_edge(UP).scale(0.6))

        # Create bars
        bars = self.create_bars(array)
        self.play(Create(bars))
        self.wait(1)

        # Sort with animation
        self.{algorithm}_sort(bars, array)

        # Done
        done_text = Text("Sorted!", font_size=36, color=GREEN)
        done_text.to_edge(DOWN)
        self.play(Write(done_text))
        self.play(bars.animate.set_color(GREEN))
        self.wait(2)

    def create_bars(self, array):
        bars = VGroup()
        max_val = max(array)
        bar_width = 0.8
        spacing = 0.9

        for i, val in enumerate(array):
            height = (val / max_val) * 4
            bar = Rectangle(
                width=bar_width,
                height=height,
                fill_opacity=0.7,
                fill_color=BLUE,
                stroke_color=WHITE
            )
            bar.move_to(
                LEFT * (len(array) / 2 - i) * spacing + DOWN * (2 - height / 2)
            )

            # Value label
            label = Text(str(val), font_size=16)
            label.next_to(bar, UP, buff=0.1)
            bar.label = label

            bars.add(VGroup(bar, label))

        return bars

    def swap_bars(self, bars, i, j):
        bar_i, bar_j = bars[i], bars[j]
        pos_i, pos_j = bar_i.get_center(), bar_j.get_center()

        self.play(
            bar_i.animate.set_color(RED),
            bar_j.animate.set_color(RED),
            run_time=0.3
        )
        self.play(
            bar_i.animate.move_to([pos_j[0], bar_i.get_center()[1], 0]),
            bar_j.animate.move_to([pos_i[0], bar_j.get_center()[1], 0]),
            run_time=0.5
        )
        self.play(
            bar_i.animate.set_color(BLUE),
            bar_j.animate.set_color(BLUE),
            run_time=0.3
        )

        # Swap in list
        bars[i], bars[j] = bars[j], bars[i]

    def bubble_sort(self, bars, array):
        n = len(array)
        for i in range(n):
            for j in range(0, n - i - 1):
                # Highlight comparison
                self.play(
                    bars[j].animate.set_color(YELLOW),
                    bars[j + 1].animate.set_color(YELLOW),
                    run_time=0.2
                )

                if array[j] > array[j + 1]:
                    self.swap_bars(bars, j, j + 1)
                    array[j], array[j + 1] = array[j + 1], array[j]
                else:
                    self.play(
                        bars[j].animate.set_color(BLUE),
                        bars[j + 1].animate.set_color(BLUE),
                        run_time=0.2
                    )

    def selection_sort(self, bars, array):
        n = len(array)
        for i in range(n):
            min_idx = i
            self.play(bars[i].animate.set_color(PURPLE), run_time=0.2)

            for j in range(i + 1, n):
                self.play(bars[j].animate.set_color(YELLOW), run_time=0.1)
                if array[j] < array[min_idx]:
                    if min_idx != i:
                        self.play(bars[min_idx].animate.set_color(BLUE), run_time=0.1)
                    min_idx = j
                    self.play(bars[min_idx].animate.set_color(RED), run_time=0.1)
                else:
                    self.play(bars[j].animate.set_color(BLUE), run_time=0.1)

            if min_idx != i:
                self.swap_bars(bars, i, min_idx)
                array[i], array[min_idx] = array[min_idx], array[i]

            self.play(bars[i].animate.set_color(GREEN), run_time=0.2)

    def insertion_sort(self, bars, array):
        for i in range(1, len(array)):
            key = array[i]
            self.play(bars[i].animate.set_color(YELLOW), run_time=0.2)

            j = i - 1
            while j >= 0 and array[j] > key:
                self.swap_bars(bars, j, j + 1)
                array[j + 1] = array[j]
                j -= 1

            array[j + 1] = key
            self.play(bars[j + 1].animate.set_color(BLUE), run_time=0.2)


if __name__ == "__main__":
    pass
'''


class PythagoreanTemplate(ManimTemplate):
    """Template for Pythagorean theorem proof."""

    def __init__(self):
        super().__init__(
            "pythagorean",
            "Visualize Pythagorean theorem proof"
        )

    def generate(self, params: Dict[str, Any]) -> str:
        return '''"""
Pythagorean Theorem Visualization
Generated by LearnViz
"""
from manim import *

class PythagoreanScene(Scene):
    def construct(self):
        # Title
        title = Text("The Pythagorean Theorem", font_size=48)
        self.play(Write(title))
        self.wait(1)
        self.play(title.animate.to_edge(UP).scale(0.6))

        # Create right triangle
        a, b = 2, 1.5  # sides
        c = np.sqrt(a**2 + b**2)  # hypotenuse

        triangle = Polygon(
            ORIGIN,
            RIGHT * a,
            RIGHT * a + UP * b,
            fill_opacity=0.3,
            fill_color=BLUE,
            stroke_color=WHITE
        )
        triangle.move_to(ORIGIN)

        # Labels
        a_label = MathTex("a").next_to(triangle, DOWN)
        b_label = MathTex("b").next_to(triangle, RIGHT)
        c_label = MathTex("c").move_to(triangle.get_center() + UP * 0.5 + LEFT * 0.5)

        # Right angle marker
        right_angle = RightAngle(
            Line(ORIGIN, RIGHT * a),
            Line(RIGHT * a, RIGHT * a + UP * b),
            length=0.2
        )

        self.play(Create(triangle))
        self.play(Write(a_label), Write(b_label), Write(c_label))
        self.play(Create(right_angle))
        self.wait(1)

        # Create squares on each side
        sq_a = Square(side_length=a, fill_opacity=0.5, fill_color=RED)
        sq_a.next_to(triangle, DOWN, buff=0)

        sq_b = Square(side_length=b, fill_opacity=0.5, fill_color=GREEN)
        sq_b.next_to(triangle, RIGHT, buff=0)

        sq_c = Square(side_length=c, fill_opacity=0.5, fill_color=YELLOW)
        # Position c square on hypotenuse (simplified positioning)
        sq_c.rotate(np.arctan(b/a))
        sq_c.move_to(triangle.get_center() + UP * 1.5 + LEFT * 1)

        # Area labels
        a2_label = MathTex("a^2", color=RED).move_to(sq_a.get_center())
        b2_label = MathTex("b^2", color=GREEN).move_to(sq_b.get_center())
        c2_label = MathTex("c^2", color=YELLOW).move_to(sq_c.get_center())

        self.play(
            Create(sq_a), Write(a2_label),
            Create(sq_b), Write(b2_label),
            run_time=2
        )
        self.wait(1)

        self.play(Create(sq_c), Write(c2_label), run_time=2)
        self.wait(1)

        # Show the equation
        equation = MathTex("a^2", "+", "b^2", "=", "c^2", font_size=56)
        equation.to_edge(DOWN)
        equation[0].set_color(RED)
        equation[2].set_color(GREEN)
        equation[4].set_color(YELLOW)

        self.play(Write(equation))
        self.wait(1)

        # Highlight equation
        box = SurroundingRectangle(equation, color=WHITE, buff=0.2)
        self.play(Create(box))
        self.wait(2)


if __name__ == "__main__":
    pass
'''


class TreeTraversalTemplate(ManimTemplate):
    """Template for binary tree traversal visualization."""

    def __init__(self):
        super().__init__(
            "tree_traversal",
            "Visualize tree traversal (inorder, preorder, postorder)"
        )

    def generate(self, params: Dict[str, Any]) -> str:
        traversal = params.get("traversal", "inorder")
        tree = params.get("tree", [1, 2, 3, 4, 5, 6, 7])

        return f'''"""
Binary Tree {traversal.title()} Traversal
Generated by LearnViz
"""
from manim import *

class TreeNode:
    def __init__(self, val, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class TreeTraversalScene(Scene):
    def construct(self):
        # Title
        title = Text("{traversal.title()} Traversal", font_size=48)
        self.play(Write(title))
        self.wait(1)
        self.play(title.animate.to_edge(UP).scale(0.6))

        # Build tree from array
        tree_vals = {tree}
        root = self.build_tree(tree_vals)

        # Draw tree
        nodes, edges = self.draw_tree(root)
        self.play(*[Create(e) for e in edges])
        self.play(*[Create(n) for n in nodes.values()])
        self.wait(1)

        # Traversal order display
        order_label = Text("Visit Order: ", font_size=24)
        order_label.to_edge(DOWN).shift(LEFT * 4)
        order_values = VGroup()
        self.play(Write(order_label))

        # Perform traversal
        visited = []
        self.{traversal}_traverse(root, nodes, visited, order_values, order_label)

        self.wait(2)

    def build_tree(self, arr):
        if not arr:
            return None
        nodes = [TreeNode(v) if v is not None else None for v in arr]
        for i, node in enumerate(nodes):
            if node:
                left_idx = 2 * i + 1
                right_idx = 2 * i + 2
                if left_idx < len(nodes):
                    node.left = nodes[left_idx]
                if right_idx < len(nodes):
                    node.right = nodes[right_idx]
        return nodes[0] if nodes else None

    def draw_tree(self, root, x=0, y=2, dx=2, dy=1.5, nodes=None, edges=None, depth=0):
        if nodes is None:
            nodes = {{}}
            edges = []

        if root is None:
            return nodes, edges

        # Create node circle
        circle = Circle(radius=0.4, fill_opacity=0.7, fill_color=BLUE, stroke_color=WHITE)
        label = Text(str(root.val), font_size=24)
        node_group = VGroup(circle, label)
        node_group.move_to([x, y, 0])

        nodes[root.val] = node_group

        # Draw children
        if root.left:
            child_x = x - dx / (depth + 1)
            child_y = y - dy
            edge = Line([x, y - 0.4, 0], [child_x, child_y + 0.4, 0], color=WHITE)
            edges.append(edge)
            self.draw_tree(root.left, child_x, child_y, dx, dy, nodes, edges, depth + 1)

        if root.right:
            child_x = x + dx / (depth + 1)
            child_y = y - dy
            edge = Line([x, y - 0.4, 0], [child_x, child_y + 0.4, 0], color=WHITE)
            edges.append(edge)
            self.draw_tree(root.right, child_x, child_y, dx, dy, nodes, edges, depth + 1)

        return nodes, edges

    def visit_node(self, node, nodes, visited, order_values, order_label):
        if node.val in nodes:
            # Highlight visited
            self.play(nodes[node.val][0].animate.set_fill(GREEN, opacity=0.8), run_time=0.5)

            # Add to order
            val_text = Text(str(node.val), font_size=24, color=GREEN)
            if not visited:
                val_text.next_to(order_label, RIGHT)
            else:
                val_text.next_to(order_values[-1], RIGHT, buff=0.3)
            order_values.add(val_text)
            visited.append(node.val)
            self.play(Write(val_text), run_time=0.3)

    def inorder_traverse(self, node, nodes, visited, order_values, order_label):
        if node is None:
            return
        self.inorder_traverse(node.left, nodes, visited, order_values, order_label)
        self.visit_node(node, nodes, visited, order_values, order_label)
        self.inorder_traverse(node.right, nodes, visited, order_values, order_label)

    def preorder_traverse(self, node, nodes, visited, order_values, order_label):
        if node is None:
            return
        self.visit_node(node, nodes, visited, order_values, order_label)
        self.preorder_traverse(node.left, nodes, visited, order_values, order_label)
        self.preorder_traverse(node.right, nodes, visited, order_values, order_label)

    def postorder_traverse(self, node, nodes, visited, order_values, order_label):
        if node is None:
            return
        self.postorder_traverse(node.left, nodes, visited, order_values, order_label)
        self.postorder_traverse(node.right, nodes, visited, order_values, order_label)
        self.visit_node(node, nodes, visited, order_values, order_label)


if __name__ == "__main__":
    pass
'''


# Registry of all templates
class ActionPotentialTemplate(ManimTemplate):
    """Template for action potential propagation visualization with text management."""

    def __init__(self):
        super().__init__(
            "action_potential",
            "Visualize action potential propagation along a neuron"
        )

    def generate(self, params: Dict[str, Any]) -> str:
        return '''"""
Action Potential Propagation
Generated by LearnViz

Visualizes how an action potential travels along an axon.
Features: Text management to prevent overlaps, detailed explanations.
"""
from manim import *
import numpy as np


class TextManager:
    """Manages text display to prevent overlaps."""

    def __init__(self, scene):
        self.scene = scene
        self.active_texts = {}
        self.regions = {
            "title": UP * 3.5,
            "top": UP * 3.0,
            "bottom": DOWN * 3.3,
            "center": ORIGIN,
        }

    def show(self, text: str, region: str = "bottom", font_size: int = 22,
             color=WHITE, duration: float = 0.5) -> Text:
        """Display text in a region, clearing existing text there first."""
        if region in self.active_texts:
            self.scene.play(FadeOut(self.active_texts[region]), run_time=0.2)

        text_obj = Text(text, font_size=font_size, color=color)
        text_obj.move_to(self.regions.get(region, ORIGIN))

        if text_obj.width > 13:
            text_obj.scale(13 / text_obj.width)

        self.scene.play(Write(text_obj), run_time=duration)
        self.active_texts[region] = text_obj
        return text_obj

    def clear(self, region: str = None):
        """Clear text from region(s)."""
        if region and region in self.active_texts:
            self.scene.play(FadeOut(self.active_texts[region]), run_time=0.2)
            del self.active_texts[region]
        elif region is None:
            if self.active_texts:
                self.scene.play(*[FadeOut(t) for t in self.active_texts.values()], run_time=0.3)
                self.active_texts.clear()


class ActionPotentialScene(Scene):
    def construct(self):
        self.tm = TextManager(self)

        # Introduction
        self.show_intro()

        # Create neuron structure
        self.show_neuron_structure()
        self.wait(1)

        # Show resting potential
        self.show_resting_potential()
        self.wait(1)

        # Animate action potential
        self.animate_action_potential()
        self.wait(1)

        # Show voltage graph
        self.show_voltage_graph()

        # Conclusion
        self.show_conclusion()
        self.wait(2)

    def show_intro(self):
        """Introduction with educational context."""
        title = Text("Action Potential Propagation", font_size=48)
        subtitle = Text("How neurons transmit electrical signals", font_size=24, color=GRAY)
        subtitle.next_to(title, DOWN, buff=0.3)

        self.play(Write(title))
        self.play(FadeIn(subtitle))
        self.wait(1)

        intro_text = Text(
            "Neurons use electrical impulses called action potentials\\n"
            "to send information quickly over long distances.",
            font_size=20, color=BLUE_B
        )
        intro_text.next_to(subtitle, DOWN, buff=0.5)
        self.play(FadeIn(intro_text))
        self.wait(2)

        self.play(FadeOut(title), FadeOut(subtitle), FadeOut(intro_text))

    def show_neuron_structure(self):
        """Create and explain neuron structure."""
        self.tm.show(
            "Let's look at the structure of a neuron and how it transmits signals",
            region="bottom", font_size=20, color=BLUE_B
        )

        # Cell body (soma)
        self.soma = Circle(radius=0.6, fill_opacity=0.7, fill_color=BLUE, stroke_color=WHITE, stroke_width=2)
        self.soma.shift(LEFT * 5)
        soma_label = Text("Soma\\n(Cell Body)", font_size=14).next_to(self.soma, DOWN, buff=0.15)

        # Axon
        self.axon = Rectangle(width=8, height=0.4, fill_opacity=0.5, fill_color=BLUE_E, stroke_color=WHITE)
        self.axon.next_to(self.soma, RIGHT, buff=0)
        axon_label = Text("Axon", font_size=14).next_to(self.axon, DOWN, buff=0.1).shift(LEFT * 2)

        # Axon terminal
        self.terminal = Circle(radius=0.35, fill_opacity=0.7, fill_color=BLUE, stroke_color=WHITE, stroke_width=2)
        self.terminal.next_to(self.axon, RIGHT, buff=0)
        terminal_label = Text("Axon\\nTerminal", font_size=14).next_to(self.terminal, DOWN, buff=0.15)

        self.play(
            Create(self.soma), Write(soma_label),
            Create(self.axon), Write(axon_label),
            Create(self.terminal), Write(terminal_label),
            run_time=1.5
        )

        self.tm.show(
            "The soma is the cell body. The axon is a long fiber that carries signals. The terminal connects to other neurons.",
            region="bottom", font_size=18, color=WHITE
        )
        self.wait(2)

        # Myelin sheaths
        self.myelin_sheaths = VGroup()
        for i in range(4):
            sheath = Rectangle(width=1.4, height=0.6, fill_opacity=0.4, fill_color=YELLOW, stroke_color=YELLOW, stroke_width=1)
            sheath.move_to(self.axon.get_center() + LEFT * 2.5 + RIGHT * i * 2)
            self.myelin_sheaths.add(sheath)

        myelin_label = Text("Myelin Sheath", font_size=14, color=YELLOW)
        myelin_label.next_to(self.myelin_sheaths[1], UP, buff=0.1)

        self.play(Create(self.myelin_sheaths), Write(myelin_label))

        self.tm.show(
            "Myelin sheaths are fatty insulation that speeds up signal transmission by up to 100x",
            region="bottom", font_size=18, color=YELLOW
        )
        self.wait(2)

        # Nodes of Ranvier
        node_arrows = VGroup()
        for i in range(3):
            pos = self.axon.get_center() + LEFT * 1.5 + RIGHT * i * 2
            arrow = Arrow(pos + UP * 0.8, pos + UP * 0.3, color=RED, stroke_width=2, max_tip_length_to_length_ratio=0.3)
            node_arrows.add(arrow)

        node_label = Text("Nodes of Ranvier (gaps)", font_size=12, color=RED)
        node_label.next_to(node_arrows[1], UP, buff=0.05)

        self.play(Create(node_arrows), Write(node_label))

        self.tm.show(
            "Nodes of Ranvier are gaps between myelin where action potentials are regenerated",
            region="bottom", font_size=18, color=RED
        )
        self.wait(2)

        self.structure_labels = VGroup(soma_label, axon_label, terminal_label, myelin_label, node_label, node_arrows)
        self.neuron_group = VGroup(self.soma, self.axon, self.terminal, self.myelin_sheaths)

    def show_resting_potential(self):
        """Explain and visualize resting potential."""
        self.tm.show(
            "At rest, the neuron maintains a voltage of -70mV (resting potential)",
            region="bottom", font_size=20, color=BLUE_B
        )

        # Show charge distribution
        self.plus_outside = VGroup(*[
            Text("+", font_size=18, color=RED).move_to(self.axon.get_top() + UP * 0.25 + LEFT * 3 + RIGHT * i * 1.4)
            for i in range(5)
        ])
        self.minus_inside = VGroup(*[
            Text("-", font_size=18, color=BLUE).move_to(self.axon.get_center() + LEFT * 3 + RIGHT * i * 1.4)
            for i in range(5)
        ])

        self.play(Write(self.plus_outside), Write(self.minus_inside))

        self.tm.show(
            "Positive ions (Na+, K+) outside, negative inside - this creates electrical potential energy",
            region="bottom", font_size=18, color=WHITE
        )
        self.wait(2)

    def animate_action_potential(self):
        """Animate and explain action potential propagation."""
        self.tm.show(
            "When triggered, an action potential begins: sodium channels open, Na+ rushes in",
            region="bottom", font_size=20, color=YELLOW
        )

        # Create propagating wave
        wave = Circle(radius=0.3, fill_opacity=0.85, fill_color=YELLOW, stroke_color=ORANGE, stroke_width=3)
        wave.move_to(self.soma.get_center())

        self.play(Create(wave))
        self.wait(0.5)

        self.tm.show(
            "The action potential propagates along the axon by jumping between nodes (saltatory conduction)",
            region="bottom", font_size=18, color=ORANGE
        )

        # Path points at nodes of Ranvier
        path_points = [
            self.soma.get_right(),
            self.axon.get_left() + RIGHT * 1.5,
            self.axon.get_center() + LEFT * 0.5,
            self.axon.get_center() + RIGHT * 1.5,
            self.axon.get_right() - RIGHT * 0.5,
            self.terminal.get_left()
        ]

        for i, point in enumerate(path_points):
            flash = Circle(radius=0.4, fill_opacity=0, stroke_color=YELLOW, stroke_width=3)
            flash.move_to(point)

            self.play(
                wave.animate.move_to(point),
                Create(flash),
                run_time=0.4
            )
            self.play(FadeOut(flash), run_time=0.15)

        # Wave reaches terminal
        self.play(
            wave.animate.scale(1.3),
            Flash(self.terminal.get_center(), color=YELLOW, line_length=0.2, num_lines=10),
            run_time=0.5
        )

        self.tm.show(
            "Signal reaches the terminal! This triggers neurotransmitter release to the next neuron.",
            region="bottom", font_size=20, color=GREEN
        )
        self.wait(1.5)

        self.wave = wave

    def show_voltage_graph(self):
        """Show voltage vs time graph with phase explanations."""
        # Clear and reposition
        self.tm.clear()
        self.play(
            FadeOut(self.plus_outside), FadeOut(self.minus_inside),
            FadeOut(self.wave), FadeOut(self.structure_labels),
            self.neuron_group.animate.shift(UP * 1.8).scale(0.6),
            run_time=0.8
        )

        self.tm.show(
            "The action potential follows a characteristic voltage pattern over time",
            region="top", font_size=18, color=BLUE_B
        )

        # Create axes
        axes = Axes(
            x_range=[0, 5, 1],
            y_range=[-90, 60, 30],
            x_length=9,
            y_length=3.5,
            axis_config={"include_tip": True, "tip_length": 0.2},
        ).shift(DOWN * 1.2)

        x_label = Text("Time (ms)", font_size=16).next_to(axes.x_axis, DOWN, buff=0.1)
        y_label = Text("mV", font_size=16).next_to(axes.y_axis, LEFT, buff=0.1)

        self.play(Create(axes), Write(x_label), Write(y_label), run_time=1)

        # Action potential curve
        def ap_curve(t):
            if t < 1:
                return -70
            elif t < 1.5:
                return -70 + 140 * (t - 1) / 0.5
            elif t < 2:
                return 40 - 120 * (t - 1.5) / 0.5
            elif t < 2.5:
                return -80 + 10 * (t - 2) / 0.5
            else:
                return -70

        graph = axes.plot(ap_curve, x_range=[0, 4], color=YELLOW, stroke_width=3)
        self.play(Create(graph), run_time=2)

        # Phase labels with explanations
        phases = [
            ("Resting\\n-70mV", 0.5, -70, WHITE, "Stable baseline"),
            ("Depolar-\\nization", 1.25, -15, GREEN, "Na+ rushes in"),
            ("Peak\\n+40mV", 1.5, 45, RED, "Maximum voltage"),
            ("Repolar-\\nization", 1.75, -20, BLUE, "K+ flows out"),
            ("Hyper-\\npolar.", 2.25, -85, PURPLE, "Brief undershoot"),
        ]

        for name, x, y, color, explanation in phases:
            point = axes.c2p(x, y)
            dot = Dot(point, color=color, radius=0.08)
            label = Text(name, font_size=11, color=color)
            label.next_to(dot, UP if y > -50 else DOWN, buff=0.1)
            self.play(Create(dot), Write(label), run_time=0.4)

        self.wait(1)
        self.axes = axes
        self.graph = graph

    def show_conclusion(self):
        """Summary of action potential."""
        self.tm.show(
            "The entire action potential takes only 1-2 milliseconds - enabling rapid brain communication!",
            region="top", font_size=20, color=GREEN
        )
        self.wait(2)


if __name__ == "__main__":
    # Render with: manim -pql action_potential.py ActionPotentialScene
    pass
'''


class SynapseTemplate(ManimTemplate):
    """Template for synaptic transmission visualization with proper text management."""

    def __init__(self):
        super().__init__(
            "synapse",
            "Visualize neurotransmitter release at a synapse"
        )

    def generate(self, params: Dict[str, Any]) -> str:
        include_narration = params.get("narration", True)
        return '''"""
Synaptic Transmission
Generated by LearnViz

Visualizes how signals cross from one neuron to another.
Features: Text management to prevent overlaps, detailed explanations.
"""
from manim import *


class TextManager:
    """Manages text display to prevent overlaps."""

    def __init__(self, scene):
        self.scene = scene
        self.active_texts = {}
        self.regions = {
            "title": UP * 3.5,
            "top": UP * 3.0,
            "bottom": DOWN * 3.3,
            "center": ORIGIN,
        }

    def show(self, text: str, region: str = "bottom", font_size: int = 22,
             color=WHITE, duration: float = 0.5) -> Text:
        """Display text in a region, clearing existing text there first."""
        if region in self.active_texts:
            self.scene.play(FadeOut(self.active_texts[region]), run_time=0.2)

        text_obj = Text(text, font_size=font_size, color=color)
        text_obj.move_to(self.regions.get(region, ORIGIN))

        # Scale if too wide
        if text_obj.width > 13:
            text_obj.scale(13 / text_obj.width)

        self.scene.play(Write(text_obj), run_time=duration)
        self.active_texts[region] = text_obj
        return text_obj

    def clear(self, region: str = None):
        """Clear text from region(s)."""
        if region and region in self.active_texts:
            self.scene.play(FadeOut(self.active_texts[region]), run_time=0.2)
            del self.active_texts[region]
        elif region is None:
            if self.active_texts:
                self.scene.play(*[FadeOut(t) for t in self.active_texts.values()], run_time=0.3)
                self.active_texts.clear()


class SynapseScene(Scene):
    def construct(self):
        self.tm = TextManager(self)

        # Introduction
        self.show_intro()

        # Build and explain synapse structure
        self.create_synapse()
        self.wait(1)

        # Signal arrival
        self.signal_arrives()
        self.wait(0.5)

        # Vesicle release with explanation
        self.release_neurotransmitters()
        self.wait(0.5)

        # Receptor binding
        self.receptor_binding()

        # Conclusion
        self.show_conclusion()
        self.wait(2)

    def show_intro(self):
        """Introduction with clear explanation."""
        title = Text("Synaptic Transmission", font_size=48)
        subtitle = Text("How neurons communicate with each other", font_size=24, color=GRAY)
        subtitle.next_to(title, DOWN, buff=0.3)

        self.play(Write(title))
        self.play(FadeIn(subtitle))
        self.wait(1.5)

        # Educational context
        intro_text = Text(
            "Neurons don't physically touch - they communicate\\n"
            "across tiny gaps called synapses using chemicals.",
            font_size=20, color=BLUE_B
        )
        intro_text.next_to(subtitle, DOWN, buff=0.5)
        self.play(FadeIn(intro_text))
        self.wait(2)

        self.play(FadeOut(title), FadeOut(subtitle), FadeOut(intro_text))

    def create_synapse(self):
        """Create synapse structure with labeled components."""
        # Explanation text at bottom
        self.tm.show(
            "The synapse has three main parts: presynaptic terminal, synaptic cleft, and postsynaptic membrane",
            region="bottom", font_size=20, color=BLUE_B
        )

        # Presynaptic terminal (top) - shifted up to make room
        self.presynaptic = RoundedRectangle(
            width=5, height=2.2, corner_radius=0.3,
            fill_opacity=0.6, fill_color=BLUE, stroke_color=WHITE, stroke_width=2
        ).shift(UP * 1.8)

        pre_label = Text("Presynaptic Terminal", font_size=16, color=BLUE_B)
        pre_label.next_to(self.presynaptic, LEFT, buff=0.2)

        # Vesicles inside presynaptic - clearly visible
        self.vesicles = VGroup()
        vesicle_label_added = False
        for i in range(6):
            vesicle = Circle(radius=0.18, fill_opacity=0.9, fill_color=YELLOW, stroke_color=ORANGE, stroke_width=2)
            vesicle.move_to(self.presynaptic.get_center() + UP * 0.4 + LEFT * 1.2 + RIGHT * i * 0.45)
            self.vesicles.add(vesicle)

        vesicle_label = Text("Vesicles (contain\\nneurotransmitters)", font_size=12, color=YELLOW)
        vesicle_label.next_to(self.vesicles, RIGHT, buff=0.3)

        # Synaptic cleft (gap) - make it more visible
        self.cleft = Rectangle(
            width=5.5, height=0.6,
            fill_opacity=0.15, fill_color=WHITE, stroke_color=GRAY, stroke_width=1
        )
        cleft_label = Text("Synaptic Cleft (~20nm gap)", font_size=14, color=GRAY)
        cleft_label.next_to(self.cleft, RIGHT, buff=0.2)

        # Postsynaptic membrane (bottom)
        self.postsynaptic = RoundedRectangle(
            width=5, height=1.8, corner_radius=0.3,
            fill_opacity=0.6, fill_color=GREEN, stroke_color=WHITE, stroke_width=2
        ).shift(DOWN * 1.8)

        post_label = Text("Postsynaptic Neuron", font_size=16, color=GREEN_B)
        post_label.next_to(self.postsynaptic, LEFT, buff=0.2)

        # Receptors on postsynaptic membrane - more prominent
        self.receptors = VGroup()
        for i in range(5):
            receptor = Rectangle(width=0.25, height=0.5, fill_opacity=0.9, fill_color=RED, stroke_color=WHITE, stroke_width=1)
            receptor.move_to(self.postsynaptic.get_top() + UP * 0.1 + LEFT * 1 + RIGHT * i * 0.55)
            self.receptors.add(receptor)

        receptor_label = Text("Receptors", font_size=12, color=RED)
        receptor_label.next_to(self.receptors, LEFT, buff=0.2).shift(UP * 0.2)

        # Animate construction with explanations
        self.play(Create(self.presynaptic), Write(pre_label), run_time=1)
        self.play(Create(self.vesicles), Write(vesicle_label), run_time=1)

        self.tm.show(
            "Vesicles are small spheres filled with neurotransmitter molecules",
            region="bottom", font_size=20, color=YELLOW
        )
        self.wait(1.5)

        self.play(Create(self.cleft), Write(cleft_label), run_time=0.8)

        self.tm.show(
            "The synaptic cleft is a tiny gap - only about 20 nanometers wide",
            region="bottom", font_size=20, color=GRAY
        )
        self.wait(1.5)

        self.play(Create(self.postsynaptic), Write(post_label), run_time=1)
        self.play(Create(self.receptors), Write(receptor_label), run_time=0.8)

        self.tm.show(
            "Receptors are proteins that recognize specific neurotransmitters",
            region="bottom", font_size=20, color=RED
        )
        self.wait(1.5)

        # Store labels for later reference
        self.structure_labels = VGroup(pre_label, vesicle_label, cleft_label, post_label, receptor_label)

    def signal_arrives(self):
        """Action potential arrives with explanation."""
        self.tm.show(
            "Step 1: An action potential arrives at the presynaptic terminal",
            region="bottom", font_size=22, color=YELLOW
        )

        # Visual indicator of signal arriving
        signal_arrow = Arrow(
            start=self.presynaptic.get_left() + LEFT * 1.5,
            end=self.presynaptic.get_left(),
            color=YELLOW, stroke_width=4
        )
        signal_label = Text("Action\\nPotential", font_size=14, color=YELLOW)
        signal_label.next_to(signal_arrow, UP, buff=0.1)

        self.play(Create(signal_arrow), Write(signal_label), run_time=0.8)

        # Flash the presynaptic terminal
        self.play(
            self.presynaptic.animate.set_fill(YELLOW, opacity=0.8),
            Flash(self.presynaptic.get_center(), color=YELLOW, line_length=0.3, num_lines=12),
            run_time=0.6
        )
        self.play(
            self.presynaptic.animate.set_fill(BLUE, opacity=0.6),
            run_time=0.4
        )

        self.tm.show(
            "This triggers calcium ions to enter, causing vesicles to move toward the membrane",
            region="bottom", font_size=20, color=BLUE_B
        )
        self.wait(1.5)

        self.play(FadeOut(signal_arrow), FadeOut(signal_label), run_time=0.3)

    def release_neurotransmitters(self):
        """Vesicle fusion and neurotransmitter release with detailed explanation."""
        self.tm.show(
            "Step 2: Vesicles fuse with the membrane and release neurotransmitters",
            region="bottom", font_size=22, color=ORANGE
        )

        # Neurotransmitters released
        self.neurotransmitters = VGroup()

        for i, vesicle in enumerate(self.vesicles[:3]):
            # Move vesicle to membrane
            target_pos = self.presynaptic.get_bottom() + RIGHT * (i - 1) * 0.8
            self.play(vesicle.animate.move_to(target_pos), run_time=0.4)

            # Vesicle "opens" - change appearance
            self.play(vesicle.animate.set_stroke(color=WHITE, width=3), run_time=0.2)

            # Release neurotransmitters
            for j in range(4):
                nt = Dot(radius=0.08, color=YELLOW)
                nt.move_to(vesicle.get_center())
                self.neurotransmitters.add(nt)

                # Scatter into cleft
                target = self.cleft.get_center() + LEFT * 0.8 + RIGHT * (i * 1.2 + j * 0.25) + UP * 0.1 * (j % 2)
                self.add(nt)
                self.play(nt.animate.move_to(target), run_time=0.15)

            # Fade vesicle (recycled)
            self.play(FadeOut(vesicle), run_time=0.2)

        self.tm.show(
            "Thousands of neurotransmitter molecules are released in microseconds",
            region="bottom", font_size=20, color=YELLOW
        )
        self.wait(1.5)

    def receptor_binding(self):
        """Neurotransmitters bind to receptors with explanation."""
        self.tm.show(
            "Step 3: Neurotransmitters diffuse across and bind to receptors",
            region="bottom", font_size=22, color=GREEN
        )

        # Move neurotransmitters to receptors
        for i, nt in enumerate(self.neurotransmitters[:5]):
            if i < len(self.receptors):
                target = self.receptors[i].get_top() + UP * 0.05
                self.play(nt.animate.move_to(target), run_time=0.25)
                # Receptor activates - color change
                self.play(
                    self.receptors[i].animate.set_fill(YELLOW, opacity=0.95),
                    run_time=0.2
                )

        self.wait(0.5)
        self.tm.show(
            "Receptor activation opens ion channels, allowing the signal to continue",
            region="bottom", font_size=20, color=GREEN
        )

        # Postsynaptic response
        self.play(
            self.postsynaptic.animate.set_fill(YELLOW, opacity=0.6),
            Flash(self.postsynaptic.get_center(), color=GREEN, line_length=0.2, num_lines=8),
            run_time=0.5
        )
        self.play(
            self.postsynaptic.animate.set_fill(GREEN, opacity=0.6),
            run_time=0.3
        )

        self.wait(1)

    def show_conclusion(self):
        """Wrap up with key takeaways."""
        self.tm.clear()

        # Fade structure slightly
        self.play(
            self.presynaptic.animate.set_opacity(0.4),
            self.postsynaptic.animate.set_opacity(0.4),
            self.cleft.animate.set_opacity(0.1),
            self.structure_labels.animate.set_opacity(0.4),
            run_time=0.5
        )

        # Summary box
        summary = VGroup(
            Text("Synaptic Transmission Summary", font_size=24, color=WHITE),
            Text("1. Action potential triggers calcium influx", font_size=18, color=YELLOW),
            Text("2. Vesicles fuse and release neurotransmitters", font_size=18, color=ORANGE),
            Text("3. Neurotransmitters bind to receptors", font_size=18, color=GREEN),
            Text("4. Ion channels open, signal continues", font_size=18, color=BLUE),
        ).arrange(DOWN, aligned_edge=LEFT, buff=0.25)

        box = SurroundingRectangle(summary, color=WHITE, buff=0.3, corner_radius=0.1)
        summary_group = VGroup(box, summary)
        summary_group.move_to(ORIGIN)

        self.play(Create(box), Write(summary), run_time=2)
        self.wait(2)

        # Final message
        final = Text("This process occurs billions of times per second in your brain!", font_size=20, color=BLUE_B)
        final.next_to(summary_group, DOWN, buff=0.5)
        self.play(FadeIn(final))


if __name__ == "__main__":
    # Render with: manim -pql synapse.py SynapseScene
    pass
'''


class MotorCortexBCITemplate(ManimTemplate):
    """Template for motor cortex and BCI neural decoding visualization."""

    def __init__(self):
        super().__init__(
            "motor_cortex_bci",
            "Visualize how BCIs decode motor intent from M1"
        )

    def generate(self, params: Dict[str, Any]) -> str:
        return '''"""
Motor Cortex & BCI Neural Decoding
Generated by LearnViz

Visualizes how brain-computer interfaces decode motor intent from M1.
"""
from manim import *
import numpy as np


class TextManager:
    """Manages text display to prevent overlaps."""

    def __init__(self, scene):
        self.scene = scene
        self.active_texts = {}
        self.regions = {
            "title": UP * 3.5,
            "top": UP * 3.0,
            "bottom": DOWN * 3.3,
            "center": ORIGIN,
        }

    def show(self, text: str, region: str = "bottom", font_size: int = 22,
             color=WHITE, duration: float = 0.5) -> Text:
        if region in self.active_texts:
            self.scene.play(FadeOut(self.active_texts[region]), run_time=0.2)
        text_obj = Text(text, font_size=font_size, color=color)
        text_obj.move_to(self.regions.get(region, ORIGIN))
        if text_obj.width > 13:
            text_obj.scale(13 / text_obj.width)
        self.scene.play(Write(text_obj), run_time=duration)
        self.active_texts[region] = text_obj
        return text_obj

    def clear(self, region: str = None):
        if region and region in self.active_texts:
            self.scene.play(FadeOut(self.active_texts[region]), run_time=0.2)
            del self.active_texts[region]
        elif region is None and self.active_texts:
            self.scene.play(*[FadeOut(t) for t in self.active_texts.values()], run_time=0.3)
            self.active_texts.clear()


class MotorCortexBCIScene(Scene):
    def construct(self):
        self.tm = TextManager(self)

        self.show_intro()
        self.show_brain_overview()
        self.show_electrode_array()
        self.show_neural_signals()
        self.show_population_coding()
        self.show_bci_decoding()
        self.show_conclusion()
        self.wait(2)

    def show_intro(self):
        title = Text("Motor Cortex & BCI Decoding", font_size=44)
        subtitle = Text("How BCIs read your motor intentions", font_size=24, color=GRAY)
        subtitle.next_to(title, DOWN, buff=0.3)

        self.play(Write(title))
        self.play(FadeIn(subtitle))
        self.wait(1)

        intro = Text(
            "Brain-Computer Interfaces decode movement intent\\n"
            "from the primary motor cortex (M1)",
            font_size=20, color=BLUE_B
        )
        intro.next_to(subtitle, DOWN, buff=0.5)
        self.play(FadeIn(intro))
        self.wait(2)
        self.play(FadeOut(title), FadeOut(subtitle), FadeOut(intro))

    def show_brain_overview(self):
        self.tm.show("The motor cortex (M1) controls voluntary movement", region="bottom", color=BLUE_B)

        # Simplified brain outline
        brain = Ellipse(width=5, height=4, color=PINK, fill_opacity=0.3, stroke_width=2)
        brain.shift(LEFT * 2)

        # Motor cortex region (precentral gyrus)
        m1 = Arc(radius=1.8, start_angle=PI/3, angle=PI/3, color=RED, stroke_width=6)
        m1.shift(LEFT * 2 + UP * 0.3)
        m1_label = Text("M1 (Motor Cortex)", font_size=16, color=RED)
        m1_label.next_to(m1, UP, buff=0.2)

        self.play(Create(brain))
        self.play(Create(m1), Write(m1_label))

        # Motor homunculus representation
        self.tm.show("M1 has a 'motor map' - different areas control different body parts", region="bottom")

        homunculus = VGroup()
        parts = [("Leg", UP*0.8), ("Arm", UP*0.3+RIGHT*0.3), ("Hand", DOWN*0.1+RIGHT*0.5), ("Face", DOWN*0.6+RIGHT*0.7)]
        for name, offset in parts:
            dot = Dot(m1.get_center() + offset, color=YELLOW, radius=0.1)
            label = Text(name, font_size=12, color=YELLOW)
            label.next_to(dot, RIGHT, buff=0.1)
            homunculus.add(VGroup(dot, label))

        self.play(Create(homunculus))
        self.wait(2)

        self.brain = brain
        self.m1 = m1
        self.homunculus = homunculus

    def show_electrode_array(self):
        self.tm.show("BCIs use electrode arrays implanted in M1 to record neural signals", region="bottom", color=ORANGE)

        # Utah array representation
        array = VGroup()
        array_bg = Square(side_length=1.2, fill_opacity=0.8, fill_color=DARK_GRAY, stroke_color=WHITE)

        # 10x10 electrode grid
        for i in range(10):
            for j in range(10):
                electrode = Dot(radius=0.03, color=GOLD)
                electrode.move_to(array_bg.get_center() + LEFT*0.5 + DOWN*0.5 + RIGHT*i*0.11 + UP*j*0.11)
                array.add(electrode)

        array.add(array_bg)
        array.move_to(self.m1.get_center())

        array_label = Text("Utah Array\\n(96 electrodes)", font_size=14, color=GOLD)
        array_label.next_to(array, RIGHT, buff=0.3)

        # Zoom effect
        self.play(
            self.brain.animate.set_opacity(0.2),
            self.homunculus.animate.set_opacity(0.3),
        )
        self.play(Create(array), Write(array_label))
        self.wait(1.5)

        self.array = array
        self.array_label = array_label

    def show_neural_signals(self):
        self.tm.show("Each electrode records 'spikes' - electrical signals from nearby neurons", region="bottom", color=GREEN)

        # Move array to left side
        self.play(
            FadeOut(self.brain), FadeOut(self.m1), FadeOut(self.homunculus),
            self.array.animate.move_to(LEFT * 4),
            self.array_label.animate.next_to(LEFT * 4, DOWN, buff=0.3)
        )

        # Signal traces
        signals = VGroup()
        axes_group = VGroup()

        for i in range(4):
            # Mini axes for each channel
            ax = Axes(x_range=[0, 2, 1], y_range=[-1, 1, 1], x_length=3, y_length=0.8,
                     axis_config={"include_tip": False, "stroke_width": 1})
            ax.move_to(RIGHT * 1 + UP * (1.5 - i))

            # Random spikes
            np.random.seed(i + 42)
            spike_times = np.random.uniform(0, 2, np.random.randint(3, 8))

            def make_spike_func(times):
                def f(t):
                    val = 0
                    for st in times:
                        if abs(t - st) < 0.05:
                            val += np.exp(-((t - st) / 0.02) ** 2)
                    return val * 0.8
                return f

            trace = ax.plot(make_spike_func(spike_times), color=GREEN, stroke_width=2)

            ch_label = Text(f"Ch {i+1}", font_size=12, color=GRAY)
            ch_label.next_to(ax, LEFT, buff=0.1)

            signals.add(trace)
            axes_group.add(VGroup(ax, ch_label))

        self.play(Create(axes_group))
        self.play(Create(signals), run_time=2)

        self.tm.show("Each spike represents a neuron firing - encoding movement information", region="bottom")
        self.wait(2)

        self.signals = signals
        self.axes_group = axes_group

    def show_population_coding(self):
        self.tm.show("Neurons encode movement DIRECTION using 'population coding'", region="bottom", color=YELLOW)

        # Clear previous
        self.play(FadeOut(self.array), FadeOut(self.array_label),
                 FadeOut(self.signals), FadeOut(self.axes_group))

        # Population vector diagram
        center = ORIGIN

        # Individual neuron preferred directions
        neurons = VGroup()
        vectors = VGroup()
        n_neurons = 8

        for i in range(n_neurons):
            angle = i * 2 * PI / n_neurons
            direction = np.array([np.cos(angle), np.sin(angle), 0])

            # Neuron dot
            neuron = Dot(center + direction * 2.5, color=BLUE, radius=0.15)
            neurons.add(neuron)

            # Contribution vector (varies by neuron)
            weight = 0.3 + 0.7 * max(0, np.cos(angle - PI/4))  # Stronger for upper-right
            vec = Arrow(center, center + direction * weight * 1.5, color=BLUE_B, stroke_width=2, buff=0)
            vectors.add(vec)

        # Population vector (sum)
        pop_vec = Arrow(center, center + np.array([1.2, 1.2, 0]), color=RED, stroke_width=4, buff=0)
        pop_label = Text("Population\\nVector", font_size=14, color=RED)
        pop_label.next_to(pop_vec.get_end(), UR, buff=0.1)

        self.play(Create(neurons))
        self.tm.show("Each neuron has a 'preferred direction' - it fires most for that movement", region="bottom")
        self.wait(1)

        self.play(Create(vectors))
        self.tm.show("Neuron firing rates combine to form a 'population vector'", region="bottom")
        self.wait(1)

        self.play(Create(pop_vec), Write(pop_label))
        self.tm.show("The population vector points in the intended movement direction!", region="bottom", color=RED)
        self.wait(2)

        self.neurons = neurons
        self.vectors = vectors
        self.pop_vec = pop_vec
        self.pop_label = pop_label

    def show_bci_decoding(self):
        self.tm.show("BCIs decode this population vector to control devices", region="bottom", color=GREEN)

        # Transition
        self.play(
            self.neurons.animate.shift(LEFT * 3).scale(0.6),
            self.vectors.animate.shift(LEFT * 3).scale(0.6),
            self.pop_vec.animate.shift(LEFT * 3).scale(0.6),
            self.pop_label.animate.shift(LEFT * 3).scale(0.6),
        )

        # Arrow to decoded output
        decode_arrow = Arrow(LEFT * 1, RIGHT * 1, color=WHITE, stroke_width=3)
        decode_label = Text("Decode", font_size=16)
        decode_label.next_to(decode_arrow, UP, buff=0.1)

        self.play(Create(decode_arrow), Write(decode_label))

        # Output: cursor movement
        cursor = Dot(RIGHT * 3, color=GREEN, radius=0.2)
        cursor_label = Text("Cursor", font_size=14, color=GREEN)
        cursor_label.next_to(cursor, DOWN, buff=0.2)

        screen = Rectangle(width=3, height=2, stroke_color=WHITE, fill_opacity=0.1)
        screen.move_to(RIGHT * 3)

        self.play(Create(screen), Create(cursor), Write(cursor_label))

        # Animate cursor moving in decoded direction
        self.play(cursor.animate.shift(UR * 0.5), run_time=0.5)
        self.play(cursor.animate.shift(UR * 0.3), run_time=0.3)

        self.tm.show("Patient thinks 'move right' → M1 neurons fire → BCI decodes → cursor moves!", region="bottom", color=GREEN)
        self.wait(2)

    def show_conclusion(self):
        self.tm.clear()

        # Summary
        summary = VGroup(
            Text("BCI Neural Decoding Summary", font_size=24, color=WHITE),
            Text("1. Electrodes record spikes from M1 neurons", font_size=18, color=GOLD),
            Text("2. Each neuron has a preferred movement direction", font_size=18, color=BLUE),
            Text("3. Population coding combines all neurons", font_size=18, color=YELLOW),
            Text("4. Decoder translates to device commands", font_size=18, color=GREEN),
        ).arrange(DOWN, aligned_edge=LEFT, buff=0.25)

        box = SurroundingRectangle(summary, color=WHITE, buff=0.3, corner_radius=0.1)
        summary_group = VGroup(box, summary)
        summary_group.move_to(ORIGIN)

        self.play(FadeOut(*self.mobjects))
        self.play(Create(box), Write(summary), run_time=2)

        final = Text("This is how paralyzed patients control robotic arms with thought!", font_size=18, color=BLUE_B)
        final.next_to(summary_group, DOWN, buff=0.5)
        self.play(FadeIn(final))


if __name__ == "__main__":
    pass
'''


class NeurotransmitterTemplate(ManimTemplate):
    """Template for neurotransmitter systems visualization."""

    def __init__(self):
        super().__init__(
            "neurotransmitter",
            "Visualize dopamine, serotonin, and norepinephrine pathways"
        )

    def generate(self, params: Dict[str, Any]) -> str:
        return '''"""
Neurotransmitter Systems
Generated by LearnViz

Visualizes dopamine, serotonin, and norepinephrine pathways.
"""
from manim import *
import numpy as np


class TextManager:
    def __init__(self, scene):
        self.scene = scene
        self.active_texts = {}
        self.regions = {"title": UP * 3.5, "top": UP * 3.0, "bottom": DOWN * 3.3, "center": ORIGIN}

    def show(self, text: str, region: str = "bottom", font_size: int = 22, color=WHITE, duration: float = 0.5):
        if region in self.active_texts:
            self.scene.play(FadeOut(self.active_texts[region]), run_time=0.2)
        text_obj = Text(text, font_size=font_size, color=color)
        text_obj.move_to(self.regions.get(region, ORIGIN))
        if text_obj.width > 13:
            text_obj.scale(13 / text_obj.width)
        self.scene.play(Write(text_obj), run_time=duration)
        self.active_texts[region] = text_obj
        return text_obj

    def clear(self, region: str = None):
        if region and region in self.active_texts:
            self.scene.play(FadeOut(self.active_texts[region]), run_time=0.2)
            del self.active_texts[region]
        elif region is None and self.active_texts:
            self.scene.play(*[FadeOut(t) for t in self.active_texts.values()], run_time=0.3)
            self.active_texts.clear()


class NeurotransmitterScene(Scene):
    def construct(self):
        self.tm = TextManager(self)

        self.show_intro()
        self.show_overview()
        self.show_dopamine()
        self.show_serotonin()
        self.show_norepinephrine()
        self.show_comparison()
        self.wait(2)

    def show_intro(self):
        title = Text("Neurotransmitter Systems", font_size=44)
        subtitle = Text("The brain's chemical messengers", font_size=24, color=GRAY)
        subtitle.next_to(title, DOWN, buff=0.3)

        self.play(Write(title))
        self.play(FadeIn(subtitle))
        self.wait(1)

        intro = Text(
            "Neurotransmitters are chemicals that transmit signals\\n"
            "between neurons, controlling mood, movement, and cognition",
            font_size=18, color=BLUE_B
        )
        intro.next_to(subtitle, DOWN, buff=0.5)
        self.play(FadeIn(intro))
        self.wait(2)
        self.play(FadeOut(title), FadeOut(subtitle), FadeOut(intro))

    def show_overview(self):
        self.tm.show("Three major neurotransmitter systems regulate brain function", region="bottom")

        # Three colored boxes
        systems = VGroup()
        colors = [("#FF6B6B", "Dopamine\\n(DA)", "Reward, Movement"),
                  ("#4ECDC4", "Serotonin\\n(5-HT)", "Mood, Sleep"),
                  ("#FFE66D", "Norepinephrine\\n(NE)", "Alertness, Stress")]

        for i, (color, name, function) in enumerate(colors):
            box = RoundedRectangle(width=3.5, height=2, corner_radius=0.2,
                                   fill_opacity=0.7, fill_color=color, stroke_color=WHITE)
            box.shift(LEFT * 4 + RIGHT * i * 4)

            title = Text(name, font_size=18, color=WHITE)
            title.move_to(box.get_center() + UP * 0.3)

            func = Text(function, font_size=14, color=WHITE)
            func.move_to(box.get_center() + DOWN * 0.4)

            systems.add(VGroup(box, title, func))

        self.play(Create(systems), run_time=2)
        self.wait(2)
        self.play(FadeOut(systems))

    def show_dopamine(self):
        self.tm.show("DOPAMINE: The reward and movement neurotransmitter", region="top", color="#FF6B6B")

        # Brain outline
        brain = Ellipse(width=6, height=5, color=PINK, fill_opacity=0.2)

        # VTA (source)
        vta = Dot(DOWN * 1.5, color="#FF6B6B", radius=0.3)
        vta_label = Text("VTA\\n(Source)", font_size=12, color="#FF6B6B")
        vta_label.next_to(vta, DOWN, buff=0.1)

        # SNc
        snc = Dot(DOWN * 1.5 + LEFT * 1, color="#FF6B6B", radius=0.25)
        snc_label = Text("SNc", font_size=10, color="#FF6B6B")
        snc_label.next_to(snc, LEFT, buff=0.1)

        # Target regions
        pfc = Dot(UP * 1.5 + LEFT * 1, color=BLUE, radius=0.25)
        pfc_label = Text("PFC\\n(Decision)", font_size=10, color=BLUE)
        pfc_label.next_to(pfc, UP, buff=0.1)

        nac = Dot(UP * 0.5, color=GREEN, radius=0.25)
        nac_label = Text("NAc\\n(Reward)", font_size=10, color=GREEN)
        nac_label.next_to(nac, RIGHT, buff=0.1)

        striatum = Dot(UP * 0.5 + LEFT * 1.5, color=YELLOW, radius=0.25)
        striatum_label = Text("Striatum\\n(Movement)", font_size=10, color=YELLOW)
        striatum_label.next_to(striatum, LEFT, buff=0.1)

        self.play(Create(brain))
        self.play(Create(vta), Write(vta_label), Create(snc), Write(snc_label))

        self.tm.show("Dopamine neurons in VTA and SNc project to multiple brain regions", region="bottom")

        # Pathways
        path1 = Arrow(vta.get_center(), pfc.get_center(), color="#FF6B6B", stroke_width=3, buff=0.3)
        path2 = Arrow(vta.get_center(), nac.get_center(), color="#FF6B6B", stroke_width=3, buff=0.3)
        path3 = Arrow(snc.get_center(), striatum.get_center(), color="#FF6B6B", stroke_width=3, buff=0.3)

        self.play(Create(pfc), Write(pfc_label), Create(path1))
        self.play(Create(nac), Write(nac_label), Create(path2))
        self.play(Create(striatum), Write(striatum_label), Create(path3))

        self.tm.show("Dopamine controls: reward seeking, motivation, and motor control", region="bottom", color="#FF6B6B")
        self.wait(2)

        self.play(FadeOut(*self.mobjects[:-1]))  # Keep TextManager text

    def show_serotonin(self):
        self.tm.show("SEROTONIN: The mood and well-being neurotransmitter", region="top", color="#4ECDC4")

        brain = Ellipse(width=6, height=5, color=PINK, fill_opacity=0.2)

        # Raphe nuclei (source)
        raphe = Dot(DOWN * 1.8, color="#4ECDC4", radius=0.3)
        raphe_label = Text("Raphe Nuclei\\n(Source)", font_size=12, color="#4ECDC4")
        raphe_label.next_to(raphe, DOWN, buff=0.1)

        self.play(Create(brain))
        self.play(Create(raphe), Write(raphe_label))

        self.tm.show("Serotonin projects WIDELY throughout the entire brain", region="bottom")

        # Wide projections
        targets = [UP * 1.5, UP * 1.5 + LEFT * 1.5, UP * 1.5 + RIGHT * 1.5,
                   UP * 0.5 + LEFT * 2, UP * 0.5 + RIGHT * 2, ORIGIN + LEFT * 1.5]

        paths = VGroup()
        dots = VGroup()
        for t in targets:
            dot = Dot(t, color="#4ECDC4", radius=0.15)
            path = Arrow(raphe.get_center(), t, color="#4ECDC4", stroke_width=2, buff=0.2)
            paths.add(path)
            dots.add(dot)

        self.play(Create(paths), Create(dots), run_time=2)

        self.tm.show("Serotonin regulates: mood, sleep, appetite, and anxiety", region="bottom", color="#4ECDC4")
        self.wait(2)

        self.play(FadeOut(*self.mobjects[:-1]))

    def show_norepinephrine(self):
        self.tm.show("NOREPINEPHRINE: The alertness and stress neurotransmitter", region="top", color="#FFE66D")

        brain = Ellipse(width=6, height=5, color=PINK, fill_opacity=0.2)

        # Locus coeruleus (source)
        lc = Dot(DOWN * 1.2 + RIGHT * 0.5, color="#FFE66D", radius=0.3)
        lc_label = Text("Locus Coeruleus\\n(Source)", font_size=12, color="#FFE66D")
        lc_label.next_to(lc, DOWN, buff=0.1)

        self.play(Create(brain))
        self.play(Create(lc), Write(lc_label))

        self.tm.show("The locus coeruleus is tiny but projects everywhere", region="bottom")

        # Widespread projections
        targets = [UP * 1.8, UP * 1.5 + LEFT * 2, UP * 1.5 + RIGHT * 2,
                   UP * 0.5 + LEFT * 2.5, UP * 0.5 + RIGHT * 2.5,
                   ORIGIN + LEFT * 2, DOWN * 0.5 + RIGHT * 2]

        paths = VGroup()
        for t in targets:
            path = Arrow(lc.get_center(), t, color="#FFE66D", stroke_width=2, buff=0.2)
            paths.add(path)

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

        # Stress response indicator
        alert = Text("!", font_size=36, color=RED)
        alert.move_to(UP * 2)
        self.play(FadeIn(alert, scale=2))

        self.tm.show("NE triggers fight-or-flight: increased heart rate, alertness, focus", region="bottom", color="#FFE66D")
        self.wait(2)

        self.play(FadeOut(*self.mobjects[:-1]))

    def show_comparison(self):
        self.tm.clear()

        # Comparison table
        title = Text("Neurotransmitter Comparison", font_size=28)
        title.to_edge(UP, buff=0.5)

        table_data = [
            ["System", "Source", "Function", "Disorders"],
            ["Dopamine", "VTA, SNc", "Reward, Movement", "Parkinson's, Addiction"],
            ["Serotonin", "Raphe", "Mood, Sleep", "Depression, Anxiety"],
            ["Norepinephrine", "LC", "Alertness", "PTSD, ADHD"],
        ]

        table = VGroup()
        colors = [WHITE, "#FF6B6B", "#4ECDC4", "#FFE66D"]

        for i, row in enumerate(table_data):
            row_group = VGroup()
            for j, cell in enumerate(row):
                text = Text(cell, font_size=14 if i == 0 else 12,
                           color=colors[i] if j == 0 else WHITE)
                text.move_to(LEFT * 4.5 + RIGHT * j * 3 + UP * 1 + DOWN * i * 0.8)
                row_group.add(text)
            table.add(row_group)

        self.play(Write(title))
        self.play(Write(table), run_time=2)

        final = Text("These systems interact and balance each other for healthy brain function", font_size=16, color=BLUE_B)
        final.to_edge(DOWN, buff=0.8)
        self.play(FadeIn(final))


if __name__ == "__main__":
    pass
'''


# Registry of all templates
TEMPLATES = {
    "binary_search": BinarySearchTemplate(),
    "search_visual": BinarySearchTemplate(),  # alias
    "sorting": SortingTemplate(),
    "sort_visual": SortingTemplate(),  # alias
    "bubble_sort": SortingTemplate(),
    "pythagorean": PythagoreanTemplate(),
    "proof_steps": PythagoreanTemplate(),  # alias
    "tree_traversal": TreeTraversalTemplate(),
    # Neuroscience templates
    "action_potential": ActionPotentialTemplate(),
    "synapse": SynapseTemplate(),
    "neuron_structure": ActionPotentialTemplate(),  # alias
    "motor_cortex_bci": MotorCortexBCITemplate(),
    "motor_cortex": MotorCortexBCITemplate(),  # alias
    "bci": MotorCortexBCITemplate(),  # alias
    "neural_decoding": MotorCortexBCITemplate(),  # alias
    "neurotransmitter": NeurotransmitterTemplate(),
    "dopamine": NeurotransmitterTemplate(),  # alias
    "serotonin": NeurotransmitterTemplate(),  # alias
    "norepinephrine": NeurotransmitterTemplate(),  # alias
}


def generate_manim_code(
    plan: Dict[str, Any],
    template_name: Optional[str] = None,
    params: Optional[Dict[str, Any]] = None
) -> str:
    """
    Generate Manim code from a visualization plan.

    Args:
        plan: Visualization plan from analyzer
        template_name: Override template selection
        params: Parameters for template

    Returns:
        Generated Manim Python code
    """
    # Determine template
    template_key = template_name or plan.get("template")

    if template_key and template_key in TEMPLATES:
        template = TEMPLATES[template_key]
        return template.generate(params or {})

    # Fallback: generate generic scene from plan
    return generate_generic_scene(plan)


def generate_generic_scene(plan: Dict[str, Any]) -> str:
    """
    Generate a generic Manim scene when no template matches.
    """
    title = plan.get("title", "Visualization")
    scenes = plan.get("scenes", [])

    scene_methods = []
    for scene in scenes:
        method = f'''
    def scene_{scene["id"]}_{scene["name"]}(self):
        """
        {scene["description"]}
        Duration: {scene["duration"]}s
        """
        # TODO: Implement {scene["name"]}
        text = Text("{scene["description"][:40]}...", font_size=24)
        self.play(Write(text))
        self.wait({scene["duration"]})
        self.play(FadeOut(text))
'''
        scene_methods.append(method)

    scene_calls = "\n        ".join(
        [f"self.scene_{s['id']}_{s['name']}()" for s in scenes]
    )

    return f'''"""
{title}
Generated by LearnViz

This is a template - fill in the scene implementations.
"""
from manim import *

class GeneratedScene(Scene):
    def construct(self):
        # Run all scenes
        {scene_calls}

{"".join(scene_methods)}

if __name__ == "__main__":
    # Render with: manim -pql generated_scene.py GeneratedScene
    pass
'''


# CLI interface
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Available templates:")
        for name, template in TEMPLATES.items():
            print(f"  {name}: {template.description}")
        print("\nUsage: python manim_gen.py <template_name> [params_json]")
        sys.exit(1)

    template_name = sys.argv[1]
    params = json.loads(sys.argv[2]) if len(sys.argv) > 2 else {}

    if template_name not in TEMPLATES:
        print(f"Unknown template: {template_name}")
        sys.exit(1)

    code = TEMPLATES[template_name].generate(params)
    print(code)
