import gradio as gr

import os
import glob
import cv2
import numpy as np
import torch
from molscribe import MolScribe
from indigo import Indigo
from indigo.renderer import IndigoRenderer

from huggingface_hub import hf_hub_download

REPO_ID = "IDEA-AI4S/IDEA-OCSR"
FILENAME = "IDEA-OCSR-v1.0.0.pth"
ckpt_path = hf_hub_download(REPO_ID, FILENAME)

device = torch.device('cpu')
model = MolScribe(ckpt_path, device)


def generate_mol_image(molblock):
    indigo = Indigo()
    render = IndigoRenderer(indigo)
    indigo.setOption('render-output-format', 'png')
    indigo.setOption('render-background-color', '1,1,1')
    indigo.setOption('render-stereo-style', 'none')
    indigo.setOption('render-label-mode', 'hetero')
    mol = indigo.loadMolecule(molblock)
    buf = render.renderToBuffer(mol)
    img = cv2.imdecode(np.asarray(bytearray(buf), dtype=np.uint8), 1)
    return img


def predict(image):
    prediction = model.predict_image(image)
    smiles = prediction['smiles']
    molfile = prediction['molfile']
    image = generate_mol_image(molfile)
    return image, smiles, molfile


import gradio as gr
import glob

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(label="Upload molecular image", show_label=False, height=256),
    outputs=[
        gr.Image(label="Prediction", height=256),
        gr.Textbox(label="SMILES"),
        gr.Textbox(label="Molfile"),
    ],
    flagging_mode="auto",
    title="IDEA-OCSR",
    description=(
        "Convert a molecular image into SMILES and Molfile. "
        "It typically takes 2-3 seconds to predict an image, but may take longer if the server is busy. "
        "To view the prediction better, copy-paste the Molfile to ChemDraw.<br>"
    ),
    examples=sorted(glob.glob('examples/*.png')) if glob.glob('examples/*.png') else None,
    examples_per_page=20,
    cache_examples=False,
)

iface.launch()