import gradio as gr
import shutil
from rdkit import Chem
from rdkit.Chem import AllChem
import pandas as pd
from ase import Atoms
from ase.optimize import BFGS
from ase.vibrations import Vibrations
from ase.thermochemistry import IdealGasThermo
from fairchem.core import FAIRChemCalculator, pretrained_mlip

predictor = pretrained_mlip.get_predict_unit("uma-s-1", device="cpu")
calculator = FAIRChemCalculator(predictor, task_name="omol")
model = "UMA-OMOL"

def calc_vibrations(smile: str, struct_type: str) -> str:
  """Calculate the vibrational frequencies of a molecule 
    using the META UMA MLIP, starting from a SMILES string.
    
    Args:
        smile: The SMILES string of the molecule
        struct_type: The type of the molecule (linear or non-linear)
        
    Returns:
        The vibrational frequencies of the molecule in cm-1, the low
        frequency motions (typically corresponding to translation or
        rotation), and the number of imaginary frequencies in a text string
  """

  mol = Chem.MolFromSmiles(smile)
  molH = Chem.AddHs(mol)
  AllChem.EmbedMolecule(molH)
  AllChem.MMFFOptimizeMolecule(molH)

  xyz_list = []
  atoms_list = ""
  for atom in molH.GetAtoms():
    atoms_list += atom.GetSymbol()
    pos = molH.GetConformer().GetAtomPosition(atom.GetIdx())
    temp_tuple = (pos[0], pos[1], pos[2])
    xyz_list.append(temp_tuple)

  atoms = Atoms(atoms_list,xyz_list) 

  atoms.calc = calculator
  energy = atoms.get_potential_energy()
  

  dyn = BFGS(atoms)
  dyn.run(fmax=0.05)
  final_energy = atoms.get_potential_energy()

  vib = Vibrations(atoms)
  vib.run()
  vib_energies = vib.get_energies()

  real_vibs = []
  number_imaginary = 0
  for vib in vib_energies:
    if vib.imag < 0.00001:
      real_vibs.append(8065.56*vib.real.item())
    else:
      number_imaginary += 1
  
  num_atoms = molH.GetNumAtoms()
  if struct_type == "non-linear" or struct_type == "nonlinear":
    non_vib_dof = 6 - number_imaginary
  else:
    non_vib_dof = 5 - number_imaginary
  vib_dof = 3*num_atoms - non_vib_dof

  out_text = ""
  for i,vib in enumerate(real_vibs[3*num_atoms - vib_dof:]):
    out_text += f"Vibrational frequency {i+1}: {vib:.3f} cm-1\n"
  
  out_text += f"Also calculated the following low frequency motions:\n"
  for i,vib in enumerate(real_vibs[:3*num_atoms - vib_dof]):
    out_text += f"Frequency {i+1}: {vib:.3f} cm-1\n"
  
  out_text += f"Also found the following number of imaginary frequencies: {number_imaginary}"

  shutil.rmtree('vib')

  return out_text

gradio_app = gr.Interface(
    calc_vibrations,
    inputs=[gr.Textbox(label="SMILES to calculate"),gr.Radio(label="Type of molecule", choices=["linear", "non-linear"])],
    outputs=[gr.Textbox(label="Frequencies: ")],
    title="Get Vibrational frequencies of a molecule from SMILES using the META UMA MLIP",
)

if __name__ == "__main__":
    gradio_app.launch(mcp_server=True)