import streamlit as st
import py3Dmol
from rdkit import Chem
from rdkit.Chem import Draw
from PIL import Image
from rdkit import Chem
from rdkit.Chem import AllChem

def smi2conf(smiles):
    '''Convert SMILES to rdkit.Mol with 3D coordinates'''
    mol = Chem.MolFromSmiles(smiles)
    if mol is not None:
        mol = Chem.AddHs(mol)
        AllChem.EmbedMolecule(mol)
        AllChem.MMFFOptimizeMolecule(mol, maxIters=200)
        return mol
    else:
        return None

smi = 'COc3nc(OCc2ccc(C#N)c(c1ccc(C(=O)O)cc1)c2P(=O)(O)O)ccc3C[NH2+]CC(I)NC(=O)C(F)(Cl)Br'
conf = smi2conf(smi)
viewer = MolTo3DView(conf, size=(600, 300), style='sphere')
viewer.show()

compound_smiles = 'c1cc(C(=O)O)c(OC(=O)C)cc1'
m = Chem.MolFromSmiles(compound_smiles)
im=Draw.MolToImage(m)
st.image(im)


def MakeMolecule(name, ingredients):
  st.write(name)
  m = Chem.MolFromSmiles(ingredients)
  im=Draw.MolToImage(m)
  st.image(im)

MakeMolecule("Ethanol", "CCO")
MakeMolecule("Acetic acid", "CC(=O)O")
MakeMolecule("Cyclohexane", "C1CCCCC1")
MakeMolecule("Pyridine", "c1cnccc1")


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')
    mblock = Chem.MolToMolBlock(mol)
    viewer = py3Dmol.view(width=size[0], height=size[1])
    viewer.addModel(mblock, 'mol')
    viewer.setStyle({style:{}})
    if surface:
        viewer.addSurface(py3Dmol.SAS, {'opacity': opacity})
    viewer.zoomTo()
    return viewer
    
viewer = MolTo3DView(conf, size=(600, 300), style='sphere')
viewer.show()

from ipywidgets import interact,fixed,IntSlider
import ipywidgets

smis = [ 'COc3nc(OCc2ccc(C#N)c(c1ccc(C(=O)O)cc1)c2P(=O)(O)O)ccc3C[NH2+]CC(I)NC(=O)C(F)(Cl)Br',
    'CC(NCCNCC1=CC=C(OCC2=C(C)C(C3=CC=CC=C3)=CC=C2)N=C1OC)=O',
    'Cc1c(COc2cc(OCc3cccc(c3)C#N)c(CN3C[C@H](O)C[C@H]3C(O)=O)cc2Cl)cccc1-c1ccc2OCCOc2c1',
    'CCCCC(=O)NCCCCC(=O)NCCCCCC(=O)[O-]',
    "CC(NCCNCC1=CC=C(OCC2=C(C)C(C3=CC=CC=C3)=CC=C2)N=C1OC)=O"]
    
confs = [smi2conf(s) for s in smis]

def conf_viewer(idx):
    mol = confs[idx]
    return MolTo3DView(mol).show()

interact(conf_viewer, idx=ipywidgets.IntSlider(min=0,max=len(class_0_list)-1, step=1))

from ipywidgets import interact,fixed,IntSlider
import ipywidgets

smis = [ 'COc3nc(OCc2ccc(C#N)c(c1ccc(C(=O)O)cc1)c2P(=O)(O)O)ccc3C[NH2+]CC(I)NC(=O)C(F)(Cl)Br',
    'CC(NCCNCC1=CC=C(OCC2=C(C)C(C3=CC=CC=C3)=CC=C2)N=C1OC)=O',
    'Cc1c(COc2cc(OCc3cccc(c3)C#N)c(CN3C[C@H](O)C[C@H]3C(O)=O)cc2Cl)cccc1-c1ccc2OCCOc2c1',
    'CCCCC(=O)NCCCCC(=O)NCCCCCC(=O)[O-]',
    "CC(NCCNCC1=CC=C(OCC2=C(C)C(C3=CC=CC=C3)=CC=C2)N=C1OC)=O"]

confs = [smi2conf(s) for s in smis]

def style_selector(idx, s):
    conf = confs[idx]
    return MolTo3DView(conf, style=s).show()

interact(style_selector, 
         idx=ipywidgets.IntSlider(min=0,max=len(class_0_list)-1, step=1),
         s=ipywidgets.Dropdown(
            options=['line', 'stick', 'sphere'],
            value='line',
            description='Style:'))
            
            
            
@interact
def smi2viewer(smi='CC=O'):
    try:
        conf = smi2conf(smi)
        return MolTo3DView(conf).show()
    except:
        return None