import gradio as gr
from src.predict import predict

def update_amino_acid(safety_param, amino_acid):
    if safety_param == "KIBA Score":
        return gr.update(visible=True, value=amino_acid)
    else:
        return gr.update(visible=False)

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            safety_param = gr.Dropdown(["Skin Reaction", "KIBA Score", "Liver Safety"], label="Parameter")
            drug_name = gr.Textbox(label="Drug Name")
            amino_acid = gr.Textbox(label="Amino Acid Name", visible=False)
        with gr.Column():
            safety_report = gr.Textbox(label="Drug Analysis")

    analyze_drug_btn = gr.Button("Get Drug Analysis")


    analyze_drug_btn.click(
        fn=predict,
        inputs=[safety_param, drug_name, amino_acid],
        outputs=safety_report,
    )

    safety_param.change(
        fn=update_amino_acid,
        inputs=[safety_param, amino_acid],
        outputs=amino_acid,
    )

    examples = [
        ["KIBA Score", "C1=CC=C(C=C1)C(=O)N", "RPDFCLEPPYTGPCKARIIRYFYNAKAGLCQTFVYGGCRAKRNNFKSAEDCMRTCGGA"],
        ["KIBA Score", "C1=CC=C(C=C1)C(=O)N", "Humanin"],
        ["Liver Safety", "benzamide", None],
        ["Liver Safety", "CN1C(=O)CN=C(C2=CCCCC2)c2cc(Cl)ccc21", None],
        ["Skin Reaction", "Amoxicillin", None],
        ["Skin Reaction", "Atorvastatin", None]
    ]

    gr.Examples(
        examples=examples,
        inputs=[safety_param, drug_name, amino_acid],
    )

if __name__ == "__main__":
    demo.queue(max_size=10).launch()