import spaces


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

import os
import urllib.request

CCD_URL = "https://huggingface.co/boltz-community/boltz-1/resolve/main/ccd.pkl"
MODEL_URL = "https://huggingface.co/boltz-community/boltz-1/resolve/main/boltz1.ckpt"

cache = "/home/user/.boltz"

os.makedirs(cache)

ccd = f"{cache}/ccd.pkl"
if not os.path.exists(ccd):
    print(
        f"Downloading the CCD dictionary to {ccd}. You may "
    )
    urllib.request.urlretrieve(CCD_URL, str(ccd))

# Download model
model =f"{cache}/boltz1.ckpt"
if not os.path.exists(model):
    print(
        f"Downloading the model weights to {model}"
    )
    urllib.request.urlretrieve(MODEL_URL, str(model))



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

    os.system("boltz predict ligand.fasta --output_format pdb")
    return "boltz_results_ligand/predictions/ligand/ligand_model_0.pdb"

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)