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 输入错误")

    mol = Chem.MolFromSmiles(smiles)
    if mol == None:
        raise gr.Error("SMILES 输入错误")
    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="请输入化合物的SMILES")
            with gr.Row():
                btn = gr.Button(variant="primary",value="提交")
                clear_btn = gr.ClearButton(value="清除")
        with gr.Column():
            outputs=gr.Textbox(lines=1, label="该物质的CHO细胞毒性为：",info="单位：摩尔浓度")

    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],
        label="参考格式"
    )


demo.launch()