import gradio as gr
import torch
import numpy as np
import tempfile
import os
from ase import Atoms
from ase.io import read, write
from ase.optimize import LBFGS
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
from ase.md.verlet import VelocityVerlet
from ase.io.trajectory import Trajectory
from ase.md import MDLogger
from ase import units
from orb_models.forcefield import pretrained
from orb_models.forcefield.calculator import ORBCalculator

# Install and import gradio_molecule3d
try:
    from gradio_molecule3d import Molecule3D
    HAS_3D = True
except ImportError:
    HAS_3D = False

# Default molecular representations for 3D visualization
DEFAULT_MOLECULAR_REPRESENTATIONS = [
    {
        "model": 0,
        "chain": "",
        "resname": "",
        "style": "sphere",
        "color": "Jmol",
        "scale": 0.3,
    },
    {
        "model": 0,
        "chain": "",
        "resname": "",
        "style": "stick",
        "color": "Jmol",
        "scale": 0.2,
    },
]

DEFAULT_MOLECULAR_SETTINGS = {
    "backgroundColor": "white",
    "orthographic": False,
    "disableFog": False,
}

# 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 run_md_simulation(input_file, md_steps, prerelax_steps, md_timestep, temperature_k, md_ensemble, charge, spin_multiplicity, explanation_buffer=""):
    """Run molecular dynamics simulation similar to FAIR Chem UMA"""
    try:
        calc = load_orbmol_model()
        if calc is None:
            return None, "❌ Error: Could not load OrbMol model", "", explanation_buffer
        
        # Parse input file content
        if isinstance(input_file, str):
            # Handle XYZ string input
            with tempfile.NamedTemporaryFile(mode='w', suffix='.xyz', delete=False) as f:
                f.write(input_file)
                xyz_file = f.name
            atoms = read(xyz_file)
            os.unlink(xyz_file)
        else:
            # Handle uploaded file
            atoms = read(input_file)
        
        # Set charge and spin
        atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
        atoms.calc = calc
        
        # Pre-relaxation
        if prerelax_steps > 0:
            opt = LBFGS(atoms)
            opt.run(fmax=0.05, steps=prerelax_steps)
        
        # Create trajectory file
        traj_path = tempfile.NamedTemporaryFile(suffix='.traj', delete=False).name
        
        # Initialize velocity distribution
        MaxwellBoltzmannDistribution(atoms, temperature_K=temperature_k)
        
        # Set up MD
        if md_ensemble == "NVE":
            dyn = VelocityVerlet(atoms, timestep=md_timestep * units.fs)
        else:  # NVT
            from ase.md.langevin import Langevin
            dyn = Langevin(atoms, md_timestep * units.fs, temperature_K=temperature_k, friction=0.001/units.fs)
        
        # Attach trajectory
        traj = Trajectory(traj_path, "w", atoms)
        dyn.attach(traj.write, interval=1)
        
        # Run simulation
        dyn.run(md_steps)
        traj.close()
        
        # Generate log
        log_content = f"""OrbMol Molecular Dynamics Simulation
=====================================
System: {len(atoms)} atoms
MD Steps: {md_steps}
Temperature: {temperature_k} K
Ensemble: {md_ensemble}
Timestep: {md_timestep} fs
Charge: {charge}
Spin Multiplicity: {spin_multiplicity}

Final Energy: {atoms.get_potential_energy():.6f} eV
Simulation completed successfully!
"""
        
        # Generate reproduction script
        script_content = f"""# OrbMol MD Simulation Script
import ase.io
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
from ase.md.verlet import VelocityVerlet
from ase.optimize import LBFGS
from ase.io.trajectory import Trajectory
from ase import units
from orb_models.forcefield import pretrained
from orb_models.forcefield.calculator import ORBCalculator

# Load structure
atoms = ase.io.read('input_structure.xyz')  # Your input file
atoms.info = {{"charge": {charge}, "spin": {spin_multiplicity}}}

# Set up OrbMol calculator
orbff = pretrained.orb_v3_conservative_inf_omat(device="cpu")
atoms.calc = ORBCalculator(orbff, device="cpu")

# Pre-relaxation
opt = LBFGS(atoms)
opt.run(fmax=0.05, steps={prerelax_steps})

# Initialize velocities
MaxwellBoltzmannDistribution(atoms, temperature_K={temperature_k})

# Set up MD
dyn = VelocityVerlet(atoms, timestep={md_timestep} * units.fs)
traj = Trajectory("md_output.traj", "w", atoms)
dyn.attach(traj.write, interval=1)

# Run simulation
dyn.run({md_steps})
"""
        
        explanation = explanation_buffer if explanation_buffer else f"Molecular dynamics simulation completed! This shows {len(atoms)} atoms moving over {md_steps} steps at {temperature_k} K. The atoms are vibrating due to thermal motion, and you can see the molecular structure evolving over time."
        
        return traj_path, log_content, script_content, explanation
        
    except Exception as e:
        return None, f"❌ Error during MD simulation: {str(e)}", "", "Error occurred"

def run_optimization(input_file, optimization_steps, fmax, charge, spin_multiplicity):
    """Run geometry optimization"""
    try:
        calc = load_orbmol_model()
        if calc is None:
            return None, "❌ Error: Could not load OrbMol model", ""
        
        # Parse input
        if isinstance(input_file, str):
            with tempfile.NamedTemporaryFile(mode='w', suffix='.xyz', delete=False) as f:
                f.write(input_file)
                xyz_file = f.name
            atoms = read(xyz_file)
            os.unlink(xyz_file)
        else:
            atoms = read(input_file)
        
        atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
        atoms.calc = calc
        
        # Create trajectory file
        traj_path = tempfile.NamedTemporaryFile(suffix='.traj', delete=False).name
        
        # Optimize
        opt = LBFGS(atoms, trajectory=traj_path)
        opt.run(fmax=fmax, steps=optimization_steps)
        
        # Generate log
        log_content = f"""OrbMol Geometry Optimization
===========================
System: {len(atoms)} atoms
Max Steps: {optimization_steps}
Force Convergence: {fmax} eV/Å
Charge: {charge}
Spin Multiplicity: {spin_multiplicity}

Final Energy: {atoms.get_potential_energy():.6f} eV
Max Force: {np.max(np.linalg.norm(atoms.get_forces(), axis=1)):.6f} eV/Å
Optimization completed!
"""
        
        script_content = f"""# OrbMol Geometry Optimization Script
import ase.io
from ase.optimize import LBFGS
from orb_models.forcefield import pretrained
from orb_models.forcefield.calculator import ORBCalculator

atoms = ase.io.read('input_structure.xyz')
atoms.info = {{"charge": {charge}, "spin": {spin_multiplicity}}}

orbff = pretrained.orb_v3_conservative_inf_omat(device="cpu")
atoms.calc = ORBCalculator(orbff, device="cpu")

opt = LBFGS(atoms, trajectory="optimization.traj")
opt.run(fmax={fmax}, steps={optimization_steps})
"""
        
        return traj_path, log_content, script_content
        
    except Exception as e:
        return None, f"❌ Error during optimization: {str(e)}", ""

def predict_single_point(xyz_content, charge=0, spin_multiplicity=1):
    """Single point energy and forces calculation"""
    try:
        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", ""
        
        with tempfile.NamedTemporaryFile(mode='w', suffix='.xyz', delete=False) as f:
            f.write(xyz_content)
            xyz_file = f.name
        
        atoms = read(xyz_file)
        atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
        atoms.calc = calc
        
        energy = atoms.get_potential_energy()
        forces = atoms.get_forces()
        
        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"
        
        max_force = np.max(np.linalg.norm(forces, axis=1))
        result += f"\n📊 **Max Force**: {max_force:.4f} eV/Å"
        
        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_md = [
    ["2\nHydrogen molecule\nH 0.0 0.0 0.0\nH 0.0 0.0 0.74", 100, 20, 1.0, 300.0, "NVE", 0, 1, "Simple H2 molecule dynamics showing thermal vibrations"],
    ["3\nWater molecule\nO 0.0 0.0 0.0\nH 0.7571 0.0 0.5864\nH -0.7571 0.0 0.5864", 200, 20, 1.0, 300.0, "NVE", 0, 1, "Water molecule thermal motion and vibrations"],
]

examples_sp = [
    ["""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],
    
    ["""5
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]
]

# Main Gradio interface
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")
                
                with gr.Tab("1. OrbMol Intro"):
                    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
                    
                    Try the examples below to see OrbMol in action!
                    """)
                    
                    # Quick examples for MD
                    gr.Examples(
                        examples=examples_md,
                        inputs=[
                            gr.Textbox(visible=False, value=""),  # Will be set by interface
                            gr.Slider(visible=False),
                            gr.Slider(visible=False),
                            gr.Slider(visible=False),
                            gr.Slider(visible=False),
                            gr.Radio(visible=False),
                            gr.Slider(visible=False),
                            gr.Slider(visible=False),
                            gr.Textbox(visible=False)
                        ],
                        outputs=[
                            gr.File(visible=False),
                            gr.Code(visible=False),
                            gr.Code(visible=False),
                            gr.Markdown(visible=False)
                        ],
                        fn=run_md_simulation,
                        cache_examples=True,
                        label="Try molecular dynamics examples!"
                    )
                
                with gr.Tab("2. Single Point Calculations"):
                    gr.Markdown("Calculate energy and forces for a single molecular geometry:")
                    
                    # Single point interface
                    with gr.Row():
                        with gr.Column():
                            xyz_input_sp = gr.Textbox(
                                label="XYZ Coordinates",
                                placeholder="""3
Water molecule
O 0.0 0.0 0.0
H 0.76 0.0 0.59
H -0.76 0.0 0.59""",
                                lines=8
                            )
                            
                            with gr.Row():
                                charge_sp = gr.Slider(value=0, label="Total Charge", minimum=-10, maximum=10, step=1)
                                spin_sp = gr.Slider(value=1, label="Spin Multiplicity", minimum=1, maximum=11, step=1)
                            
                            predict_btn_sp = gr.Button("Calculate Energy & Forces", variant="primary")
                        
                        with gr.Column():
                            results_sp = gr.Textbox(label="Results", lines=12, interactive=False)
                            status_sp = gr.Textbox(label="Status", max_lines=1, interactive=False)
                    
                    gr.Examples(
                        examples=examples_sp,
                        inputs=[xyz_input_sp, charge_sp, spin_sp],
                        label="Single point examples"
                    )
                
                with gr.Tab("3. Molecular Dynamics"):
                    gr.Markdown("Run molecular dynamics simulations with OrbMol:")
                    
                    xyz_input_md = gr.Textbox(
                        label="XYZ Coordinates",
                        placeholder="""3
Water molecule
O 0.0 0.0 0.0
H 0.76 0.0 0.59
H -0.76 0.0 0.59""",
                        lines=8
                    )
                    
                    with gr.Row():
                        md_steps = gr.Slider(minimum=10, maximum=500, value=100, label="MD Steps")
                        prerelax_steps = gr.Slider(minimum=0, maximum=100, value=20, label="Pre-Relaxation Steps")
                    
                    with gr.Row():
                        temperature_k = gr.Slider(minimum=0, maximum=1500, value=300, label="Temperature [K]")
                        md_timestep = gr.Slider(minimum=0.1, maximum=5.0, value=1.0, label="Timestep [fs]")
                    
                    md_ensemble = gr.Radio(choices=["NVE", "NVT"], value="NVE", label="Ensemble")
                    
                    with gr.Row():
                        charge_md = gr.Slider(value=0, label="Total Charge", minimum=-10, maximum=10, step=1)
                        spin_md = gr.Slider(value=1, label="Spin Multiplicity", minimum=1, maximum=11, step=1)
                    
                    md_button = gr.Button("Run MD Simulation", variant="primary")
                
                with gr.Tab("4. Geometry Optimization"):
                    gr.Markdown("Optimize molecular geometries with OrbMol:")
                    
                    xyz_input_opt = gr.Textbox(
                        label="XYZ Coordinates",
                        placeholder="""3
Water molecule
O 0.0 0.0 0.0
H 0.76 0.0 0.59
H -0.76 0.0 0.59""",
                        lines=8
                    )
                    
                    with gr.Row():
                        opt_steps = gr.Slider(minimum=1, maximum=500, value=300, label="Max Steps")
                        fmax = gr.Slider(minimum=0.001, maximum=0.5, value=0.05, label="Force Tolerance [eV/Å]")
                    
                    with gr.Row():
                        charge_opt = gr.Slider(value=0, label="Total Charge", minimum=-10, maximum=10, step=1)
                        spin_opt = gr.Slider(value=1, label="Spin Multiplicity", minimum=1, maximum=11, step=1)
                    
                    opt_button = gr.Button("Run Optimization", variant="primary")
        
        # Results panel
        with gr.Column(variant="panel", elem_id="results", min_width=500):
            gr.Markdown("## OrbMol Simulation Results")
            
            with gr.Tab("Visualization"):
                if HAS_3D:
                    output_structure = Molecule3D(
                        label="Simulation Visualization",
                        reps=DEFAULT_MOLECULAR_REPRESENTATIONS,
                        config=DEFAULT_MOLECULAR_SETTINGS,
                        height=500,
                        interactive=False,
                    )
                else:
                    gr.Markdown("3D visualization not available. Install gradio-molecule3d for 3D viewing.")
                
                output_traj = gr.File(label="Trajectory File", interactive=False)
                
                explanation = gr.Markdown("Run a simulation to see results here!")
            
            with gr.Tab("Log"):
                output_text = gr.Code(lines=20, max_lines=30, label="Simulation Log", interactive=False)
            
            with gr.Tab("Script"):
                reproduction_script = gr.Code(
                    interactive=False,
                    max_lines=30,
                    language="python",
                    label="Reproduction Script",
                )
    
    # Sidebar
    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
            * Built on Orb-v3 architecture, trained on OMol25 dataset (100M+ DFT calculations)
            * Supports charge and spin multiplicity for accurate molecular modeling
            * Optimized for biomolecules, metal complexes, and electrolytes
            
            [Read more about OrbMol](https://orbitalmaterials.com/posts/orbmol-extending-orb-to-molecular-systems)
            """)
        
        with gr.Accordion("Model Disclaimers", open=False):
            gr.Markdown("""
            * OrbMol has limitations and may not work perfectly for all systems
            * Always validate results for your specific use case
            * Consider the limitations of the training data and methodology
            """)
        
        with gr.Accordion("Open source packages", open=False):
            gr.Markdown("""
            * Model: [orbital-materials/orb-models](https://github.com/orbital-materials/orb-models)
            * Uses ASE, Gradio, and other open source packages
            * Licensed under Apache 2.0
            """)
    
    # Connect buttons to functions
    predict_btn_sp.click(
        predict_single_point,
        inputs=[xyz_input_sp, charge_sp, spin_sp],
        outputs=[results_sp, status_sp]
    )
    
    md_button.click(
        run_md_simulation,
        inputs=[xyz_input_md, md_steps, prerelax_steps, md_timestep, temperature_k, md_ensemble, charge_md, spin_md],
        outputs=[output_traj, output_text, reproduction_script, explanation]
    )
    
    opt_button.click(
        run_optimization,
        inputs=[xyz_input_opt, opt_steps, fmax, charge_opt, spin_opt],
        outputs=[output_traj, output_text, reproduction_script]
    )
    
    # Update 3D visualization when trajectory changes
    if HAS_3D:
        output_traj.change(
            lambda x: x,
            inputs=[output_traj],
            outputs=[output_structure]
        )

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

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