File size: 11,626 Bytes
6fc3143
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
"""
Code Post-Processor for Manim

Fixes common issues in AI-generated Manim code, particularly around
SurroundingRectangle with indexed MathTex elements.
"""

import re
from typing import List


def fix_surrounding_rectangles(code: str) -> str:
    """
    Fix or remove SurroundingRectangle calls that use indexed access to MathTex.
    
    Problem: AI often generates code like:
        SurroundingRectangle(equation[0][5], ...)
    
    This doesn't work reliably because MathTex indexing is unpredictable.
    
    Solution: Comment out these problematic lines with an explanation.
    
    Args:
        code: Raw generated Manim code
    
    Returns:
        Fixed code with problematic SurroundingRectangle calls commented out
    """
    lines = code.split('\n')
    fixed_lines = []
    skip_next = False  # Initialize before the loop
    
    for line in lines:
        # Pattern: SurroundingRectangle with indexed access like equation[0][5]
        # Matches: SurroundingRectangle(something[X][Y], ...) or SurroundingRectangle(something[X:Y], ...)
        if 'SurroundingRectangle(' in line:
            # Check if it has indexed or sliced access
            # Pattern: variable[number] or variable[number:number] inside SurroundingRectangle()
            if re.search(r'SurroundingRectangle\([^,\)]*\[\d+\]', line) or \
               re.search(r'SurroundingRectangle\([^,\)]*\[\d+:\d*\]', line):
                # Replace with empty VGroup to prevent NameError in subsequent animations
                # while avoiding the rendering crash
                indent = len(line) - len(line.lstrip())
                # Extract variable name if present (e.g. "box = SurroundingRectangle(...)")
                var_name = ""
                if "=" in line:
                    var_name = line.split("=")[0].strip()
                    replacement = ' ' * indent + f'{var_name} = VGroup()  # Replaced invalid SurroundingRectangle'
                else:
                    # If no assignment, just comment it out as it won't be referenced
                    replacement = ' ' * indent + '# ' + line.lstrip() + '  # Auto-disabled: indexed SurroundingRectangle'
                
                fixed_lines.append(replacement)
            else:
                fixed_lines.append(line)
        else:
            fixed_lines.append(line)
    
    return '\n'.join(fixed_lines)


def remove_problematic_indexing(code: str) -> str:
    """
    More aggressive approach: Remove entire blocks that use indexed MathTex highlighting.
    
    This removes the entire voiceover block if it contains indexed SurroundingRectangle.
    """
    # For now, use the commenting approach which is safer
    return fix_surrounding_rectangles(code)


def fix_undefined_colors(code: str) -> str:
    """
    Fix undefined color constants by replacing them with valid Manim colors.
    
    Common issues:
    - ORANGE_A, ORANGE_B, etc. -> ORANGE
    - RED_A, RED_B, etc. -> RED
    - Similar patterns for other colors
    
    Args:
        code: Raw generated code
    
    Returns:
        Code with undefined colors replaced
    """
    # Color mappings: undefined variants -> standard colors
    color_replacements = {
        # Orange variants
        r'\bORANGE_[A-Z]\b': 'ORANGE',
        # Red variants
        r'\bRED_[A-Z]\b': 'RED',
        # Blue variants
        r'\bBLUE_[A-Z]\b': 'BLUE',
        # Green variants
        r'\bGREEN_[A-Z]\b': 'GREEN',
        # Yellow variants
        r'\bYELLOW_[A-Z]\b': 'YELLOW',
        # Purple variants
        r'\bPURPLE_[A-Z]\b': 'PURPLE',
        # Pink variants
        r'\bPINK_[A-Z]\b': 'PINK',
        # Teal variants
        r'\bTEAL_[A-Z]\b': 'TEAL',
        # Gray variants
        r'\bGRAY_[A-Z]\b': 'GRAY',
    }
    
    for pattern, replacement in color_replacements.items():
        code = re.sub(pattern, replacement, code)
    
    return code


def clean_duplicate_imports(code: str) -> str:
    """
    Remove duplicate and incorrect imports.
    Keep only the correct imports for our architecture.
    """
    lines = code.split('\n')
    cleaned_lines = []
    seen_imports = set()
    
    # Imports to remove (old/wrong patterns)
    bad_imports = [
        'from manim_voiceover import VoiceoverScene',
        'from manimator.services import ElevenLabsService',
        'from manim_voiceover.services.gtts import GTTSService',
        'from manim_voiceover.services.elevenlabs import ElevenLabsService',
    ]
    
    for line in lines:
        stripped = line.strip()
        
        # Skip bad imports
        if any(bad_import in line for bad_import in bad_imports):
            continue
        
        # Track and deduplicate good imports
        if stripped.startswith('from ') or stripped.startswith('import '):
            if stripped not in seen_imports:
                seen_imports.add(stripped)
                cleaned_lines.append(line)
        else:
            cleaned_lines.append(line)
    
    return '\n'.join(cleaned_lines)


def fix_invalid_manim_parameters(code: str) -> str:
    """
    Fix invalid Manim parameters that cause runtime errors.
    
    Common issues:
    - rounded_corners parameter in Rectangle (doesn't exist)
    - corner_radius parameter in Rectangle (doesn't exist in older Manim versions)
    - Invalid parameter names
    """
    # Remove corner_radius from Rectangle/RoundedRectangle calls (line by line)
    lines = code.split('\n')
    fixed_lines = []
    
    for line in lines:
        # Skip lines that are just the parameter
        if 'corner_radius' in line and '=' in line:
            # Check if this is a standalone parameter line
            stripped = line.strip()
            if stripped.startswith('corner_radius=') or stripped.startswith('rounded_corners='):
                # Skip this line entirely
                continue
        
        # Remove inline corner_radius/rounded_corners parameters
        line = re.sub(r',\s*corner_radius\s*=\s*[^,)]+', '', line)
        line = re.sub(r',\s*rounded_corners\s*=\s*[^,)]+', '', line)
        
        # Remove invalid scale_tips parameter (often hallucinated for scale() method)
        if 'scale_tips' in line and '=' in line:
            stripped = line.strip()
            if stripped.startswith('scale_tips='):
                continue
                
        # Handle inline scale_tips removal (both with and without leading comma)
        line = re.sub(r',\s*scale_tips\s*=\s*[^,)]+', '', line)
        line = re.sub(r'\bscale_tips\s*=\s*[^,)]+\s*,?', '', line)
        
        fixed_lines.append(line)
    
    return '\n'.join(fixed_lines)


def fix_visual_leaks(code: str) -> str:
    """
    Prevent 'visual memory leaks' where reassigning a variable leaves the old Mobject on screen.
    Also enforces 'Title Exclusivity' to prevent multiple titles from stacking.
    
    Problem 1 (Variable Reuse):
        text = Text("A")
        self.play(Write(text))
        text = Text("B") # "A" is still on screen!
        
    Problem 2 (Title Stacking):
        title1 = Text("Intro")
        self.play(Write(title1))
        title2 = Text("Chapter 1") # "Intro" is still on screen!
        
    Solution:
        Inject `self.remove(var)` checks.
    """
    lines = code.split('\n')
    fixed_lines = []
    
    # Regex to detect Mobject assignments
    # Matches: var = ClassName(...)
    # We broaden this to catch any Capitalized class instantiation to be safer
    assignment_pattern = re.compile(r'^\s*([a-zA-Z_]\w*)\s*=\s*([A-Z]\w*)\(')
    
    # Track variables that have been assigned
    assigned_vars = set()
    # Track variables that look like titles
    active_titles = []
    
    for line in lines:
        match = assignment_pattern.match(line)
        if match:
            var_name = match.group(1)
            class_name = match.group(2)
            
            indent = len(line) - len(line.lstrip())
            indent_str = ' ' * indent
            
            # 1. Handle Variable Reuse (Same variable name)
            # We inject this for EVERY assignment to handle loops correctly.
            # In a loop, the line `t = Text(...)` appears once but runs multiple times.
            # By injecting the check, we ensure the previous iteration's object is removed.
            removal = f"{indent_str}if '{var_name}' in locals(): self.remove({var_name})"
            fixed_lines.append(removal)
            
            # 2. Handle Title Exclusivity (Different variable names, but both are titles)
            # Check if this looks like a title (contains 'title', 'header', 'heading')
            # But exclude 'subtitle' or 'sub_title'
            is_title = re.search(r'title|header|heading', var_name, re.IGNORECASE) and \
                       not re.search(r'sub', var_name, re.IGNORECASE)
            
            if is_title:
                # If we are creating a new title, remove ALL previous titles
                for old_title in active_titles:
                    # Don't remove if it's the same variable (handled above)
                    if old_title != var_name:
                        removal = f"{indent_str}if '{old_title}' in locals(): self.remove({old_title})"
                        fixed_lines.append(removal)
                
                # Add this to active titles (if not already there)
                if var_name not in active_titles:
                    active_titles.append(var_name)
            
            assigned_vars.add(var_name)
            fixed_lines.append(line)
        else:
            fixed_lines.append(line)
            
    return '\n'.join(fixed_lines)


def post_process_code(code: str) -> str:
    """
    Main entry point for code post-processing.
    
    Applies all fixes to AI-generated Manim code.
    
    Args:
        code: Raw generated code
    
    Returns:
        Cleaned and fixed code
    """
    # Check if we need to add header (before making changes)
    has_undefined_colors = bool(re.search(r'\b(ORANGE|RED|BLUE|GREEN|YELLOW|PURPLE|PINK|TEAL|GRAY)_[A-Z]\b', code))
    
    # Apply fixes in order
    code = clean_duplicate_imports(code)
    code = fix_undefined_colors(code)
    code = fix_invalid_manim_parameters(code)
    code = fix_invalid_manim_parameters(code)
    code = fix_surrounding_rectangles(code)
    code = fix_visual_leaks(code)
    
    # Add header comment explaining post-processing
    header = """# NOTE: This code has been automatically post-processed to fix common issues.
# Indexed SurroundingRectangle calls have been disabled as they don't reliably
# highlight the intended equation parts in MathTex objects.
# Undefined color constants have been replaced with standard Manim colors.
# Invalid Manim parameters have been removed or corrected.

"""
    
    # Only add header if we actually made changes
    if '# Auto-disabled:' in code or has_undefined_colors or 'rounded_corners' in code:
        code = header + code
    
    return code


def validate_code_structure(code: str) -> List[str]:
    """
    Validate the generated code for common issues.
    
    Returns:
        List of warning messages (empty if no issues)
    """
    warnings = []
    
    # Check for common issues
    if 'SurroundingRectangle(' in code:
        if re.search(r'SurroundingRectangle\([^,\)]*\[\d+\]', code):
            warnings.append("Code contains indexed SurroundingRectangle calls (will be auto-fixed)")
    
    if 'from manim_voiceover.services.gtts import GTTSService' in code:
        warnings.append("Code still uses deprecated GTTSService (should use ElevenLabsService)")
    
    return warnings