import os

import spaces
import gradio as gr 



from gradio_molecule3d import Molecule3D

os.system("pip install allmetal3d-0.7.4-py3-none-any.whl")

os.system("mkdir -p /home/user/.local/share/allmetal3d/weights")

os.system("mv *.pth /home/user/.local/share/allmetal3d/weights/")


from allmetal3d.utils.helpersNew import *
from allmetal3d.utils.main import predict, predict_cli, update_mode

@spaces.GPU
def predict_zero_gpu(pdb, models, pthreshold, threshold, batch_size, mode, central_residue, radius):
    return predict(pdb, models, pthreshold, threshold, batch_size, mode, central_residue, radius)

with gr.Blocks() as blocks: 
    gr.Markdown("## AllMetal3D and Water3D")
    pdb = Molecule3D(label="Input PDB", showviewer=False) #gr.File("2cba.pdb",label="Upload PDB file")
    gr.Markdown("Metals might bind anywhere in the protein, choose how to sample the residues in the protein. <br>Fast uses blocked sampling of residues to reduce required computational time, full uses all residues, site allows you to look around a specific site in the protein")
    with gr.Row("Prediction mode"):
        mode = gr.Dropdown(["fast", "all", "site"], value="fast", label="Mode")
        central_residue = gr.Textbox(label="Central residue", info="add multiple residues with space e.g 101 203", visible=False)
        radius = gr.Slider(value=8, minimum=4, maximum=50, label="Distance threshold", visible=False)
    mode.change(update_mode, mode, [central_residue, radius])
    models = gr.Radio(["AllMetal3D + Water3D", "Only AllMetal3D", "Only Water3D"], value="AllMetal3D + Water3D", label="Which models to run?")
    with gr.Accordion("Settings"):
        threshold = gr.Slider(value=7,minimum=0, maximum=10,  label="Threshold")
        pthreshold = gr.Slider(value=0.25,minimum=0.1, maximum=1,  label="Probability Threshold")
        batch_size = gr.Slider(value=50, minimum=0, maximum=100, label="Batch Size")
    btn = gr.Button("Predict")
    with gr.Row():
        metal_pdb = gr.File(label="predicted metals (PDB)", visible=False)
        metal_cube = gr.File(label="predicted metal density (CUBE)",visible=False)
        water_pdb = gr.File(label="predicted waters (PDB)", visible=False)
        water_cube = gr.File(label="predicted water density (CUBE)", visible=False)
    out = gr.HTML("")
    results_json = gr.JSON(visible=False)
    waters_json = gr.JSON(visible=False)
    btn.click(predict_zero_gpu, inputs=[pdb, models, pthreshold, threshold, batch_size, mode, central_residue, radius], outputs=[out, metal_pdb, metal_cube, water_pdb, water_cube, results_json, waters_json])

blocks.launch(allowed_paths=["frontend"])