import gradio as gr
from gradio_molecule2d import molecule2d
from synformer.chem.mol import Molecule
from synformer.sampler.analog.parallel import run_sampling_one_cpu
from huggingface_hub import hf_hub_download

REPO_ID = "whgao/synformer"
CKPT_FILENAME = "sf_ed_default.ckpt"
MAT_FILENAME = "matrix.pkl"
FPI_FILENAME = "fpindex.pkl"
ckpt_path = hf_hub_download(REPO_ID, CKPT_FILENAME)
mat_path = hf_hub_download(REPO_ID, MAT_FILENAME)
fpi_path = hf_hub_download(REPO_ID, FPI_FILENAME)

last_result = {}

# Function to clear all inputs
def clear_inputs():
    # Return default or empty values to reset each input component
    return None, 24, 64, 0

def sample(smi, search_width, exhaustiveness):
    result_df = run_sampling_one_cpu(
        input=Molecule(smi),
        model_path=ckpt_path,
        mat_path=mat_path,
        fpi_path=fpi_path,
        search_width=search_width,
        exhaustiveness=exhaustiveness,
        time_limit=180,
        max_results=100,
        max_evolve_steps=24,
        sort_by_scores=True,
    )
    result_df = result_df[:30]
    last_result["results_df"] = result_df
    smiles = result_df.iloc[0]["smiles"]
    similarity = result_df.iloc[0]["score"]
    synthesis = result_df.iloc[0]["synthesis"]
    return smiles, similarity, synthesis, gr.update(maximum=len(result_df)-1)


def select_from_output(index):
    df_results = last_result["results_df"]
    return df_results.iloc[index]["smiles"], df_results.iloc[index]["score"], df_results.iloc[index]["synthesis"]

examples = [
    "Nc1cccc(S(=O)(=O)N2CCCN(S(=O)(=O)c3ccc4c(c3)OCCO4)CC2)c1",
    "CN1C[C@H](Nc2cnn(C)c(=O)c2)C[C@H](c2ccccc2)C1",
    "COc1ccc(-c2ccnc(Nc3ccccc3)n2)cc1",
    "CC[C@@H]1OC[C@@]23Cc4cc(F)c(N)cc4-c4ccc5c(c42)C(=CC(F)(F)O5)[C@@H]1C3=O",
    "O=C(OCC(=O)N1[C@H](C(=O)O)C[C@@H]2CCCC[C@@H]21)[C@H](Cc1cbccc1)NC(I)c1bcccc1",
]


with gr.Blocks() as demo:
    gr.Markdown(f"""
        # Demo of [SynFormer](https://github.com/wenhao-gao/synformer/tree/main)
        This page demonstrates the SynFormer-ED model, which takes a molecule as input—regardless of its synthetic accessibility—and outputs 
                identical or approximate molecules along with their associated synthetic paths. The demo runs on CPUs and typically takes about 
                one minute per run but can be accelerated by reducing the search width and exhaustiveness. The model may take longer if the server 
                is busy. Since the sampling is stochastic, you may run the demo multiple times to explore different results, with a maximum of 
                30 molecules displayed at once.
        To learn more about SynFormer’s architecture and applications, check out [our paper](https://github.com/wenhao-gao/synformer/tree/main).
                
        Authors: [Wenhao Gao](mailto:gaowh19@gmail.com), Shitong Luo, Connor W. Coley
        """)
    with gr.Row():
        with gr.Column(scale=0.5):
            input_molecule = molecule2d(label="SMILES Input")
            slider_1 = gr.Slider(minimum=1, maximum=100, step=1, label="Search Width", value=24)
            slider_2 = gr.Slider(minimum=1, maximum=100, step=1, label="Exhaustiveness", value=64)

            with gr.Row():
                with gr.Column(scale=0.5):
                    run_btn = gr.Button("Run on sample")
                with gr.Column(scale=0.5):
                    clear_btn = gr.Button("Clear")

        with gr.Column(scale=0.5):
            index_slider = gr.Slider(minimum=0, maximum=10, step=1, label="Select Output Index", value=0, interactive=True) 
            output_similarity = gr.Text(label="Tanimoto Similarity")
            output_molecule = molecule2d(label="Output")
            output_synpath = gr.Textbox(label="Synthetic Path")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Examples")
            gr.Examples(
                examples = examples,
                inputs = [input_molecule]
            )
    
    run_btn.click(
        fn=sample, 
        inputs=[
            input_molecule, 
            slider_1, 
            slider_2
        ], 
        outputs=[
            output_molecule, 
            output_similarity,
            output_synpath,
            index_slider
        ],
        api_name="Run"
    )

    index_slider.change(
        fn=select_from_output,
        inputs=[index_slider],
        outputs=[output_molecule, output_similarity, output_synpath],
    )

    clear_btn.click(
        fn=clear_inputs,
        inputs=[],
        outputs=[input_molecule, slider_1, slider_2, index_slider]
    )

demo.launch()