# import os
# import py3Dmol
# from rdkit import Chem


# def visualize_complex(pdb_block, sdf_block, show_protein_surface=True, show_ligand=True, show_ligand_surface=True):
#     view = py3Dmol.view()

#     # Add protein to the canvas
#     view.addModel(pdb_block, 'pdb')
#     if show_protein_surface:
#         view.addSurface(py3Dmol.VDW, {'opacity': 0.7, 'color': 'white'}, {'model': -1})
#     else:
#         view.setStyle({'model': -1}, {'cartoon': {'color': 'spectrum'}, 'line': {}})
#     view.setStyle({'model': -1}, {"cartoon": {"style": "edged", 'opacity': 0}})

#     # Add ligand to the canvas
#     if show_ligand:
#         view.addModel(sdf_block, 'sdf')
#         view.setStyle({'model': -1}, {'stick': {}})
#         # view.setStyle({'model': -1}, {'cartoon': {'color': 'spectrum'}, 'line': {}})
#         if show_ligand_surface:
#             view.addSurface(py3Dmol.VDW, {'opacity': 0.8}, {'model': -1})

#     view.zoomTo()
#     return view


# def visualize_data(data, root, show_ligand=True, show_surface=True):
#     protein_path = os.path.join(root, data.protein_filename)
#     ligand_path = os.path.join(root, data.ligand_filename)
#     with open(protein_path, 'r') as f:
#         pdb_block = f.read()
#     with open(ligand_path, 'r') as f:
#         sdf_block = f.read()
#     return visualize_complex(pdb_block, sdf_block, show_ligand=show_ligand, show_surface=show_surface)


# def visualize_generated_mol(protein_filename, mol, root, show_surface=False, opacity=0.5):
#     protein_path = os.path.join(root, protein_filename)
#     with open(protein_path, 'r') as f:
#         pdb_block = f.read()

#     view = py3Dmol.view()

#     # Add protein to the canvas
#     view.addModel(pdb_block, 'pdb')
#     view.setStyle({'model': -1}, {'cartoon': {'color': 'spectrum'}, 'line': {}})

#     mblock = Chem.MolToMolBlock(mol)
#     view.addModel(mblock, 'mol')
#     view.setStyle({'model': -1}, {'stick': {}, 'sphere': {'radius': 0.35}})
#     if show_surface:
#         view.addSurface(py3Dmol.SAS, {'opacity': opacity}, {'model': -1})

#     view.zoomTo()
#     return view


# def MolTo3DView(mol, size=(300, 300), style="stick", surface=False, opacity=0.5):
#     """Draw molecule in 3D

#     Args:
#     ----
#         mol: rdMol, molecule to show
#         size: tuple(int, int), canvas size
#         style: str, type of drawing molecule
#                style can be 'line', 'stick', 'sphere', 'carton'
#         surface, bool, display SAS
#         opacity, float, opacity of surface, range 0.0-1.0
#     Return:
#     ----
#         viewer: py3Dmol.view, a class for constructing embedded 3Dmol.js views in ipython notebooks.
#     """
#     assert style in ('line', 'stick', 'sphere', 'carton')

#     viewer = py3Dmol.view(width=size[0], height=size[1])
#     if isinstance(mol, list):
#         for i, m in enumerate(mol):
#             mblock = Chem.MolToMolBlock(m)
#             viewer.addModel(mblock, 'mol' + str(i))
#     elif len(mol.GetConformers()) > 1:
#         for i in range(len(mol.GetConformers())):
#             mblock = Chem.MolToMolBlock(mol, confId=i)
#             viewer.addModel(mblock, 'mol' + str(i))
#     else:
#         mblock = Chem.MolToMolBlock(mol)
#         viewer.addModel(mblock, 'mol')
#     viewer.setStyle({style: {}})
#     if surface:
#         viewer.addSurface(py3Dmol.SAS, {'opacity': opacity})
#     viewer.zoomTo()
#     return viewer