import gradio as gr
from gradio_molecule3d import Molecule3D
import deepchem as dc
import os

example = Molecule3D().example_inputs()


reps =    [
    {
      "model": 0,
      "chain": "",
      "resname": "",
      "style": "stick",
      "color": "whiteCarbon",
      "residue_range": "",
      "around": 0,
      "height": 150,
      "byres": False,
      "visible": False
    }
  ]

def get_file_name_without_extension(file_path):
    base_name = os.path.basename(file_path)
    file_name = os.path.splitext(base_name)[0]
    return file_name


def predict(input_receptor,input_ligand):
    print(input_ligand)
    vpg = dc.dock.pose_generation.VinaPoseGenerator()
    saved_filename = get_file_name_without_extension(input_ligand)
    complexes, scores = vpg.generate_poses(molecular_complex=(input_receptor,input_ligand),  # protein-ligand files for docking,
                                            out_dir='./output_directory',
                                            generate_scores=True
                                            )

    # print(complexes[0])
    return scores


with gr.Blocks() as docking:
    inp = [Molecule3D(label="Receptor", reps=reps),Molecule3D(label="Ligand", reps=reps)]
    out = gr.Textbox(label="Docking Output")#Molecule3D(label="Output", reps=reps)
    # out = Molecule3D(label="Output", reps=reps)
    btn = gr.Button("Docking")
    btn.click(predict, inputs=inp, outputs=out)


if __name__ == '__main__':
  docking.launch()