import streamlit as st

st.set_page_config(layout="wide")
st.header("Demo: Proteins in a grid")

sequences = """MINDLLDISRIISGKMTLDRAEVNLTAIARQVVEEQRQAAEAKSIQLLCSTPDTNHYVFGDFDRLKQTLWNLLSNAVKFTPSGGTVELELGY
MQGDSSISSSNRMFTLCKPLTVANETSTLSTTRNSKSNKRVSKQRVNLAESPERNAPSPASIKTNETEEFSTIKTTNNEVLGYEPNYVSYDF
MSTHVSLENTLASLQATFFSLEARHTALETQLLSTRTELAATKQELVRVQAEISRADAQAQDLKAQILTLKEKADQAEVEAAAATQRAEESQ
MVLLSTGPLPILFLGPSLAELNQKYQVVSDTLLRFTNTVTFNTLKFLGSDS
MNNDEQPFIMSTSGYAGNTTSSMNSTSDFNTNNKSNTWSNRFSNFIAYFSGVGWFIGAISVIFFIIYVIVFLSRKTKPSGQKQYSRTERNNR
MEAVYSFTITETGTGTVEVTPLDRTISGADIVYPPDTACVPLTVQPVINANGTWTLGSGCTGHFSVDTTGHVNCLTGGFGAAGVHTVIYTVE
MGLTTSGGARGFCSLAVLQELVPRPELLFVIDRAFHSGKHAVDMQVVDQEGLGDGVATLLYAHQGLYTCLLQAEARLLGREWAAVPALEPNF
MGLTTSGGARGFCSLAVLQELVPRPELLFVIDRAFHSGKHAVDMQVVDQEGLGDGVATLLYAHQGLYTCLLQAEARLLGREWAAVPALEPNF
MGAAGYTGSLILAALKQNPDIAVYALNRNDEKLKDVCGQYSNLKGQVCDLSNESQVEALLSGPRKTVVNLVGPYSFYGSRVLNACIEANCHY""".split(
    "\n"
)

import time

import requests


@st.cache_data
def get_pdb(sequence):
    retries = 0
    pdb_str = None
    url = "https://api.esmatlas.com/foldSequence/v1/pdb/"
    while retries < 3 and pdb_str is None:
        response = requests.post(url, data=sequence)
        pdb_str = response.text
        if pdb_str == "INTERNAL SERVER ERROR":
            retries += 1
            time.sleep(0.1)
            pdb_str = None
    return pdb_str


pdb_strings = []

for seq in sequences:
    pdb = get_pdb(seq)
    if pdb:
        pdb_strings.append(pdb)

grid_size = 3

import py3Dmol

dim = st.number_input("Viewer dim (pk)", value=1300)
color = st.sidebar.selectbox("Color", ["spectrum", "black", "white"])
st.markdown("""
# See protein, cool result
Herer are my things
""")

view = py3Dmol.view(width=dim, height=dim, viewergrid=(grid_size, grid_size))

for i, pdb in enumerate(pdb_strings):
    view.addModel(pdb, "pdb", viewer=(i // grid_size, i % grid_size))

view.zoomTo()
view.setStyle({"cartoon": {"color": color}})

import stmol

stmol.showmol(view, width=dim, height=dim)