import streamlit as st
# for the small mols/editor ------
from streamlit_ketcher import st_ketcher

# for the big mols --------
# based on: https://napoles-uach-stmol-home-pom051.streamlit.app/Examples
# further documentation on params etc.:
# https://github.com/napoles-uach/stmol/blob/master/functions.md
# https://github.com/napoles-uach/stmol/blob/master/pages/3_%F0%9F%93%9A_Examples.py
# https://github.com/napoles-uach/stmol/blob/master/pages/2_%E2%9C%A8_Demo%20stmol.py

# from stmol import showmol
from stmol import *

import py3Dmol

# 1A2C
# Structure of thrombin inhibited by AERUGINOSIN298-A from a BLUE-GREEN ALGA
# xyzview = py3Dmol.view(query='pdb:3LPP') 
# xyzview.setStyle({'cartoon':{'color':'spectrum'}})
# showmol(xyzview, height = 500,width=800)

# EZ OPTION, it's this guy's wrapper, not too customizable
# showmol(render_pdb(id = '1A2C'))

# PARAMS
# pdbID = '3LPP' # add interactivity
# backgroundColor = 'black'


# view = py3Dmol.view(query='pdb:'+pdbID)
# # make the mol rainbow
# view.setStyle({style:{
#     'color':'spectrum'
# }})
# # set the background
# # can also take args etc. but w/e, accepts str object
# view.setBackgroundColor(backgroundColor)



# Code Block
# prot_str='1A2C,1BML,1D5M,1D5X,1D5Z,1D6E,1DEE,1E9F,1FC2,1FCC,1G4U,1GZS,1HE1,1HEZ,1HQR,1HXY,1IBX,1JBU,1JWM,1JWS'
# prot_list=prot_str.split(',')
# bcolor = st.color_picker('Pick A Color','#89cff0')
# protein=st.selectbox('select protein',prot_list)
# style = st.selectbox('style',['cartoon','line','cross','stick','sphere'])
# xyzview = py3Dmol.view(query='pdb:'+protein)
# xyzview.setStyle({style:{'color':'spectrum'}})
# xyzview.setBackgroundColor(bcolor)
# showmol(xyzview, height = 500,width=800)

# CSID:
# 3LPO other option
with st.echo():
    # try to replicate Chimera/PyMol behavior
    pdbCode = '3LPP'
    backgroundColor = "black"
    showmol(render_pdb_resi(
        viewer = render_pdb(id= '3LPP'),
        resi_lst=['']
    ))
    

# showmol(render_pdb(id="3LPP"))

# showmol(xyz)



# And small mol app, taken from: 
# https://github.com/streamlit/mol-demo/blob/main/streamlit_app.py

DEFAULT_COMPOUND = "CHEMBL141739"

if "molfile" not in st.session_state:import streamlit as st

from streamlit_ketcher import st_ketcher

# st.set_page_config(layout="wide")
# st.title("`streamlit-ketcher`")

st.header("Component with user input")
# DEFAULT_MOL = (
#     # r"C[N+]1=CC=C(/C2=C3\C=CC(=N3)/C(C3=CC=CC(C(N)=O)=C3)=C3/C=C/C(=C(\C4=CC=[N+]"
#     # "(C)C=C4)C4=N/C(=C(/C5=CC=CC(C(N)=O)=C5)C5=CC=C2N5)C=C4)N3)C=C1"
#     "CNCC(=O)c1ccc2c(c1)OCO2" #serotonin
# )
with st.echo():
    DEFAULT_MOL = "CNCC(=O)c1ccc2c(c1)OCO2" #serotonin
    molecule = st.text_input("Molecule", DEFAULT_MOL)
    smile_code = st_ketcher(molecule)
    st.markdown(f"Smile code: ``{smile_code}``")

st.write("---")

# st.header("Components with custom height")
# with st.echo():
#     st_ketcher("CCO", height=400)
#     st_ketcher("CCO", height=800)

# st.header("Component with `molfile` format")
# with st.echo():
#     molfile = st_ketcher(molecule_format="MOLFILE")
#     st.markdown("molfile:")
#     st.code(molfile)