app.py

import os
os.environ["TORCH_DYNAMO_DISABLE"] = "1"
import tempfile
import numpy as np
import gradio as gr
from ase.io import read, write
from ase.io.trajectory import Trajectory
from gradio_molecule3d import Molecule3D
from simulation_scripts_orbmol import load_orbmol_model
import hashlib
import shutil

# ==== Configuración UNIVERSAL para Molecule3D ====
UNIVERSAL_MOLECULE3D_REPS = [
    {
        "model": 0,
        "chain": "",
        "resname": "",
        "style": "sphere",
        "color": "element",  # Colores automáticos por elemento químico
        "residue_range": "",
        "around": 0,
        "byres": False,
        "visible": True,
        "opacity": 0.8
    },
    {
        "model": 0,
        "chain": "",
        "resname": "",
        "style": "stick",
        "color": "element",
        "residue_range": "",
        "around": 0,
        "byres": False,
        "visible": True,
        "opacity": 1.0
    }
]

# ==== Función mejorada para crear PDB universal ====
def create_universal_pdb(atoms, cache_key):
    """
    Crea un PDB optimizado que funciona bien con gradio_molecule3d para cualquier molécula
    """
    cache_dir = os.path.join(tempfile.gettempdir(), "gradio")
    os.makedirs(cache_dir, exist_ok=True)
    pdb_path = os.path.join(cache_dir, f"mol_{cache_key}.pdb")
    
    if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
        return pdb_path
    
    try:
        positions = atoms.get_positions()
        symbols = atoms.get_chemical_symbols()
        
        with open(pdb_path, "w") as f:
            # Header estándar
            f.write("HEADER    MOLECULAR STRUCTURE                     01-JAN-25   MOL \n")
            f.write("TITLE     CALCULATION RESULT\n")
            f.write("MODEL        1\n")
            
            # Escribir átomos con formato PDB correcto
            for i, (symbol, pos) in enumerate(zip(symbols, positions)):
                # Asegurar formato correcto del elemento
                element = symbol.strip().upper()
                
                # Nombre del átomo (4 caracteres, alineado a la izquierda)
                atom_name = f"{element:<4s}"
                
                # Línea ATOM con formato PDB estricto
                f.write(
                    f"ATOM  {i+1:5d} {atom_name} UNL A   1    "
                    f"{pos[0]:8.3f}{pos[1]:8.3f}{pos[2]:8.3f}"
                    f"  1.00 30.00           {element:>2s}  \n"
                )
            
            f.write("ENDMDL\n")
            f.write("END\n")
        
        # Verificar creación
        if os.path.exists(pdb_path) and os.path.getsize(pdb_path) > 0:
            print(f"✅ Universal PDB created: {pdb_path} ({os.path.getsize(pdb_path)} bytes)")
            
            # Debug: mostrar contenido
            with open(pdb_path, 'r') as f:
                content = f.read()
                print(f"PDB content preview:\n{content[:300]}...")
            
            return pdb_path
        else:
            print(f"❌ Universal PDB creation failed")
            return None
            
    except Exception as e:
        print(f"❌ Error in create_universal_pdb: {e}")
        import traceback
        traceback.print_exc()
        return None

# ==== Función actualizada para preparar cualquier molécula ====
def prepare_universal_molecule_viewer(traj_path_or_atoms):
    """
    Función universal que funciona con cualquier molécula
    """
    try:
        # Determinar si es trayectoria o átomos
        if isinstance(traj_path_or_atoms, str):
            if not os.path.exists(traj_path_or_atoms):
                print("Trajectory file doesn't exist")
                return None
            
            traj = Trajectory(traj_path_or_atoms)
            if len(traj) == 0:
                print("Empty trajectory")
                return None
            
            atoms = traj[-1]
            cache_key = hashlib.md5(str(traj_path_or_atoms).encode()).hexdigest()[:12]
        else:
            atoms = traj_path_or_atoms
            cache_key = hashlib.md5(atoms.get_positions().tobytes()).hexdigest()[:12]
        
        print(f"Preparing universal viewer: {len(atoms)} atoms")
        print(f"Chemical symbols: {atoms.get_chemical_symbols()}")
        
        # Crear PDB universal
        pdb_path = create_universal_pdb(atoms, cache_key)
        
        return pdb_path
        
    except Exception as e:
        print(f"Error in prepare_universal_molecule_viewer: {e}")
        import traceback
        traceback.print_exc()
        return None

# ==== Funciones de preparación universales ====
def prepare_input_for_viewer(structure_file):
    """Universal input preparation"""
    if not structure_file or not os.path.exists(structure_file):
        return None
    try:
        atoms = read(structure_file)
        return prepare_universal_molecule_viewer(atoms)
    except Exception as e:
        print(f"Error preparing input for viewer: {e}")
        return None

def prepare_molecule_for_viewer(traj_path):
    """Universal trajectory preparation"""
    return prepare_universal_molecule_viewer(traj_path)

# ==== OrbMol SPE ====
from orb_models.forcefield import pretrained
from orb_models.forcefield.calculator import ORBCalculator

def predict_molecule(structure_file, task_name, charge=0, spin_multiplicity=1):
    """
    Single Point Energy + fuerzas (OrbMol). Acepta archivos subidos.
    """
    try:
        calc = load_orbmol_model(task_name)
        if not structure_file:
            return "Error: Please upload a structure file", "Error", None

        file_path = structure_file
        if not os.path.exists(file_path):
            return f"Error: File not found: {file_path}", "Error", None
        if os.path.getsize(file_path) == 0:
            return f"Error: Empty file: {file_path}", "Error", None

        atoms = read(file_path)
        atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
        atoms.calc = calc

        energy = atoms.get_potential_energy()
        forces = atoms.get_forces()

        lines = [f"Total Energy: {energy:.6f} eV", "", "Atomic Forces:"]
        for i, fc in enumerate(forces):
            lines.append(f"Atom {i+1}: [{fc[0]:.4f}, {fc[1]:.4f}, {fc[2]:.4f}] eV/Å")
        max_force = float(np.max(np.linalg.norm(forces, axis=1)))
        lines += ["", f"Max Force: {max_force:.4f} eV/Å"]

        # Preparar PDB para visualización
        pdb_file = prepare_input_for_viewer(file_path)

        return "\n".join(lines), f"Calculation completed with {task_name}", pdb_file
    except Exception as e:
        import traceback
        traceback.print_exc()
        return f"Error during calculation: {e}", "Error", None

# ==== Simulaciones (helpers) ====
from simulation_scripts_orbmol import (
    run_md_simulation,
    run_relaxation_simulation,
)

# ==== Wrappers con debug y Molecule3D ====
def md_wrapper(structure_file, task_name, charge, spin, steps, tempK, timestep_fs, ensemble):
    try:
        if not structure_file:
            return ("Error: Please upload a structure file", None, "", "", "", None)

        file_path = structure_file
        print(f"MD Wrapper: Processing {file_path}")

        traj_path, log_text, script_text, explanation = run_md_simulation(
            file_path,
            int(steps),
            20,  # pre-relax steps
            float(timestep_fs),
            float(tempK),
            "NVT" if ensemble == "NVT" else "NVE",
            str(task_name),
            int(charge),
            int(spin),
        )
        status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
        print(f"MD completed, trajectory: {traj_path}")

        pdb_file = prepare_molecule_for_viewer(traj_path)
        print(f"PDB file for Molecule3D: {pdb_file}")

        return (status, traj_path, log_text, script_text, explanation, pdb_file)

    except Exception as e:
        print(f"MD Wrapper Error: {e}")
        import traceback
        traceback.print_exc()
        return (f"Error: {e}", None, "", "", "", None)

def relax_wrapper(structure_file, task_name, steps, fmax, charge, spin, relax_cell):
    try:
        if not structure_file:
            return ("Error: Please upload a structure file", None, "", "", "", None)

        file_path = structure_file
        print(f"Relax Wrapper: Processing {file_path}")

        traj_path, log_text, script_text, explanation = run_relaxation_simulation(
            file_path,
            int(steps),
            float(fmax),
            str(task_name),
            int(charge),
            int(spin),
            bool(relax_cell),
        )
        status = f"Relaxation finished (≤ {int(steps)} steps, fmax={float(fmax)} eV/Å)"
        print(f"Relaxation completed, trajectory: {traj_path}")

        pdb_file = prepare_molecule_for_viewer(traj_path)
        print(f"PDB file for Molecule3D: {pdb_file}")

        return (status, traj_path, log_text, script_text, explanation, pdb_file)

    except Exception as e:
        print(f"Relax Wrapper Error: {e}")
        import traceback
        traceback.print_exc()
        return (f"Error: {e}", None, "", "", "", None)

# ==== UI ====
with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
    with gr.Tabs():
        # -------- SPE --------
        with gr.Tab("Single Point Energy"):
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("# OrbMol — Quantum-Accurate Molecular Predictions")
                    gr.Markdown("Upload molecular structure files (.xyz, .pdb, .cif, .traj) for energy and force calculations.")
                    
                    xyz_input = gr.File(
                        label="Upload Structure File (.xyz/.pdb/.cif/.traj)",
                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                        file_count="single"
                    )
                    with gr.Row():
                        task_name_spe = gr.Radio(
                            ["OMol", "OMat", "OMol-Direct"],
                            value="OMol",
                            label="Model Type",
                            info="Choose the OrbMol model variant for the calculation."
                        )
                    with gr.Row():
                        charge_input = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
                        spin_input = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin Multiplicity")
                    
                    run_spe = gr.Button("Run OrbMol Prediction", variant="primary")
                    
                with gr.Column(variant="panel", min_width=500):
                    spe_out = gr.Textbox(label="Energy & Forces", lines=15, interactive=False)
                    spe_status = gr.Textbox(label="Status", interactive=False, max_lines=1)
                    spe_viewer = Molecule3D(
                        label="Input Structure Viewer",
                        reps=UNIVERSAL_MOLECULE3D_REPS
                    )
                    
                    # Charge and Spin are only applicable to OMol and OMol-Direct
                    task_name_spe.input(
                        lambda x: (
                            (gr.Number(visible=True), gr.Number(visible=True))
                            if x == "OMol" or x == "OMol-Direct"
                            else (gr.Number(visible=False), gr.Number(visible=False))
                        ),
                        [task_name_spe],
                        [charge_input, spin_input],
                    )
                    
            run_spe.click(predict_molecule, [xyz_input, task_name_spe, charge_input, spin_input], [spe_out, spe_status, spe_viewer])

        # -------- MD --------
        with gr.Tab("Molecular Dynamics"):
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("## Molecular Dynamics Simulation")
                    gr.Markdown("Upload your molecular structure and configure MD parameters.")
                    
                    xyz_md = gr.File(
                        label="Upload Structure File (.xyz/.pdb/.cif/.traj)",
                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                        file_count="single"
                    )
                    with gr.Row():
                        task_name_md = gr.Radio(
                            ["OMol", "OMat", "OMol-Direct"],
                            value="OMol",
                            label="Model Type",
                            info="Choose the OrbMol model variant for the calculation."
                        )
                    with gr.Row():
                        charge_md = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
                        spin_md = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin Multiplicity")
                    with gr.Row():
                        steps_md = gr.Slider(minimum=10, maximum=2000, value=100, step=10, label="Steps")
                        temp_md = gr.Slider(minimum=10, maximum=1500, value=300, step=10, label="Temperature (K)")
                    with gr.Row():
                        timestep_md = gr.Slider(minimum=0.1, maximum=5.0, value=1.0, step=0.1, label="Timestep (fs)")
                        ensemble_md = gr.Radio(["NVE", "NVT"], value="NVE", label="Ensemble")
                    run_md_btn = gr.Button("Run MD Simulation", variant="primary")

                with gr.Column(variant="panel", min_width=520):
                    md_status = gr.Textbox(label="MD Status", interactive=False)
                    md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
                    md_viewer = Molecule3D(
                        label="Final Structure Viewer (Last MD Frame)",
                        reps=UNIVERSAL_MOLECULE3D_REPS
                    )
                    md_log = gr.Textbox(label="Log", interactive=False, lines=15, max_lines=25)
                    md_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30)
                    md_explain = gr.Markdown()
                    
                    # Charge and Spin are only applicable to OMol and OMol-Direct
                    task_name_md.input(
                        lambda x: (
                            (gr.Number(visible=True), gr.Number(visible=True))
                            if x == "OMol" or x == "OMol-Direct"
                            else (gr.Number(visible=False), gr.Number(visible=False))
                        ),
                        [task_name_md],
                        [charge_md, spin_md],
                    )
                    
            run_md_btn.click(
                md_wrapper,
                inputs=[xyz_md, task_name_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md],
                outputs=[md_status, md_traj, md_log, md_script, md_explain, md_viewer],
            )

        # -------- Relax --------
        with gr.Tab("Relaxation / Optimization"):
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("## Structure Relaxation/Optimization")
                    gr.Markdown("Upload your molecular structure for geometry optimization.")
                    
                    xyz_rlx = gr.File(
                        label="Upload Structure File (.xyz/.pdb/.cif/.traj)",
                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                        file_count="single"
                    )
                    with gr.Row():
                        task_name_rlx = gr.Radio(
                            ["OMol", "OMat", "OMol-Direct"],
                            value="OMol",
                            label="Model Type",
                            info="Choose the OrbMol model variant for the calculation."
                        )
                    with gr.Row():   
                        steps_rlx = gr.Slider(minimum=1, maximum=2000, value=300, step=1, label="Max Steps")
                        fmax_rlx = gr.Slider(minimum=0.001, maximum=0.5, value=0.05, step=0.001, label="Fmax (eV/Å)")
                    with gr.Row():
                        charge_rlx = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
                        spin_rlx = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin")
                    relax_cell = gr.Checkbox(False, label="Relax Unit Cell")
                    run_rlx_btn = gr.Button("Run Optimization", variant="primary")

                with gr.Column(variant="panel", min_width=520):
                    rlx_status = gr.Textbox(label="Status", interactive=False)
                    rlx_traj = gr.File(label="Trajectory (.traj)", interactive=False)
                    rlx_viewer = Molecule3D(
                        label="Optimized Structure Viewer",
                        reps=UNIVERSAL_MOLECULE3D_REPS
                    )
                    rlx_log = gr.Textbox(label="Log", interactive=False, lines=15, max_lines=25)
                    rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30)
                    rlx_explain = gr.Markdown()

                    # Charge and Spin are only applicable to OMol and OMol-Direct
                    task_name_rlx.input(
                        lambda x: (
                            (gr.Number(visible=True), gr.Number(visible=True))
                            if x == "OMol" or x == "OMol-Direct"
                            else (gr.Number(visible=False), gr.Number(visible=False))
                        ),
                        [task_name_rlx],
                        [charge_rlx, spin_rlx],
                    )
                    
            run_rlx_btn.click(
                relax_wrapper,
                inputs=[xyz_rlx, task_name_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell],
                outputs=[rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_viewer],
            )

print("Starting OrbMol model loading…")

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