import gradio as gr
import torch
import numpy as np
from ase import Atoms
from ase.io import read, write
import tempfile
import os
from orb_models.forcefield import pretrained
from orb_models.forcefield.calculator import ORBCalculator

# Global variable for the model
model_calc = None

def load_orbmol_model():
    """Load OrbMol model once"""
    global model_calc
    if model_calc is None:
        try:
            print("Loading OrbMol model...")
            orbff = pretrained.orb_v3_conservative_inf_omat(
                device="cpu",  
                precision="float32-high"
            )
            model_calc = ORBCalculator(orbff, device="cpu")
            print("✅ OrbMol model loaded successfully")
        except Exception as e:
            print(f"❌ Error loading model: {e}")
            model_calc = None
    return model_calc

def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
    """
    Main function: XYZ → OrbMol → Results
    """
    try:
        # Load model
        calc = load_orbmol_model()
        if calc is None:
            return "❌ Error: Could not load OrbMol model", ""
        
        if not xyz_content.strip():
            return "❌ Error: Please enter XYZ coordinates", ""
        
        # Create temporary file with XYZ
        with tempfile.NamedTemporaryFile(mode='w', suffix='.xyz', delete=False) as f:
            f.write(xyz_content)
            xyz_file = f.name
        
        # Read molecular structure
        atoms = read(xyz_file)
        
        # Configure charge and spin (IMPORTANT for OrbMol!)
        atoms.info = {
            "charge": int(charge), 
            "spin": int(spin_multiplicity)
        }
        
        # Assign OrbMol calculator
        atoms.calc = calc
        
        # Make the prediction!
        energy = atoms.get_potential_energy()  # In eV
        forces = atoms.get_forces()  # In eV/Å
        
        # Format results nicely
        result = f"""
🔋 **Total Energy**: {energy:.6f} eV

⚡ **Atomic Forces**:
"""
        
        for i, force in enumerate(forces):
            result += f"Atom {i+1}: [{force[0]:.4f}, {force[1]:.4f}, {force[2]:.4f}] eV/Å\n"
        
        # Additional statistics
        max_force = np.max(np.linalg.norm(forces, axis=1))
        result += f"\n📊 **Max Force**: {max_force:.4f} eV/Å"
        
        # Clean up temporary file
        os.unlink(xyz_file)
        
        return result, "✅ Calculation completed with OrbMol"
        
    except Exception as e:
        return f"❌ Error during calculation: {str(e)}", "Error"

# Predefined examples
examples = [
    ["""2
Hydrogen molecule
H 0.0 0.0 0.0
H 0.0 0.0 0.74""", 0, 1],
    
    ["""3
Water molecule  
O 0.0000 0.0000 0.0000
H 0.7571 0.0000 0.5864
H -0.7571 0.0000 0.5864""", 0, 1],
    
    ["""4
Methane
C 0.0000 0.0000 0.0000
H 1.0890 0.0000 0.0000
H -0.3630 1.0267 0.0000
H -0.3630 -0.5133 0.8887
H -0.3630 -0.5133 -0.8887""", 0, 1]
]

# Gradio interface - using FAIR Chem UMA style
with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
    
    with gr.Row():
        with gr.Column(scale=2):
            with gr.Column(variant="panel"):
                gr.Markdown("# OrbMol Demo - Quantum-Accurate Molecular Predictions")
                
                gr.Markdown("""
                **OrbMol** is a neural network potential trained on the **OMol25** dataset (100M+ high-accuracy DFT calculations).
                
                Predicts **energies** and **forces** with quantum accuracy, optimized for:
                * 🧬 Biomolecules  
                * ⚗️ Metal complexes
                * 🔋 Electrolytes
                """)
                
                gr.Markdown("## Simulation inputs")
                
                with gr.Column(variant="panel"):
                    gr.Markdown("### Input molecular structure")
                    
                    xyz_input = gr.Textbox(
                        label="XYZ Coordinates",
                        placeholder="""3
Water molecule
O 0.0000 0.0000 0.0000  
H 0.7571 0.0000 0.5864
H -0.7571 0.0000 0.5864""",
                        lines=12,
                        info="Paste XYZ coordinates of your molecule here"
                    )
                    
                    gr.Markdown("OMol-specific settings for total charge and spin multiplicity")
                    with gr.Row():
                        charge_input = gr.Slider(
                            value=0, label="Total Charge", minimum=-10, maximum=10, step=1
                        )
                        spin_input = gr.Slider(
                            value=1, maximum=11, minimum=1, step=1, label="Spin Multiplicity"
                        )
                    
                    predict_btn = gr.Button("Run OrbMol Prediction", variant="primary", size="lg")
                    
        with gr.Column(variant="panel", elem_id="results", min_width=500):
            gr.Markdown("## OrbMol Prediction Results")
            
            results_output = gr.Textbox(
                label="Energy & Forces",
                lines=15,
                interactive=False,
                info="OrbMol energy and force predictions"
            )
            
            status_output = gr.Textbox(
                label="Status",
                interactive=False,
                max_lines=1
            )
    
    # Examples section
    gr.Markdown("### 🧪 Try These Examples")
    gr.Examples(
        examples=examples,
        inputs=[xyz_input, charge_input, spin_input],
        label="Click any example to load it"
    )
    
    # Connect button to function
    predict_btn.click(
        predict_molecule,
        inputs=[xyz_input, charge_input, spin_input],
        outputs=[results_output, status_output]
    )
    
    # Footer info - matching FAIR Chem UMA style
    with gr.Sidebar(open=True):
        gr.Markdown("## Learn more about OrbMol")
        with gr.Accordion("What is OrbMol?", open=False):
            gr.Markdown("""
            * OrbMol is a neural network potential for molecular property prediction with quantum-level accuracy
            * Built on the Orb-v3 architecture and trained on OMol25 dataset (100M+ DFT calculations)  
            * Optimized for biomolecules, metal complexes, and electrolytes
            * Supports configurable charge and spin multiplicity
            
            [Read more about OrbMol](https://orbitalmaterials.com/posts/orbmol-extending-orb-to-molecular-systems)
            """)
            
        with gr.Accordion("Model Disclaimers", open=False):
            gr.Markdown("""
            * While OrbMol represents significant progress in molecular ML potentials, the model has limitations
            * Always validate results for your specific use case
            * Consider the limitations of the ωB97M-V/def2-TZVPD level of theory used in training
            """)
            
        with gr.Accordion("Open source packages", open=False):
            gr.Markdown("""
            * Model code available at [orbital-materials/orb-models](https://github.com/orbital-materials/orb-models)
            * This demo uses ASE, Gradio, and other open source packages
            """)

# Load model on startup
print("🚀 Starting OrbMol model loading...")
load_orbmol_model()

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )