import gradio as gr
import glob
# pip install gradio-molecule3d
from gradio_molecule3d import Molecule3D

files = [
  "../MolCRAFT/data/test_set/ABL2_HUMAN_274_551_0/4xli_B_rec.pdb", 
  "../MolCRAFT/data/test_set/ABL2_HUMAN_274_551_0/4xli_B_rec_4xli_1n1_lig_tt_min_0.sdf",
]
files += sorted(glob.glob('output/*.sdf'))


from rdkit import Chem

# for f in files[1:]:
#   mol = Chem.SDMolSupplier(f, removeHs=False)[0]
#   # get coordinates of rdkit mol object
#   coords = mol.GetConformer().GetPositions()
#   print(coords.mean(axis=0))

# exit(0)

reps = [
    {
      "model": 0,
      "style": "cartoon",
      "color": "whiteCarbon",
    },
    {
      "model": 1,
      "style": "stick",
      "color": "redCarbon",
      "residue_range": "",
    },
] + [
    {
      "model": _ + 2,
      "style": "stick",
      "color": "greenCarbon",
      "residue_range": "",
    } for _ in range(10)
]

with gr.Blocks() as app:
  Molecule3D(files, reps=reps)

app.launch()