from stmol import showmol
import py3Dmol
import streamlit as st
import requests

sequence = st.text_input(label="Protein Sequence", value="MGSSHHHHHHSSGLVPRGSHMRGPNPTAASLEASAGPFTVRSFTVSRPSGYGAGTVYYPTNAGGTVGAIAIVPGYTARQSSIKWWGPRLASHGFVVITIDTNSTLDQPSSRSSQQMAALRQVASLNGTSSSPIYGKVDTARMGVMGWSMGGGGSLISAANNPSLKAAAPQAPWDSSTNFSSVTVPTLIFACENDSIAPVNSSALPIYDSMSRNAKQFLEINGGSHSCANSGNSNQALIGKKGVAWMKRFMDNDTRYSTFACENPNSTRVSDFRTANCSLEDPAANKARKEAELAAATAEQ")

headers = {
    'Content-Type': 'application/x-www-form-urlencoded',
}

response = requests.post('https://api.esmatlas.com/foldSequence/v1/pdb/', headers=headers, data=sequence)
name = sequence[:3] + sequence[-3:] 
pdb_filename = "test.pdb"
pdb_string = response.content.decode('utf-8')
with open(pdb_filename, "w") as out:
  out.write(pdb_string)

# 1A2C
# Structure of thrombin inhibited by AERUGINOSIN298-A from a BLUE-GREEN ALGA
xyzview = py3Dmol.view(width=400, height=300)
xyzview.addModelsAsFrames(pdb_filename)
xyzview.setStyle({'cartoon':{'color':'spectrum'}})
print("Set style")
showmol(xyzview, height = 500,width=800)
print("Displayed")