app.py

import gradio as gr
import os
from dotenv import load_dotenv
import google.generativeai as genai
import ase
from ase.build import molecule # Often used, good to have for exec scope if needed
from ase.calculators.emt import EMT # Common default
from ase.optimize import BFGS # Common optimizer
from ase.visualize.plot import plot_atoms
import matplotlib
matplotlib.use('Agg') # Use Agg backend for non-interactive plotting
import matplotlib.pyplot as plt
from PIL import Image
import io
import base64
import traceback
import contextlib

# Load environment variables from .env file (for GEMINI_API_KEY)
load_dotenv()

# --- Gemini API Configuration ---
API_KEY = os.getenv("GEMINI_API_KEY")
# MODEL = 'gemini-2.0-flash'
MODEL = 'gemini-2.5-flash-preview-04-17'
# MODEL = 'gemini-2.5-pro-exp-03-25'

if not API_KEY:
    print("WARNING: GEMINI_API_KEY environment variable not set. App will not function.")
else:
    try:
        genai.configure(api_key=API_KEY)
        MODEL = genai.GenerativeModel(MODEL)
        print("Gemini model initialized successfully.")
    except Exception as e:
        print(f"ERROR: Could not initialize Gemini model: {e}")
        MODEL = None

# --- Prompts for Gemini ---
PROMPT_FOR_ASE_CODE_GENERATION = """
You are an expert in molecular simulations using the Atomic Simulation Environment (ASE) and Python.
A user wants to perform the following task: "{user_query}"

Generate a complete and runnable Python script using ASE to perform this task.
The script MUST:
1.  Include all necessary import statements (e.g., `import ase`, `from ase.build import molecule`, `from ase.calculators.emt import EMT`, `from ase.optimize import BFGS`, etc.). Do not assume any modules are pre-imported.
2.  Define the atomic structure.
3.  Set up an appropriate ASE calculator. 
4.  Perform the simulation as requested (e.g., geometry optimization, molecular dynamics, single point energy calculation). If non is specified perform an appropriate simulation. 
5.  Ensure the final ASE `Atoms` object, representing the result of the simulation (e.g., optimized structure), is assigned to a variable named `atoms_object`. This is crucial for visualization.
6.  Optionally, print any relevant information or results from the simulation to standard output (e.g., final energy, bond lengths, status messages).


Provide ONLY the Python code block. Do not include any explanations, comments, or markdown formatting before or after the code block itself.
The script will be executed in an environment where ASE and its dependencies are installed."""

PROMPT_FOR_ANALYSIS = """
You are an AI assistant specialized in interpreting results from molecular simulations performed with the Atomic Simulation Environment (ASE).

The user's original simulation request was:
"{user_query}"

The following Python ASE code was generated and executed to address this request:
```python
{generated_code}
```

The execution of this code produced the following standard output (stdout/stderr):
```
{simulation_output}
```

{structure_info}

Based on all the above information, please provide a concise analysis:
1.  Briefly explain what the simulation aimed to do and what steps were performed by the code.
2.  Summarize the key results obtained from the simulation output (e.g., final energy, convergence status, important structural parameters if mentioned in output).
3.  If a molecular structure was visualized, comment on it in relation to the simulation (e.g., "The visualized structure shows the molecule after optimization").
4.  Conclude whether the simulation likely achieved the user's original request.
5.  If there were any errors or issues, please mention them.

Present the analysis in a clear, easy-to-understand manner. Avoid overly technical jargon where possible.
Provide ONLY the Markdown code block. Do not include any explanations, comments before or after the code block itself.
"""

# --- Helper Functions ---
def plot_atoms_to_pil_image(atoms_obj):
    """Converts an ASE Atoms object to a PIL Image for display."""
    if not isinstance(atoms_obj, ase.Atoms):
        return None
    try:
        fig, ax = plt.subplots()
        # Use ase.visualize.plot_atoms for plotting
        plot_atoms(atoms_obj, ax, radii=0.3, rotation=('0x,0y,0z'))
        ax.set_xlabel("X (Å)")
        ax.set_ylabel("Y (Å)")
        ax.axis('equal')
        
        buf = io.BytesIO()
        fig.savefig(buf, format='png', bbox_inches='tight')
        plt.close(fig)  # Close the figure to free up memory
        buf.seek(0)
        return Image.open(buf)
    except Exception as e:
        print(f"Error plotting atoms: {traceback.format_exc()}")
        return None

# --- Core Simulation Logic ---
def simulate_and_analyze(user_query_str):
    if not MODEL:
        error_msg = "Gemini model not initialized. Please check API key and configuration."
        return "# Error: Model not initialized.", error_msg, None, error_msg

    generated_code_str = ""
    simulation_stdout_str = ""
    pil_image = None
    analysis_text_str = ""

    try:
        # 1. Generate ASE code from user query
        prompt_for_code = PROMPT_FOR_ASE_CODE_GENERATION.format(user_query=user_query_str)
        response_code = MODEL.generate_content(prompt_for_code)
        generated_code_str = response_code.text.strip()
        # Remove potential markdown backticks if Gemini adds them
        if generated_code_str.startswith("```python"):
            generated_code_str = generated_code_str[len("```python"):].strip()
        if generated_code_str.endswith("```"):
            generated_code_str = generated_code_str[:-len("```")].strip()

    except Exception as e:
        error_msg = f"Error generating ASE code with Gemini: {traceback.format_exc()}"
        return f"# Error: Could not generate ASE code.\n{error_msg}", error_msg, None, error_msg

    # 2. Execute generated ASE code
    atoms_object_from_exec = None
    execution_namespace = {}
    captured_output = io.StringIO()
    exec_error_msg = None

    try:
        with contextlib.redirect_stdout(captured_output):
            # Provide a minimal, controlled global scope for exec
            # ASE and its modules should be imported by the generated code itself
            exec(generated_code_str, {'__builtins__': __builtins__, 'ase': ase}, execution_namespace)
        
        simulation_stdout_str = captured_output.getvalue()
        atoms_object_from_exec = execution_namespace.get('atoms_object')

        if atoms_object_from_exec is not None and not isinstance(atoms_object_from_exec, ase.Atoms):
            warning_msg = f"Warning: 'atoms_object' was found but is not a valid ASE Atoms object (type: {type(atoms_object_from_exec)}). Visualization might fail."
            simulation_stdout_str += "\n" + warning_msg
            atoms_object_from_exec = None # Invalidate if not correct type
        elif 'atoms_object' not in execution_namespace:
            simulation_stdout_str += "\nWarning: 'atoms_object' variable not found after code execution. Visualization will not be available."

    except Exception as e:
        exec_error_msg = f"Error during ASE code execution:\n{traceback.format_exc()}"
        simulation_stdout_str = captured_output.getvalue() # Get any output before error
        if simulation_stdout_str:
            simulation_stdout_str += f"\n\n{exec_error_msg}"
        else:
            simulation_stdout_str = exec_error_msg
        atoms_object_from_exec = None # Ensure it's None on error

    # 3. Plot atoms object if available
    if atoms_object_from_exec and isinstance(atoms_object_from_exec, ase.Atoms):
        pil_image = plot_atoms_to_pil_image(atoms_object_from_exec)
        structure_info_for_analysis = "A molecular structure was generated and visualized."
    else:
        structure_info_for_analysis = "No molecular structure was available or successfully generated for visualization."
        if exec_error_msg and not atoms_object_from_exec : # If exec failed before atoms_object could be assigned
             structure_info_for_analysis += f" This might be due to the execution error: {exec_error_msg.splitlines()[0]}"

    # 4. Generate analysis using Gemini
    try:
        prompt_for_analysis = PROMPT_FOR_ANALYSIS.format(
            user_query=user_query_str,
            generated_code=generated_code_str,
            simulation_output=simulation_stdout_str,
            structure_info=structure_info_for_analysis
        )
        response_analysis = MODEL.generate_content(prompt_for_analysis)
        analysis_text_str = response_analysis.text

        analysis_text_str = analysis_text_str.strip()

        if analysis_text_str.startswith("```markdown"):
            analysis_text_str = analysis_text_str[len("```markdown"):].strip()
        if analysis_text_str.endswith("```"):
            analysis_text_str = analysis_text_str[:-len("```")].strip()

    except Exception as e:
        analysis_text_str = f"Error generating analysis with Gemini: {traceback.format_exc()}"
        if exec_error_msg: # If code exec failed, make analysis reflect that primarily
            analysis_text_str = f"Could not perform analysis as the ASE code execution failed.\nDetails: {exec_error_msg}\n\nGemini analysis generation also failed: {e}"
        elif not generated_code_str:
            analysis_text_str = f"Could not perform analysis as ASE code generation failed.\nGemini analysis generation also failed: {e}"

    return generated_code_str, simulation_stdout_str, pil_image, analysis_text_str

# --- Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal", secondary_hue="orange")) as demo:
    gr.Markdown("# ⚛️ GeminiSim: ASE Simulations via Natural Language 💬")
    gr.Markdown(
        "Describe your atomic/molecular simulation task in plain English. "
        "Gemini AI will generate ASE Python code, run the simulation, and provide an analysis of the results.\n\n"
        "**⚠️ SECURITY WARNING:** This tool executes AI-generated Python code using `exec()`. "
        "This is inherently unsafe if the AI generates malicious code. Use with caution and only with trusted inputs, "
        "especially in non-sandboxed environments."
    )

    with gr.Row():
        with gr.Column(scale=2):
            user_query_input = gr.Textbox(
                label="Your Simulation Request:",
                placeholder="e.g., 'Optimize a water molecule (H2O) using EMT and show the final geometry and energy.'",
                lines=10,
                elem_id="user_query_textbox"
            )
        with gr.Column(scale=1):
            gr.Markdown("### ✨ Example Prompts:")
            with gr.Tabs():
                with gr.TabItem("Simple Molecules"):
                    gr.Examples(
                        examples=[
                            "Optimize a water molecule (H2O) using EMT and show the final geometry and energy.",
                            "Optimize a CO molecule using BFGS and show its bond length.",
                            "Build a methane molecule (CH4) and calculate its energy with EMT.",
                            "Simulate a single gold (Au) atom.",
                            "Calculate the binding energy of an O2 molecule."
                        ],
                        inputs=[user_query_input],
                    )
                with gr.TabItem("Complex Simulations"):
                    gr.Examples(
                        examples=[
                            "Simulate a basic alcohol molecule (e.g., methanol) and report key bond lengths.",
                            "Relax an ammonia molecule (NH3) and display its bond angles.",
                            "Optimize a CO2 molecule and display its linearity.",
                            "Simulate a water dimer and analyze the hydrogen bond length.",
                            "Create a small gold cluster (Au3) and calculate its binding energy."
                        ],
                        inputs=[user_query_input],
                    )

    submit_button = gr.Button("🚀 Simulate & Analyze", variant="primary")
    
    with gr.Accordion("🔬 Results", open=True):
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 🤖 Generated ASE Code")
                generated_code_display = gr.Code(language="python", label="Generated ASE Python Code", lines=15, interactive=False)
                gr.Markdown("### 🖥️ Simulation Standard Output")
                simulation_stdout_display = gr.Textbox(label="Simulation stdout/stderr", lines=8, interactive=False, show_copy_button=True)
            
            with gr.Column(scale=1):
                gr.Markdown("### 🖼️ Molecular Visualization")
                molecule_plot_display = gr.Image(label="Molecular Structure", type="pil", interactive=False, height=350, show_download_button=True)
                gr.Markdown("### 🧠 AI Analysis")
                ai_analysis_display = gr.Markdown(label="Gemini's Analysis ", show_copy_button=True)

    gr.Markdown("---")
    gr.Markdown("Built with [Gradio](https://gradio.app/), [ASE](https://wiki.fysik.dtu.dk/ase/), and [Google Gemini](https://ai.google.dev/).")

    submit_button.click(
        fn=simulate_and_analyze,
        inputs=[user_query_input],
        outputs=[generated_code_display, simulation_stdout_display, molecule_plot_display, ai_analysis_display]
    )

if __name__ == "__main__":
    if not API_KEY or not MODEL:
        print("ERROR: Gemini API key not set or model failed to initialize. Gradio server will not start.")
        # Optionally, could display an error in Gradio itself if it were to launch.
    else:
        print("Launching Gradio server...")
        demo.launch() # Add share=True if you want a public link (be mindful of security warning!)