import gradio as gr
import pandas as pd
from rdkit import Chem
from rdkit.Chem import AllChem
import joblib

model = joblib.load('CHO.pkl')

def predict(smiles):
    if smiles.strip() == "":
        raise gr.Error("SMILES input error")

    mol = Chem.MolFromSmiles(smiles)
    if mol == None:
        raise gr.Error("SMILES input error")
    mol_ECFP4 = list(AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=1024).ToBitString())
    preprocess_data = pd.DataFrame([mol_ECFP4])
    result = model.predict(preprocess_data)
    postprocess_data = '{:.2e}'.format(pow(10, result[0]))
    return postprocess_data


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            inputs=gr.Textbox(lines=2, label="Please enter SMILES for the compound")
            with gr.Row():
                btn = gr.Button(variant="primary",value="submit")
                clear_btn = gr.ClearButton(value="clear")
        with gr.Column():
            outputs=gr.Textbox(lines=1, label="Predicted CHO cytotoxicity of the chemical is:",info="Unit: mol/L")

    btn.click(predict, inputs=[inputs], outputs=[outputs])
    clear_btn.add([inputs,outputs])

    gr.Examples(
        [["O=C(O)CBr"],["O=CC(Br)(Br)Br"],["IC(Br)Br"]],
        [inputs],
    )


demo.launch()