from chroma import api
from chroma import Chroma,conditioners,Protein
import streamlit as st
from stmol import *


# Register your Chroma API key
api.register_key("2cdade6d058b4fd1b85fa5badb501312")

def load_message(protein):
    st.header("Protein Information:")
    st.write(f"Device: {protein.device}")
    st.write(f"Protein Length: {len(protein)} residues")
    st.write(f"Structured Residue Count: {protein.length(structured=True)}")

    # 显示 Protein 的序列
    st.header("Protein Sequence:")
    protein_sequence = protein.sequence(format="one-letter-string")
    st.write(protein_sequence)

# Display title
st.title("Unconditional sample display")

st.write("slide to choose your chain_lengths:")
chain_lengths = st.slider(label="chain_lengths", min_value=10, max_value=500, key="chain_lengths")

#Generate protein sample

# 显示 Protein 的一些信息
load_message(Protein('sample.cif'))

if st.button("Run Code with Button",key="Unconditional"):
    chroma = Chroma()
    protein = chroma.sample(chain_lengths=[chain_lengths])
    protein.to("sample.pdb")
    load_message(protein)

# 显示 Protein 的结构
with open("sample.pdb", "r") as file:
    pdb_content = file.read()
obj = makeobj(pdb_content)

# 使用 stmol 展示蛋白质结构
st.header("Protein Structure:")
showmol(obj, width=1800)

st.title("Conditional sample display")

st.write("slide to choose your num_chain_neighbors:")
num_chain_neighbors = st.slider(label="num_chain_neighbors",min_value=1,max_value=10, key="num_chain_neighbors")

conditioner = conditioners.SymmetryConditioner(G="C_3", num_chain_neighbors=num_chain_neighbors)


st.write("slide to choose your chain_lengths:")
chain_lengths_2 = st.slider(label="chain_lengths_2",min_value=10,max_value=500,key="chain_lengths_2")

st.write("slide to choose your langevin_factor:")
langevin_factor = st.slider(label="langevin_factor",min_value=8,max_value=32, key="langevin_factor")

st.write("slide to choose your inverse_temperature:")
inverse_temperature = st.slider(label="inverse_temperature",min_value=8,max_value=64,key="inverse_temperature")

load_message(Protein('sample-C3.pdb'))
if st.button("Run Code with Button",key="Conditional"):
    chroma = Chroma()
    protein = chroma.sample(
        chain_lengths=[chain_lengths_2],
        conditioner=conditioner,
        langevin_factor=langevin_factor,
        inverse_temperature=inverse_temperature,
        sde_func="langevin",
        potts_symmetry_order=conditioner.potts_symmetry_order)

    protein.to("sample-C3.pdb")
    load_message(protein)

# 显示 Protein 的结构
with open("sample-C3.pdb", "r") as file:
    pdb_content = file.read()
obj = makeobj(pdb_content)

# 使用 stmol 展示蛋白质结构
st.header("Protein Structure:")
showmol(obj, width=1800)


# Display title
st.title("Design display")
load_message(Protein('1GFP'))
if st.button("Run Code with Button",key="Design"):
    chroma = Chroma()

    protein = Protein('1GFP',device='cuda')
    protein = chroma.design(protein)
    protein.to("1GFP-redesign.cif")



# 显示 Protein 的结构
with open("1GFP-redesign.cif", "r") as file:
    pdb_content = file.read()
obj = makeobj(pdb_content)

# 使用 stmol 展示蛋白质结构
st.header("Protein Structure:")
showmol(obj, width=1800)


# Display title
st.title("Redesign display")
load_message(Protein('my_favorite_protein_redesign.cif'))
if st.button("Run Code with Button",key="Redesign"):
    chroma = Chroma()

    protein = Protein('sample.cif',device='cuda') # PDB is fine too
    protein = chroma.design(protein, design_selection="resid 20-50 around 5.0") #  5 angstrom bubble around indices 20-50

    protein.to("my_favorite_protein_redesign.cif")



# 显示 Protein 的结构
with open("my_favorite_protein_redesign.cif", "r") as file:
    pdb_content = file.read()
obj = makeobj(pdb_content)

# 使用 stmol 展示蛋白质结构
st.header("Protein Structure:")
showmol(obj, width=1800)