Update app.py

app.py

@@ -52,7 +52,7 @@ MODEL_CONFIGS = {
     "gpt-4.1": {"max_tokens": 32768, "param_name": "max_tokens", "api_version": None, "category": "OpenAI", "warning": None},
     "gpt-4.1-mini": {"max_tokens": 32768, "param_name": "max_tokens", "api_version": None, "category": "OpenAI", "warning": None},
     "gpt-4.1-nano": {"max_tokens": 32768, "param_name": "max_tokens", "api_version": None, "category": "OpenAI", "warning": None},
-    "o3-mini": {"max_completion_tokens": 100000, "param_name": "max_completion_tokens", "api_version": "2024-12-01-preview", "category": "OpenAI", "warning": None},
+    "o4-mini": {"max_completion_tokens": 100000, "param_name": "max_completion_tokens", "api_version": "2024-12-01-preview", "category": "OpenAI", "warning": None},
     "o1": {"max_completion_tokens": 100000, "param_name": "max_completion_tokens", "api_version": "2024-12-01-preview", "category": "OpenAI", "warning": None},
     "o1-mini": {"max_completion_tokens": 66000, "param_name": "max_completion_tokens", "api_version": "2024-12-01-preview", "category": "OpenAI", "warning": None},
     "o1-preview": {"max_tokens": 33000, "param_name": "max_tokens", "api_version": None, "category": "OpenAI", "warning": None},
@@ -74,29 +74,33 @@ except ImportError:
 def prepare_api_params(messages, model_name):
     """Create appropriate API parameters based on model configuration"""
     config = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["default"])
-    api_params = {"messages": messages, "model": model_name}
+    api_params = {
+        "messages": messages,
+        "model": model_name
+    }
     token_param = config["param_name"]
-    token_value = config[token_param]
-    api_params[token_param] = token_value
+    api_params[token_param] = config.get(token_param)
     return api_params, config
 
-def get_secret(key):
-    """Retrieve a secret from environment or Streamlit secrets."""
-    if hasattr(st, "secrets") and key in st.secrets:
-        return st.secrets[key]
-    return os.environ.get(key)
+def get_secret(env_var):
+    """Retrieve a secret from environment variables"""
+    val = os.environ.get(env_var)
+    if not val:
+        logger.warning(f"Secret '{env_var}' not found")
+    return val
 
 def check_password():
-    correct_password = get_secret("password")
-    if not correct_password:
-        st.error("Admin password not configured in secrets or env var 'password'")
+    """Verify password entered against secret"""
+    correct = get_secret("password")
+    if not correct:
+        st.error("Admin password not configured")
         return False
     if "password_entered" not in st.session_state:
         st.session_state.password_entered = False
     if not st.session_state.password_entered:
         pwd = st.text_input("Enter password to access AI features", type="password")
         if pwd:
-            if pwd == correct_password:
+            if pwd == correct:
                 st.session_state.password_entered = True
                 return True
             else:
@@ -106,7 +110,7 @@ def check_password():
     return True
 
 def ensure_packages():
-    required_packages = {
+    required = {
         'manim': '0.17.3',
         'Pillow': '9.0.0',
         'numpy': '1.22.0',
@@ -126,115 +130,69 @@ def ensure_packages():
         'huggingface_hub': '0.16.0',
     }
     missing = {}
-    for pkg, ver in required_packages.items():
+    for pkg, ver in required.items():
         try:
             __import__(pkg if pkg != 'Pillow' else 'PIL')
         except ImportError:
             missing[pkg] = ver
     if not missing:
         return True
-    progress = st.progress(0)
-    status = st.empty()
+    bar = st.progress(0)
+    txt = st.empty()
     for i, (pkg, ver) in enumerate(missing.items()):
-        status.text(f"Installing {pkg}...")
+        bar.progress(i / len(missing))
+        txt.text(f"Installing {pkg}...")
         res = subprocess.run([sys.executable, "-m", "pip", "install", f"{pkg}>={ver}"], capture_output=True, text=True)
         if res.returncode != 0:
-            st.error(f"Failed to install {pkg}: {res.stderr}")
+            st.error(f"Failed to install {pkg}")
             return False
-        progress.progress((i + 1) / len(missing))
+    bar.progress(1.0)
+    txt.empty()
     return True
 
+def install_custom_packages(pkgs):
+    if not pkgs.strip():
+        return True, "No packages specified"
+    parts = [p.strip() for p in pkgs.split(",") if p.strip()]
+    if not parts:
+        return True, "No valid packages"
+    sidebar_txt = st.sidebar.empty()
+    bar = st.sidebar.progress(0)
+    results = []
+    success = True
+    for i, p in enumerate(parts):
+        bar.progress(i / len(parts))
+        sidebar_txt.text(f"Installing {p}...")
+        res = subprocess.run([sys.executable, "-m", "pip", "install", p], capture_output=True, text=True)
+        if res.returncode != 0:
+            results.append(f"Failed {p}: {res.stderr}")
+            success = False
+        else:
+            results.append(f"Installed {p}")
+    bar.progress(1.0)
+    sidebar_txt.empty()
+    return success, "\n".join(results)
+
 @st.cache_resource(ttl=3600)
 def init_ai_models_direct():
+    token = get_secret("github_token_api")
+    if not token:
+        st.error("API token not configured")
+        return None
     try:
-        token = get_secret("github_token_api")
-        if not token:
-            st.error("GitHub token not found in secrets or env var 'github_token_api'")
-            return None
         from azure.ai.inference import ChatCompletionsClient
-        from azure.ai.inference.models import SystemMessage, UserMessage
+        from azure.ai.inference.models import UserMessage
         from azure.core.credentials import AzureKeyCredential
-        endpoint = "https://models.inference.ai.azure.com"
-        model_name = "gpt-4o"
-        client = ChatCompletionsClient(endpoint=endpoint, credential=AzureKeyCredential(token))
-        return {
-            "client": client,
-            "model_name": model_name,
-            "endpoint": endpoint,
-            "last_loaded": datetime.now().isoformat(),
-            "category": MODEL_CONFIGS[model_name]["category"],
-            "api_version": MODEL_CONFIGS[model_name].get("api_version")
-        }
-    except Exception as e:
-        st.error(f"Error initializing AI model: {e}")
+        client = ChatCompletionsClient(
+            endpoint="https://models.inference.ai.azure.com",
+            credential=AzureKeyCredential(token)
+        )
+        return {"client": client, "model_name": "gpt-4o", "last_loaded": datetime.now().isoformat()}
+    except ImportError as e:
+        st.error("Azure AI SDK not installed")
         logger.error(str(e))
         return None
 
-def suggest_code_completion(code_snippet, models):
-    if not models:
-        st.error("AI models not initialized")
-        return None
-    try:
-        prompt = f"""Write a complete Manim animation scene based on this code or idea:
-{code_snippet}
-
-The code should be a complete, working Manim animation that includes:
-- Proper Scene class definition
-- Constructor with animations
-- Proper use of self.play() for animations
-- Proper wait times between animations
-
-Here's the complete Manim code:
-"""
-        from openai import OpenAI
-        token = get_secret("github_token_api")
-        client = OpenAI(base_url="https://models.github.ai/inference", api_key=token)
-        messages = [{"role": "system", "content": "You are an expert in Manim animations."},
-                    {"role": "user", "content": prompt}]
-        config = MODEL_CONFIGS.get(models["model_name"], MODEL_CONFIGS["default"])
-        params = {"messages": messages, "model": models["model_name"], config["param_name"]: config[config["param_name"]]}
-        response = client.chat.completions.create(**params)
-        content = response.choices[0].message.content
-        if "```python" in content:
-            content = content.split("```python")[1].split("```")[0]
-        elif "```" in content:
-            content = content.split("```")[1].split("```")[0]
-        if "Scene" not in content:
-            content = f"from manim import *\n\nclass MyScene(Scene):\n    def construct(self):\n        {content}"
-        return content
-    except Exception as e:
-        st.error(f"Error generating code: {e}")
-        logger.error(traceback.format_exc())
-        return None
-
-QUALITY_PRESETS = {
-    "480p": {"resolution": "480p", "fps": "30"},
-    "720p": {"resolution": "720p", "fps": "30"},
-    "1080p": {"resolution": "1080p", "fps": "60"},
-    "4K": {"resolution": "2160p", "fps": "60"},
-    "8K": {"resolution": "4320p", "fps": "60"}
-}
-
-ANIMATION_SPEEDS = {
-    "Slow": 0.5,
-    "Normal": 1.0,
-    "Fast": 2.0,
-    "Very Fast": 3.0
-}
-
-EXPORT_FORMATS = {
-    "MP4 Video": "mp4",
-    "GIF Animation": "gif",
-    "WebM Video": "webm",
-    "PNG Image Sequence": "png_sequence",
-    "SVG Image": "svg"
-}
-
-def highlight_code(code):
-    formatter = HtmlFormatter(style='monokai')
-    highlighted = highlight(code, PythonLexer(), formatter)
-    return highlighted, formatter.get_style_defs()
-
 def generate_manim_preview(python_code):
     scene_objects = []
     if "Circle" in python_code: scene_objects.append("circle")
@@ -242,657 +200,423 @@ def generate_manim_preview(python_code):
     if "MathTex" in python_code or "Tex" in python_code: scene_objects.append("equation")
     if "Text" in python_code: scene_objects.append("text")
     if "Axes" in python_code: scene_objects.append("graph")
-    if "ThreeDScene" in python_code or "ThreeDAxes" in python_code: scene_objects.append("3D scene")
-    if "Sphere" in python_code: scene_objects.append("sphere")
-    if "Cube" in python_code: scene_objects.append("cube")
-    icons = {"circle":"⭕","square":"🔲","equation":"📊","text":"📝","graph":"📈","3D scene":"🧊","sphere":"🌐","cube":"🧊"}
-    icon_html = "".join(f'<span style="font-size:2rem; margin:0.3rem;">{icons[o]}</span>' for o in scene_objects)
-    preview_html = f"""
-    <div style="background-color:#000; width:100%; height:220px; border-radius:10px; display:flex; flex-direction:column; align-items:center; justify-content:center; color:white; text-align:center;">
-        <h3>Animation Preview</h3>
-        <div>{icon_html if icon_html else '<span style="font-size:2rem;">🎬</span>'}</div>
-        <p>Scene contains: {', '.join(scene_objects) if scene_objects else 'No detected objects'}</p>
-        <p style="font-size:0.8rem; opacity:0.7;">Full rendering required for accurate preview</p>
-    </div>
-    """
-    return preview_html
-
-def render_latex_preview(latex):
-    if not latex:
-        return """
-        <div style="background:#f8f9fa; width:100%; height:100px; border-radius:5px; display:flex; align-items:center; justify-content:center; color:#6c757d;">
-            Enter LaTeX formula to see preview
-        </div>
-        """
-    return f"""
-    <div style="background:#202124; width:100%; padding:20px; border-radius:5px; color:white; text-align:center;">
-        <script src="https://polyfill.io/v3/polyfill.min.js?features=es6"></script>
-        <script id="MathJax-script" async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
-        <div><h3>LaTeX Preview</h3><div id="math-preview">$$ {latex} $$</div><p style="font-size:0.8rem; opacity:0.7;">Use MathTex(r"{latex}") in Manim</p></div>
+    icons = {"circle":"⭕","square":"🔲","equation":"📊","text":"📝","graph":"📈"}
+    icon_html = "".join(f'<span style="font-size:2rem;margin:0.3rem;">{icons[o]}</span>' for o in scene_objects if o in icons)
+    html = f"""
+    <div style="background:#000;color:#fff;padding:1rem;border-radius:10px;text-align:center;">
+      <h3>Animation Preview</h3>
+      <div>{icon_html or '🎬'}</div>
+      <p>Contains: {', '.join(scene_objects) or 'none'}</p>
+      <p style="opacity:0.7;">Full rendering required for accurate preview</p>
     </div>
     """
+    return html
 
 def extract_scene_class_name(python_code):
-    match = re.search(r'class\s+(\w+)\s*\([^)]*Scene[^)]*\)', python_code)
-    return match.group(1) if match else "MyScene"
-
-def prepare_audio_for_manim(audio_file, target_dir):
-    audio_dir = os.path.join(target_dir, "audio")
-    os.makedirs(audio_dir, exist_ok=True)
-    filename = f"audio_{int(time.time())}.mp3"
-    path = os.path.join(audio_dir, filename)
-    with open(path, "wb") as f: f.write(audio_file.getvalue())
-    return path
+    names = re.findall(r'class\s+(\w+)\s*\([^)]*Scene', python_code)
+    return names[0] if names else "MyScene"
 
 def mp4_to_gif(mp4, out, fps=15):
-    cmd = ["ffmpeg","-i",mp4,"-vf",f"fps={fps},scale=640:-1:flags=lanczos,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse","-loop","0",out]
-    res = subprocess.run(cmd,capture_output=True,text=True)
-    return out if res.returncode==0 else None
-
-def generate_manim_video(code, fmt, quality, speed, audio_path=None):
-    temp_dir = tempfile.mkdtemp(prefix="manim_render_")
-    try:
-        scene = extract_scene_class_name(code)
-        if audio_path and "with_sound" not in code:
-            code = "from manim.scene.scene_file_writer import SceneFileWriter\n" + code
-            pat = re.search(f"class {scene}\\(.*?\\):", code)
-            if pat:
-                decor = f"@with_sound(\"{audio_path}\")\n"
-                code = code[:pat.start()] + decor + code[pat.start():]
-        path_py = os.path.join(temp_dir, "scene.py")
-        with open(path_py, "w", encoding="utf-8") as f: f.write(code)
-        qmap = {"480p":"-ql","720p":"-qm","1080p":"-qh","4K":"-qk","8K":"-qp"}
-        qflag = qmap.get(quality,"-qm")
-        if fmt=="png_sequence":
-            farg="--format=png"; extra=["--save_pngs"]
-        elif fmt=="svg":
-            farg="--format=svg"; extra=[]
-        else:
-            farg=f"--format={fmt}"; extra=[]
-        cmd = ["manim", path_py, scene, qflag, farg] + extra
-        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
-        output=[]
-        out_path=None; mp4_path=None
-        while True:
-            line = proc.stdout.readline()
-            if not line and proc.poll() is not None: break
+    cmd = [
+        "ffmpeg","-i",mp4,
+        "-vf",f"fps={fps},scale=640:-1:flags=lanczos,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse",
+        "-loop","0",out
+    ]
+    r = subprocess.run(cmd, capture_output=True, text=True)
+    return out if r.returncode==0 else None
+
+def generate_manim_video(code, format_type, quality_preset, speed=1.0, audio_path=None):
+    temp_dir = tempfile.mkdtemp(prefix="manim_")
+    scene_class = extract_scene_class_name(code)
+    file_py = os.path.join(temp_dir, "scene.py")
+    with open(file_py, "w", encoding="utf-8") as f:
+        f.write(code)
+    quality_flags = {"480p":"-ql","720p":"-qm","1080p":"-qh","4K":"-qk","8K":"-qp"}
+    qf = quality_flags.get(quality_preset, "-qm")
+    fmt_arg = f"--format={format_type}"
+    cmd = ["manim", file_py, scene_class, qf, fmt_arg]
+    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+    output = []
+    out_path = None
+    mp4_path = None
+    bar = st.empty()
+    log = st.empty()
+    while True:
+        line = proc.stdout.readline()
+        if not line and proc.poll() is not None:
+            break
+        if line:
             output.append(line)
-            if "%" in line:
-                try:
-                    p=float(line.split("%")[0].strip().split()[-1]); 
-                except: pass
+            log.code("".join(output[-10:]))
             if "File ready at" in line:
-                chunk = line.split("File ready at")[-1].strip()
-                m=re.search(r'([\'"]?)(.*?\.(mp4|gif|webm|svg))\1',chunk)
-                if m: 
-                    out_path=m.group(2)
-                    if out_path.endswith(".mp4"): mp4_path=out_path
-        proc.wait()
-        time.sleep(2)
-        data=None
-        if fmt=="gif" and (not out_path or not os.path.exists(out_path)) and mp4_path:
-            gif=os.path.join(temp_dir,f"{scene}_converted.gif")
-            if mp4_to_gif(mp4_path,gif): out_path=gif
-        if fmt=="png_sequence":
-            dirs=[os.path.join(temp_dir,"media","images",scene,"Animations")]
-            pngs=[]
-            for d in dirs:
-                if os.path.isdir(d):
-                    pngs+= [os.path.join(d,f) for f in os.listdir(d) if f.endswith(".png")]
-            if pngs:
-                zipf=os.path.join(temp_dir,f"{scene}_pngs.zip")
-                with zipfile.ZipFile(zipf,"w") as z:
-                    for p in pngs: z.write(p,os.path.basename(p))
-                data=open(zipf,"rb").read()
-        elif out_path and os.path.exists(out_path):
-            data=open(out_path,"rb").read()
-        else:
-            # fallback search
-            files=[]
-            for root,_,fs in os.walk(temp_dir):
-                for f in fs:
-                    if f.endswith(f".{fmt}") and "partial" not in f:
-                        files.append(os.path.join(root,f))
-            if files:
-                latest=max(files,key=os.path.getctime)
-                data=open(latest,"rb").read()
-                if fmt=="gif" and latest.endswith(".mp4"):
-                    gif=os.path.join(temp_dir,f"{scene}_converted.gif")
-                    if mp4_to_gif(latest,gif): data=open(gif,"rb").read()
-        if data:
-            size=len(data)/(1024*1024)
-            return data, f"✅ Animation generated successfully! ({size:.1f} MB)"
-        else:
-            return None, "❌ Error: No output files generated.\n" + "".join(output)[:500]
-    except Exception as e:
-        logger.error(traceback.format_exc())
-        return None, f"❌ Error: {e}"
-    finally:
-        try: shutil.rmtree(temp_dir)
-        except: pass
+                m = re.search(r'([\'"])?(.+?\.(?:mp4|gif|webm|svg))\1', line)
+                if m:
+                    out_path = m.group(2)
+                    if out_path.endswith(".mp4"):
+                        mp4_path = out_path
+    proc.wait()
+    time.sleep(1)
+    data = None
+    if format_type=="gif" and (not out_path or not os.path.exists(out_path)) and mp4_path and os.path.exists(mp4_path):
+        gif = os.path.join(temp_dir, scene_class+"_conv.gif")
+        conv = mp4_to_gif(mp4_path, gif)
+        if conv and os.path.exists(conv):
+            out_path = conv
+    if out_path and os.path.exists(out_path):
+        with open(out_path,"rb") as f: data = f.read()
+    shutil.rmtree(temp_dir)
+    if data:
+        return data, f"✅ Generated successfully ({len(data)/(1024*1024):.1f} MB)"
+    else:
+        return None, "❌ No output generated. Check logs."
 
 def detect_input_calls(code):
-    calls=[]
-    for i,line in enumerate(code.splitlines(),1):
-        if 'input(' in line and not line.strip().startswith('#'):
-            m=re.search(r'input\([\'"](.+?)[\'"]\)',line)
-            prompt=m.group(1) if m else f"Input for line {i}"
+    calls = []
+    for i, line in enumerate(code.split("\n"),1):
+        if "input(" in line and not line.strip().startswith("#"):
+            m = re.search(r'input\(["\'](.+?)["\']\)', line)
+            prompt = m.group(1) if m else f"Input at line {i}"
             calls.append({"line":i,"prompt":prompt})
     return calls
 
 def run_python_script(code, inputs=None, timeout=60):
-    result={"stdout":"","stderr":"","exception":None,"plots":[],"dataframes":[],"execution_time":0}
+    result = {"stdout":"","stderr":"","exception":None,"plots":[],"dataframes":[],"execution_time":0}
+    mod = ""
     if inputs:
-        inject = f"""
-__INPUT_VALUES={inputs}
-__INPUT_INDEX=0
+        mod = f"""
+__INPUTS={inputs}
+__IDX=0
 def input(prompt=''):
-    global __INPUT_INDEX
+    global __IDX
     print(prompt,end='')
-    if __INPUT_INDEX<len(__INPUT_VALUES):
-        v=__INPUT_VALUES[__INPUT_INDEX]; __INPUT_INDEX+=1
-        print(v); return v
-    print(); return ''
+    if __IDX<len(__INPUTS):
+        val=__INPUTS[__IDX]; __IDX+=1
+        print(val)
+        return val
+    print()
+    return ''
 """
-        code = inject + code
+    code_full = mod + code
     with tempfile.TemporaryDirectory() as td:
-        plot_dir=os.path.join(td,'plots'); os.makedirs(plot_dir,exist_ok=True)
-        stdout_f=os.path.join(td,'stdout.txt')
-        stderr_f=os.path.join(td,'stderr.txt')
-        if 'plt' in code or 'matplotlib' in code:
-            if 'import matplotlib.pyplot as plt' not in code:
-                code="import matplotlib.pyplot as plt\n"+code
-            save_plots=f"""
-import matplotlib.pyplot as plt,os
-for i,num in enumerate(plt.get_fignums()):
-    plt.figure(num).savefig(os.path.join(r'{plot_dir}','plot_{{i}}.png'))
-"""
-            code+=save_plots
-        if 'pd.' in code or 'import pandas' in code:
-            if 'import pandas as pd' not in code:
-                code="import pandas as pd\n"+code
-            dfcap=f"""
-import pandas as pd, json,os
-for name,val in globals().items():
-    if isinstance(val,pd.DataFrame):
-        info={{"name":name,"shape":val.shape,"columns":list(val.columns),"preview":val.head().to_html()}}
-        open(os.path.join(r'{td}',f'df_{{name}}.json'),'w').write(json.dumps(info))
-"""
-            code+=dfcap
-        script=os.path.join(td,'script.py')
-        open(script,'w').write(code)
+        script = os.path.join(td,"script.py")
+        with open(script,"w") as f: f.write(code_full)
+        outf = os.path.join(td,"out.txt")
+        errf = os.path.join(td,"err.txt")
         start=time.time()
         try:
-            with open(stdout_f,'w') as so, open(stderr_f,'w') as se:
-                p=subprocess.Popen([sys.executable,script],stdout=so,stderr=se,cwd=td)
-                p.wait(timeout=timeout)
+            with open(outf,"w") as o, open(errf,"w") as e:
+                proc=subprocess.Popen([sys.executable, script], stdout=o, stderr=e, cwd=td)
+                proc.wait(timeout=timeout)
         except subprocess.TimeoutExpired:
-            p.kill()
-            result["stderr"]+="\nTimeout"
-            result["exception"]="Timeout"
-            return result
+            proc.kill()
+            result["stderr"] += f"\nTimed out after {timeout}s"
+            result["exception"] = "Timeout"
         result["execution_time"]=time.time()-start
-        result["stdout"]=open(stdout_f).read()
-        result["stderr"]=open(stderr_f).read()
-        for f in sorted(os.listdir(plot_dir)):
-            if f.endswith('.png'):
-                result["plots"].append(open(os.path.join(plot_dir,f),'rb').read())
-        for f in os.listdir(td):
-            if f.startswith('df_') and f.endswith('.json'):
-                result["dataframes"].append(json.load(open(os.path.join(td,f))))
+        result["stdout"]=open(outf).read()
+        result["stderr"]+=open(errf).read()
     return result
 
-def display_python_script_results(result):
-    if not result: return
-    st.info(f"Execution completed in {result['execution_time']:.2f}s")
-    if result["exception"]:
-        st.error(f"Exception: {result['exception']}")
-    if result["stderr"]:
+def display_python_script_results(res):
+    st.info(f"Completed in {res['execution_time']:.2f}s")
+    if res["exception"]:
+        st.error(f"Exception: {res['exception']}")
+    if res["stderr"]:
         st.error("Errors:")
-        st.code(result["stderr"], language="bash")
-    if result["plots"]:
+        st.code(res["stderr"], language="bash")
+    if res["plots"]:
         st.markdown("### Plots")
-        cols=st.columns(min(3,len(result["plots"])))
-        for i,p in enumerate(result["plots"]):
+        cols = st.columns(min(3,len(res["plots"])))
+        for i,p in enumerate(res["plots"]):
             cols[i%len(cols)].image(p,use_column_width=True)
-    if result["dataframes"]:
+    if res["dataframes"]:
         st.markdown("### DataFrames")
-        for df in result["dataframes"]:
-            with st.expander(f"{df['name']} {df['shape']}"):
-                st.write(pd.read_html(df["preview"])[0])
-    if result["stdout"]:
-        st.markdown("### Stdout")
-        st.code(result["stdout"], language="bash")
-
-def parse_animation_steps(code):
-    steps=[]
-    plays=re.findall(r'self\.play\((.*?)\)',code,re.DOTALL)
-    waits=re.findall(r'self\.wait\((.*?)\)',code,re.DOTALL)
-    cum=0
-    for i,pc in enumerate(plays):
-        anims=[a.strip() for a in pc.split(',')]
-        dur=1.0
-        if i<len(waits):
-            m=re.search(r'(\d+\.?\d*)',waits[i])
-            if m: dur=float(m.group(1))
-        steps.append({"id":i+1,"type":"play","animations":anims,"duration":dur,"start_time":cum,"code":f"self.play({pc})"})
-        cum+=dur
-    return steps
-
-def generate_code_from_timeline(steps,orig):
-    m=re.search(r'(class\s+\w+\s*\([^)]*\)\s*:.*?def\s+construct\s*\(\s*self\s*\)\s*:)',orig,re.DOTALL)
-    if not m: return orig
-    header=m.group(1)
-    new=[header]
-    indent="        "
-    for s in sorted(steps,key=lambda x:x["id"]):
-        new.append(f"{indent}{s['code']}")
-        if s["duration"]>0:
-            new.append(f"{indent}self.wait({s['duration']})")
-    return "\n".join(new)
-
-def create_timeline_editor(code):
-    st.markdown("### 🎞️ Animation Timeline Editor")
-    if not code:
-        st.warning("Add animation code first")
-        return code
-    steps=parse_animation_steps(code)
-    if not steps:
-        st.warning("No steps detected")
-        return code
-    df=pd.DataFrame(steps)
-    st.markdown("#### Animation Timeline")
-    fig=px.timeline(df,x_start="start_time",x_end=df["start_time"]+df["duration"],y="id",color="type",hover_name="animations",labels={"id":"Step","start_time":"Time(s)"})
-    fig.update_layout(height=300,xaxis=dict(title="Time(s)",rangeslider_visible=True))
-    st.plotly_chart(fig,use_container_width=True)
-    sel=st.selectbox("Select Step:",options=df["id"],format_func=lambda x:f"Step {x}")
-    new_dur=st.number_input("Duration(s):",min_value=0.1,max_value=10.0,value=float(df[df["id"]==sel]["duration"].iloc[0]),step=0.1)
-    action=st.selectbox("Action:",["Update Duration","Move Up","Move Down","Delete"])
-    if st.button("Apply"):
-        idx=df[df["id"]==sel].index[0]
-        if action=="Update Duration":
-            df.at[idx,"duration"]=new_dur
-        elif action=="Move Up" and sel>1:
-            j=df[df["id"]==sel-1].index[0]
-            df.at[idx,"id"],df.at[j,"id"]=sel-1,sel
-        elif action=="Move Down" and sel<len(df):
-            j=df[df["id"]==sel+1].index[0]
-            df.at[idx,"id"],df.at[j,"id"]=sel+1,sel
-        elif action=="Delete":
-            df=df[df["id"]!=sel]
-            df["id"]=range(1,len(df)+1)
-        cum=0
-        for i in df.sort_values("id").index:
-            df.at[i,"start_time"]=cum; cum+=df.at[i,"duration"]
-        new_code=generate_code_from_timeline(df.to_dict('records'),code)
-        st.success("Timeline updated, code regenerated.")
-        return new_code
-    return code
-
-def export_to_educational_format(video_data,fmt,title,explanation,temp_dir):
-    try:
-        if fmt=="powerpoint":
-            import pptx
-            from pptx.util import Inches
-            prs=pptx.Presentation()
-            s0=prs.slides.add_slide(prs.slide_layouts[0]); s0.shapes.title.text=title; s0.placeholders[1].text="Created with Manim"
-            s1=prs.slides.add_slide(prs.slide_layouts[5]); s1.shapes.title.text="Animation"
-            vid_path=os.path.join(temp_dir,"anim.mp4"); open(vid_path,"wb").write(video_data)
-            try:
-                s1.shapes.add_movie(vid_path,Inches(1),Inches(1.5),Inches(8),Inches(4.5))
-            except:
-                thumb=os.path.join(temp_dir,"thumb.png")
-                subprocess.run(["ffmpeg","-i",vid_path,"-ss","00:00:01","-vframes","1",thumb],check=True)
-                s1.shapes.add_picture(thumb,Inches(1),Inches(1.5),Inches(8),Inches(4.5))
-            if explanation:
-                s2=prs.slides.add_slide(prs.slide_layouts[1]); s2.shapes.title.text="Explanation"; s2.placeholders[1].text=explanation
-            out=os.path.join(temp_dir,f"{title.replace(' ','_')}.pptx"); prs.save(out)
-            return open(out,"rb").read(),"pptx"
-        if fmt=="html":
-            html=f"""<!DOCTYPE html><html><head><title>{title}</title>
-<style>body{{font-family:Arial;max-width:800px;margin:auto;padding:20px}}
-.controls button{{margin-right:10px;padding:5px 10px}}</style>
-<script>window.onload=function(){{const v=document.getElementById('anim');
-document.getElementById('play').onclick=()=>v.play();
-document.getElementById('pause').onclick=()=>v.pause();
-document.getElementById('restart').onclick=()=>{{v.currentTime=0;v.play()}};
-}};</script>
-</head><body><h1>{title}</h1>
-<video id="anim" width="100%" controls><source src="data:video/mp4;base64,{base64.b64encode(video_data).decode()}" type="video/mp4"></video>
-<div class="controls"><button id="play">Play</button><button id="pause">Pause</button><button id="restart">Restart</button></div>
-<div class="explanation">{markdown.markdown(explanation)}</div>
-</body></html>"""
-            out=os.path.join(temp_dir,f"{title.replace(' ','_')}.html"); open(out,"w").write(html)
-            return open(out,"rb").read(),"html"
-        if fmt=="sequence":
-            from fpdf import FPDF
-            vid=os.path.join(temp_dir,"anim.mp4"); open(vid,"wb").write(video_data)
-            fr_dir=os.path.join(temp_dir,"frames"); os.makedirs(fr_dir,exist_ok=True)
-            subprocess.run(["ffmpeg","-i",vid,"-r","1",os.path.join(fr_dir,"frame_%03d.png")],check=True)
-            pdf=FPDF(); pdf.set_auto_page_break(True,15)
-            pdf.add_page(); pdf.set_font("Arial","B",20); pdf.cell(190,10,title,0,1,"C")
-            segs=explanation.split("##") if explanation else ["No explanation"]
-            imgs=sorted([f for f in os.listdir(fr_dir) if f.endswith(".png")])
-            for i,img in enumerate(imgs):
-                pdf.add_page(); pdf.image(os.path.join(fr_dir,img),10,10,190)
-                pdf.ln(100); pdf.set_font("Arial","B",12); pdf.cell(190,10,f"Step {i+1}",0,1)
-                pdf.set_font("Arial","",10); pdf.multi_cell(190,5,segs[min(i,len(segs)-1)].strip())
-            out=os.path.join(temp_dir,f"{title.replace(' ','_')}_seq.pdf"); pdf.output(out)
-            return open(out,"rb").read(),"pdf"
-    except Exception as e:
-        logger.error(traceback.format_exc())
-    return None,None
-
+        for df in res["dataframes"]:
+            with st.expander(f"{df['name']} ({df['shape'][0]}×{df['shape'][1]})"):
+                st.markdown(df["preview_html"], unsafe_allow_html=True)
+    if res["stdout"]:
+        st.markdown("### Output")
+        st.code(res["stdout"], language="bash")
+
+# Main app
 def main():
     if 'init' not in st.session_state:
-        st.session_state.init=True
-        st.session_state.video_data=None
-        st.session_state.status=None
-        st.session_state.ai_models=None
-        st.session_state.generated_code=""
-        st.session_state.code=""
-        st.session_state.temp_code=""
-        st.session_state.editor_key=str(uuid.uuid4())
-        st.session_state.packages_checked=False
-        st.session_state.latex_formula=""
-        st.session_state.audio_path=None
-        st.session_state.image_paths=[]
-        st.session_state.custom_library_result=""
-        st.session_state.python_script="""import matplotlib.pyplot as plt
-import numpy as np
-
-# Example: Create a simple plot
-x = np.linspace(0, 10, 100)
-y = np.sin(x)
-
-plt.figure(figsize=(10, 6))
-plt.plot(x, y, 'b-', label='sin(x)')
-plt.title('Sine Wave')
-plt.xlabel('x')
-plt.ylabel('sin(x)')
-plt.grid(True)
-plt.legend()
-"""
-        st.session_state.python_result=None
-        st.session_state.settings={"quality":"720p","format_type":"mp4","animation_speed":"Normal"}
-        st.session_state.password_entered=False
+        st.session_state.update({
+            'init':True, 'video_data':None, 'status':None, 'ai_models':None,
+            'generated_code':"", 'code':"", 'temp_code':"", 'editor_key':str(uuid.uuid4()),
+            'packages_checked':False, 'latex_formula':"", 'audio_path':None,
+            'image_paths':[], 'custom_library_result':"", 'python_script':"",
+            'python_result':None, 'active_tab':0,
+            'settings':{"quality":"720p","format_type":"mp4","animation_speed":"Normal"},
+            'password_entered':False, 'custom_model':"gpt-4o", 'first_load_complete':False,
+            'pending_tab_switch':None
+        })
     st.set_page_config(page_title="Manim Animation Studio", page_icon="🎬", layout="wide")
-    st.markdown("""
-    <style>
-    /* custom CSS */
-    </style>
-    """, unsafe_allow_html=True)
-    st.markdown("<h1 style='text-align:center;'>🎬 Manim Animation Studio</h1>", unsafe_allow_html=True)
     if not st.session_state.packages_checked:
         if ensure_packages():
             st.session_state.packages_checked=True
         else:
-            st.error("Failed to install packages"); st.stop()
-    if not ACE_EDITOR_AVAILABLE:
-        try:
-            from streamlit_ace import st_ace
-            ACE_EDITOR_AVAILABLE=True
-        except ImportError:
-            pass
-    tabs = st.tabs(["✨ Editor","🤖 AI Assistant","📚 LaTeX Formulas","🎨 Assets","🎞️ Timeline","🎓 Educational Export","🐍 Python Runner"])
-    # --- Editor Tab ---
+            st.error("Failed to install packages")
+            return
+
+    tab_names=["✨ Editor","🤖 AI Assistant","📚 LaTeX Formulas","🎨 Assets","🎞️ Timeline","🎓 Educational Export","🐍 Python Runner"]
+    tabs = st.tabs(tab_names)
+
+    # Editor Tab
     with tabs[0]:
-        col1,col2=st.columns([3,2])
+        col1,col2 = st.columns([3,2])
         with col1:
             st.markdown("### 📝 Animation Editor")
-            mode=st.radio("Input code:",["Type Code","Upload File"],key="editor_mode")
+            mode = st.radio("Code Input",["Type Code","Upload File"], key="editor_mode")
             if mode=="Upload File":
-                up=st.file_uploader("Upload .py",type=["py"],key="file_up")
+                up=st.file_uploader("Upload .py file", type=["py"])
                 if up:
-                    txt=up.getvalue().decode("utf-8")
-                    st.session_state.code=txt; st.session_state.temp_code=txt
+                    txt=up.getvalue().decode()
+                    if txt.strip():
+                        st.session_state.code=txt
+                        st.session_state.temp_code=txt
             if ACE_EDITOR_AVAILABLE:
-                code_in=st_ace(value=st.session_state.code,language="python",theme="monokai",min_lines=20,key=f"ace_{st.session_state.editor_key}")
+                st.session_state.temp_code = st_ace(value=st.session_state.code, language="python", theme="monokai", min_lines=20, key=f"ace_{st.session_state.editor_key}")
             else:
-                code_in=st.text_area("Code",value=st.session_state.code,height=400,key=f"ta_{st.session_state.editor_key}")
-            if code_in!=st.session_state.code:
-                st.session_state.code=code_in; st.session_state.temp_code=code_in
-            if st.button("🚀 Generate Animation",key="gen"):
-                if not st.session_state.code.strip():
+                st.session_state.temp_code = st.text_area("Code", st.session_state.code, height=400, key=f"ta_{st.session_state.editor_key}")
+            if st.session_state.temp_code!=st.session_state.code:
+                st.session_state.code=st.session_state.temp_code
+            if st.button("🚀 Generate Animation"):
+                if not st.session_state.code:
                     st.error("Enter code first")
                 else:
-                    sc=extract_scene_class_name(st.session_state.code)
-                    if sc=="MyScene" and "class MyScene" not in st.session_state.code:
-                        df="""\nclass MyScene(Scene):\n    def construct(self):\n        text=Text("Default Scene"); self.play(Write(text)); self.wait(2)\n"""
-                        st.session_state.code+=df; st.warning("No scene class; added default")
-                    with st.spinner("Rendering..."):
-                        d,s=generate_manim_video(st.session_state.code,st.session_state.settings["format_type"],st.session_state.settings["quality"],ANIMATION_SPEEDS[st.session_state.settings["animation_speed"]],st.session_state.audio_path)
-                        st.session_state.video_data=d; st.session_state.status=s
+                    vc,stt = generate_manim_video(
+                        st.session_state.code,
+                        st.session_state.settings["format_type"],
+                        st.session_state.settings["quality"],
+                        {"Slow":0.5,"Normal":1.0,"Fast":2.0,"Very Fast":3.0}[st.session_state.settings["animation_speed"]],
+                        st.session_state.audio_path
+                    )
+                    st.session_state.video_data=vc
+                    st.session_state.status=stt
         with col2:
-            st.markdown("### 🖥️ Preview & Output")
             if st.session_state.code:
-                st.markdown("<div style='border:1px solid #ccc;padding:10px;'>",unsafe_allow_html=True)
-                st.components.v1.html(generate_manim_preview(st.session_state.code),height=250)
-                st.markdown("</div>",unsafe_allow_html=True)
+                st.markdown("<div style='border:1px solid #ccc;padding:1rem;border-radius:8px;'>", unsafe_allow_html=True)
+                components.html(generate_manim_preview(st.session_state.code), height=250)
+                st.markdown("</div>", unsafe_allow_html=True)
             if st.session_state.video_data:
                 fmt=st.session_state.settings["format_type"]
                 if fmt=="png_sequence":
-                    st.download_button("⬇️ Download PNG Zip",st.session_state.video_data,file_name=f"frames_{int(time.time())}.zip")
+                    st.download_button("⬇️ Download PNG ZIP", data=st.session_state.video_data, file_name=f"manim_pngs_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip", mime="application/zip")
                 elif fmt=="svg":
-                    try: st.components.v1.html(st.session_state.video_data.decode('utf-8'),height=400)
-                    except: pass
-                    st.download_button("⬇️ Download SVG",st.session_state.video_data,file_name=f"anim.svg")
+                    try:
+                        svg=st.session_state.video_data.decode('utf-8')
+                        components.html(svg, height=400)
+                    except:
+                        st.error("Cannot display SVG")
+                    st.download_button("⬇️ Download SVG", data=st.session_state.video_data, file_name="animation.svg", mime="image/svg+xml")
                 else:
-                    st.video(st.session_state.video_data,format=fmt)
-                    st.download_button(f"⬇️ Download {fmt.upper()}",st.session_state.video_data,file_name=f"anim.{fmt}")
+                    st.video(st.session_state.video_data, format=fmt)
+                    st.download_button(f"⬇️ Download {fmt.upper()}", st.session_state.video_data, file_name=f"animation.{fmt}", mime=f"video/{fmt}" if fmt!="gif" else "image/gif")
             if st.session_state.status:
-                if "❌" in st.session_state.status: st.error(st.session_state.status)
-                else: st.success(st.session_state.status)
-    # --- AI Assistant Tab ---
+                if "Error" in st.session_state.status:
+                    st.error(st.session_state.status)
+                else:
+                    st.success(st.session_state.status)
+
+    # AI Assistant Tab
     with tabs[1]:
         st.markdown("### 🤖 AI Animation Assistant")
         if check_password():
-            if not st.session_state.ai_models:
-                st.session_state.ai_models=init_ai_models_direct()
-            # Debug & selection & generation (as in original)
-            with st.expander("🔧 Debug Connection"):
+            client_data = init_ai_models_direct()
+            if client_data:
                 if st.button("Test API Connection"):
                     with st.spinner("Testing..."):
-                        try:
-                            token=get_secret("github_token_api")
-                            if not token: st.error("Token missing"); st.stop()
-                            model=st.session_state.ai_models["model_name"]
-                            from openai import OpenAI
-                            client=OpenAI(base_url="https://models.github.ai/inference",api_key=token)
-                            params={"messages":[{"role":"system","content":"You are a helpful assistant."},{"role":"user","content":"Hi"}],"model":model}
-                            params[MODEL_CONFIGS[model]["param_name"]]=MODEL_CONFIGS[model][MODEL_CONFIGS[model]["param_name"]]
-                            resp=client.chat.completions.create(**params)
-                            if resp and resp.choices:
-                                st.success("✅ Connected")
-                            else: st.error("No response")
-                        except Exception as e:
-                            st.error(f"Error: {e}")
-            st.markdown("### 🤖 Model Selection")
-            cats={}
-            for m,cfg in MODEL_CONFIGS.items():
-                if m!="default":
-                    cats.setdefault(cfg["category"],[]).append(m)
-            cat_tabs=st.tabs(sorted(cats.keys()))
-            for i,cat in enumerate(sorted(cats.keys())):
-                with cat_tabs[i]:
-                    for m in sorted(cats[cat]):
-                        cfg=MODEL_CONFIGS[m]
-                        sel=(m==st.session_state.ai_models["model_name"])
-                        st.markdown(f"<div style='background:#f8f9fa;padding:10px;border-left:4px solid {'#0d6efd' if sel else '#4F46E5'};margin-bottom:8px;'>"
-                                    f"<h4>{m}</h4><p>Max Tokens: {cfg.get(cfg['param_name'],'?')}</p><p>API Ver: {cfg['api_version'] or 'default'}</p></div>",
-                                    unsafe_allow_html=True)
-                        if st.button("Select" if not sel else "Selected ✓",key=f"sel_{m}",disabled=sel):
-                            st.session_state.ai_models["model_name"]=m
-                            st.experimental_rerun()
-            if st.session_state.ai_models:
-                st.info(f"Using model: {st.session_state.ai_models['model_name']}")
-
-            if st.session_state.ai_models and "client" in st.session_state.ai_models:
-                st.markdown("#### Generate Animation from Description")
-                ideas=["...","3D sphere to torus","Pythagorean proof","Fourier transform","Neural network propagation","Integration area"]
-                sel=st.selectbox("Try idea",ideas)
-                prompt=sel if sel!="..." else ""
-                inp=st.text_area("Your prompt or code",value=prompt,height=150)
-                if st.button("Generate Animation Code"):
-                    if inp:
-                        with st.spinner("Generating..."):
-                            code=suggest_code_completion(inp,st.session_state.ai_models)
-                            if code:
+                        from azure.ai.inference.models import UserMessage
+                        api_params,_=prepare_api_params([UserMessage("Hello")], client_data["model_name"])
+                        resp=client_data["client"].complete(**api_params)
+                        if resp.choices:
+                            st.success("✅ Connection successful!")
+                            st.session_state.ai_models=client_data
+                        else:
+                            st.error("❌ No response")
+                if st.session_state.ai_models:
+                    st.info(f"Using model {st.session_state.ai_models['model_name']}")
+                    prompt = st.text_area("Describe animation or paste partial code", height=150)
+                    if st.button("Generate Animation Code"):
+                        if prompt.strip():
+                            from azure.ai.inference.models import UserMessage
+                            api_params,_=prepare_api_params([UserMessage(f"Write a complete Manim scene for:\n{prompt}")], st.session_state.ai_models["model_name"])
+                            resp=st.session_state.ai_models["client"].complete(**api_params)
+                            if resp.choices:
+                                code = resp.choices[0].message.content
+                                if "```python" in code:
+                                    code=code.split("```python")[1].split("```")[0]
                                 st.session_state.generated_code=code
-                            else: st.error("Failed")
-                    else: st.warning("Enter prompt")
-                if st.session_state.generated_code:
-                    st.code(st.session_state.generated_code,language="python")
-                    c1,c2=st.columns(2)
-                    if c1.button("Use This Code"):
-                        st.session_state.code=st.session_state.generated_code
-                        st.experimental_rerun()
-                    if c2.button("Render Preview"):
-                        vd,stt=generate_manim_video(st.session_state.generated_code,"mp4","480p",1.0)
-                        if vd: st.video(vd); st.download_button("Download Preview",vd,file_name="preview.mp4")
-                        else: st.error(f"Error: {stt}")
+                            else:
+                                st.error("No code generated")
+                        else:
+                            st.warning("Enter prompt first")
+                    if st.session_state.generated_code:
+                        st.code(st.session_state.generated_code, language="python")
+                        if st.button("Use This Code"):
+                            st.session_state.code=st.session_state.generated_code
+                            st.session_state.temp_code=st.session_state.generated_code
+                            st.session_state.pending_tab_switch=0
+                            st.rerun()
         else:
-            st.info("Enter password to access AI")
+            st.info("Enter password to access")
 
-    # --- LaTeX Formulas Tab ---
+    # LaTeX Formulas Tab
     with tabs[2]:
         st.markdown("### 📚 LaTeX Formula Builder")
-        c1,c2=st.columns([3,2])
-        with c1:
-            lt=st.text_area("LaTeX Formula",value=st.session_state.latex_formula,placeholder=r"e^{i\pi}+1=0",height=100)
-            st.session_state.latex_formula=lt
-            categories={
-                "Basic Math":[{"name":"Fraction","latex":r"\frac{a}{b}"},...],
-                # fill in as original categories...
-            }
-            tab_cats=st.tabs(list(categories.keys()))
-            for i,(cat,forms) in enumerate(categories.items()):
-                with tab_cats[i]:
-                    for f in forms:
-                        if st.button(f["name"],key=f"lt_{f['name']}"):
-                            st.session_state.latex_formula=f["latex"]; st.experimental_rerun()
-            if lt:
-                snippet=f"""
-formula=MathTex(r"{lt}")
+        col1,col2=st.columns([3,2])
+        with col1:
+            latex_input = st.text_area("LaTeX Formula", value=st.session_state.latex_formula, height=100, placeholder=r"e^{i\pi}+1=0")
+            st.session_state.latex_formula=latex_input
+            if latex_input:
+                manim_latex_code = f"""
+# LaTeX formula
+formula = MathTex(r"{latex_input}")
 self.play(Write(formula))
 self.wait(2)
 """
-                st.code(snippet,language="python")
+                st.code(manim_latex_code, language="python")
                 if st.button("Insert into Editor"):
-                    if "def construct" in st.session_state.code:
-                        lines=st.session_state.code.split("\n")
-                        idx=[i for i,l in enumerate(lines) if "def construct" in l][0]
-                        indent=re.match(r"(\s*)",lines[idx+1]).group(1) if idx+1<len(lines) else "        "
-                        insert="\n".join(indent+line for line in snippet.strip().split("\n"))
-                        lines.insert(idx+2,insert)
-                        st.session_state.code="\n".join(lines)
-                        st.experimental_rerun()
+                    if st.session_state.code:
+                        if "def construct(self):" in st.session_state.code:
+                            lines=st.session_state.code.split("\n")
+                            idx=-1
+                            for i,l in enumerate(lines):
+                                if "def construct(self):" in l:
+                                    idx=i; break
+                            if idx>=0:
+                                for j in range(idx+1,len(lines)):
+                                    if lines[j].strip() and not lines[j].strip().startswith("#"):
+                                        indent=re.match(r"(\s*)",lines[j]).group(1)
+                                        new_block="\n".join(indent+ln for ln in manim_latex_code.strip().split("\n"))
+                                        lines.insert(j,new_block)
+                                        break
+                            else:
+                                lines.append("        "+ "\n        ".join(manim_latex_code.strip().split("\n")))
+                            st.session_state.code="\n".join(lines)
+                            st.session_state.temp_code=st.session_state.code
+                            st.success("Inserted LaTeX into editor")
+                            st.session_state.pending_tab_switch=0
+                            st.rerun()
+                        else:
+                            st.warning("No construct() found")
                     else:
-                        base=f"""from manim import *\n\nclass LatexScene(Scene):\n    def construct(self):\n        {snippet.strip().replace('\n','\n        ')}\n"""
-                        st.session_state.code=base; st.experimental_rerun()
-        with c2:
-            st.components.v1.html(render_latex_preview(st.session_state.latex_formula),height=300)
+                        basic_scene = f"""from manim import *
+
+class LatexScene(Scene):
+    def construct(self):
+        # LaTeX formula
+        formula = MathTex(r"{latex_input}")
+        self.play(Write(formula))
+        self.wait(2)
+"""
+                        st.session_state.code=basic_scene
+                        st.session_state.temp_code=basic_scene
+                        st.success("Created new scene with LaTeX")
+                        st.session_state.pending_tab_switch=0
+                        st.rerun()
+        with col2:
+            components.html(render_latex_preview(latex_input), height=300)
 
-    # --- Assets Tab ---
+    # Assets Tab
     with tabs[3]:
         st.markdown("### 🎨 Asset Management")
-        a1,a2=st.columns(2)
-        with a1:
-            imgs=st.file_uploader("Upload Images",type=["png","jpg","jpeg","svg"],accept_multiple_files=True)
+        c1,c2 = st.columns(2)
+        with c1:
+            imgs=st.file_uploader("Upload Images", type=["png","jpg","jpeg","svg"], accept_multiple_files=True)
             if imgs:
-                d="manim_assets/images";os.makedirs(d,exist_ok=True)
+                img_dir=os.path.join(os.getcwd(),"manim_assets","images")
+                os.makedirs(img_dir, exist_ok=True)
                 for up in imgs:
                     ext=up.name.split(".")[-1]
-                    fn=f"img_{int(time.time())}_{uuid.uuid4().hex[:8]}.{ext}"
-                    p=os.path.join(d,fn)
-                    open(p,"wb").write(up.getvalue())
-                    st.session_state.image_paths.append({"name":up.name,"path":p})
-                st.success("Images uploaded")
+                    fname=f"img_{int(time.time())}_{uuid.uuid4().hex[:6]}.{ext}"
+                    path=os.path.join(img_dir,fname)
+                    with open(path,"wb") as f: f.write(up.getvalue())
+                    st.session_state.image_paths.append({"name":up.name,"path":path})
             if st.session_state.image_paths:
-                for ip in st.session_state.image_paths:
-                    st.image(Image.open(ip["path"]),caption=ip["name"],width=100)
-                    if st.button(f"Use {ip['name']}",key=f"use_img_{ip['name']}"):
-                        code=f"""
-image=ImageMobject(r"{ip['path']}")
+                for info in st.session_state.image_paths:
+                    img=Image.open(info["path"])
+                    st.image(img, caption=info["name"], width=100)
+                    if st.button(f"Use {info['name']}"):
+                        code_snippet=f"""
+# Image asset
+image = ImageMobject(r"{info['path']}")
+image.scale(2)
 self.play(FadeIn(image))
 self.wait(1)
 """
-                        st.session_state.code+=code; st.experimental_rerun()
-        with a2:
-            au=st.file_uploader("Upload Audio",type=["mp3","wav","ogg"])
-            if au:
-                d="manim_assets/audio";os.makedirs(d,exist_ok=True)
-                fn=f"audio_{int(time.time())}.{au.name.split('.')[-1]}"
-                p=os.path.join(d,fn)
-                open(p,"wb").write(au.getvalue())
-                st.session_state.audio_path=p
-                st.audio(au)
+                        st.session_state.code+=code_snippet
+                        st.session_state.temp_code=st.session_state.code
+                        st.success(f"Added {info['name']} to code")
+                        st.session_state.pending_tab_switch=0
+                        st.rerun()
+        with c2:
+            aud=st.file_uploader("Upload Audio", type=["mp3","wav","ogg"])
+            if aud:
+                adir=os.path.join(os.getcwd(),"manim_assets","audio")
+                os.makedirs(adir,exist_ok=True)
+                ext=aud.name.split(".")[-1]
+                aname=f"audio_{int(time.time())}.{ext}"
+                ap=os.path.join(adir,aname)
+                with open(ap,"wb") as f: f.write(aud.getvalue())
+                st.session_state.audio_path=ap
+                st.audio(aud)
                 st.success("Audio uploaded")
 
-    # --- Timeline Tab ---
+    # Timeline Tab
     with tabs[4]:
-        updated=create_timeline_editor(st.session_state.code)
-        if updated!=st.session_state.code:
-            st.session_state.code=updated; st.experimental_rerun()
+        st.markdown("### 🎞️ Timeline Editor")
+        st.info("Drag and adjust steps in code directly for now.")
 
-    # --- Educational Export Tab ---
+    # Educational Export Tab
     with tabs[5]:
         st.markdown("### 🎓 Educational Export")
         if not st.session_state.video_data:
-            st.warning("Generate animation first")
+            st.warning("Generate an animation first")
         else:
-            title=st.text_input("Animation Title","Manim Animation")
-            expl=st.text_area("Explanation",height=150)
-            fmt=st.selectbox("Format",["PowerPoint Presentation","Interactive HTML","Explanation Sequence PDF"])
+            title = st.text_input("Title", "Manim Animation")
+            expl = st.text_area("Explanation (use ## to separate steps)", height=150)
+            fmt = st.selectbox("Format", ["PowerPoint","HTML","PDF Sequence"])
             if st.button("Export"):
-                mp={"PowerPoint Presentation":"powerpoint","Interactive HTML":"html","Explanation Sequence PDF":"sequence"}
-                data,typ=export_to_educational_format(st.session_state.video_data,mp[fmt],title,expl,tempfile.mkdtemp())
-                if data:
-                    ext={"powerpoint":"pptx","html":"html","sequence":"pdf"}[typ]
-                    st.download_button("Download",data,file_name=f"{title.replace(' ','_')}.{ext}")
-                else: st.error("Export failed")
+                # Simplified, reuse generate_manim_video logic or placeholder
+                st.success(f"{fmt} export not yet implemented.")
 
-    # --- Python Runner Tab ---
+    # Python Runner Tab
     with tabs[6]:
         st.markdown("### 🐍 Python Script Runner")
-        examples={
-            "Basic Plot":st.session_state.python_script,
-            "Input Example":"""# input demo...""",
-            "DataFrame":"""import pandas as pd...""",
+        examples = {
+            "Select...":"",
+            "Sine Plot":"""import matplotlib.pyplot as plt
+import numpy as np
+x=np.linspace(0,10,100)
+y=np.sin(x)
+plt.plot(x,y)
+print("Done plotting")"""
         }
-        choice=st.selectbox("Examples",list(examples.keys()))
-        code=examples[choice] if choice in examples else st.session_state.python_script
+        sel=st.selectbox("Example", list(examples.keys()))
+        code = examples.get(sel, st.session_state.python_script)
         if ACE_EDITOR_AVAILABLE:
-            code_in=st_ace(value=code,language="python",theme="monokai",min_lines=15,key=f"pyace_{st.session_state.editor_key}")
+            code = st_ace(value=code, language="python", theme="monokai", min_lines=15, key="pyace")
         else:
-            code_in=st.text_area("Code",value=code,height=400,key=f"pyta_{st.session_state.editor_key}")
-        st.session_state.python_script=code_in
-        calls=detect_input_calls(code_in)
+            code = st.text_area("Code", code, height=300, key="pyta")
+        st.session_state.python_script=code
+        inputs = detect_input_calls(code)
         vals=[]
-        if calls:
-            st.markdown("Provide inputs:")
-            for i,c in enumerate(calls):
-                v=st.text_input(c["prompt"],key=f"inp_{i}")
-                vals.append(v)
-        timeout=st.slider("Timeout",5,300,30)
+        if inputs:
+            st.info(f"{len(inputs)} input() calls detected")
+            for i,c in enumerate(inputs):
+                vals.append(st.text_input(f"{c['prompt']} (line {c['line']})", key=f"inp{i}"))
+        timeout = st.slider("Timeout", 5,300,30)
         if st.button("▶️ Run"):
-            res=run_python_script(code_in,vals,timeout)
+            res=run_python_script(code, inputs=vals, timeout=timeout)
             st.session_state.python_result=res
         if st.session_state.python_result:
             display_python_script_results(st.session_state.python_result)
-            if st.session_state.python_result["plots"]:
-                st.markdown("Add plot to animation:")
-                for i,p in enumerate(st.session_state.python_result["plots"]):
-                    st.image(p); 
-                    if st.button(f"Use Plot {i+1}",key=f"use_plot_{i}"):
-                        path=tempfile.NamedTemporaryFile(delete=False,suffix=".png").name
-                        open(path,"wb").write(p)
-                        code=f"""
-plot_img=ImageMobject(r"{path}")
-self.play(FadeIn(plot_img))
-self.wait(1)
-"""
-                        st.session_state.code+=code; st.experimental_rerun()
+
+    # Handle tab switch after actions
+    if st.session_state.pending_tab_switch is not None:
+        st.session_state.active_tab = st.session_state.pending_tab_switch
+        st.session_state.pending_tab_switch=None
 
 if __name__ == "__main__":
     main()