import gradio as gr
import os
from rdkit import Chem
from motif_sample import demo
import tempfile
from rdkit.Chem import AllChem
import numpy as np
from PIL import Image
import io
from rdkit.Chem import Draw


# Function to serve the file via Gradio
def create_and_return_sdf(Protein_index: int):
    # Ensure input is an integer, as it's coming from the interface.
    number = Protein_index
    number = int(number)

    # Generate SDF file (you'll replace this with your actual logic)
    sdf_filename = demo(number)

    suppl = Chem.SDMolSupplier(sdf_filename)
    mol = next(suppl)
    # AllChem.UFFOptimizeMolecule(mol)
    mol = Chem.MolFromSmiles(Chem.MolToSmiles(mol))

    for atom in mol.GetAtoms():
        atom.SetAtomMapNum(0)

    mol_image = Draw.MolToImage(mol)

    np_image = np.array(mol_image)
    np_image = np_image[:, :, :3]

    return np_image, sdf_filename


# Define Gradio interface
iface = gr.Interface(
    fn=create_and_return_sdf,
    inputs="text",
    outputs=[
        gr.outputs.Image(type="numpy", label="Molecule Image"),
        gr.outputs.File(label="Download SDF")
    ],
    live=False  # The function should only be called when the user submits the form
)

# Launch the interface
iface.launch(share=True)