import spaces


import gradio as gr
from gradio_molecule3d import Molecule3D
from gradio_cofoldinginput import CofoldingInput

import os

@spaces.GPU(duration=200)
def predict(jobname, inputs, recycling_steps, sampling_steps, diffusion_samples):

    os.system("boltz predict ligand.fasta")
    return "boltz_results_ligand/predictions/ligand/ligand_model_0.cif"

with gr.Blocks() as blocks:
    gr.Markdown("# Boltz-1")
    with gr.Tab("Main"):
        jobname = gr.Textbox(label="Jobname")
        inp = CofoldingInput(label="Input")
        out = Molecule3D(label="Output")
    with gr.Tab("Settings"):
        recycling_steps =gr.Slider(value=3, minimum=0, label="Recycling steps")
        sampling_steps = gr.Slider(value=200, minimum=0, label="Sampling steps")
        diffusion_samples = gr.Slider(value=1, label="Diffusion samples")

    btn = gr.Button("predict")

    btn.click(fn=predict, inputs=[jobname,inp, recycling_steps, sampling_steps, diffusion_samples], outputs=[out],  api_name="predict")

blocks.launch(ssr_mode=False)