import streamlit as st
from transformers import TFAutoModelForSequenceClassification
from transformers import AutoTokenizer
import pandas as pd

# title
st.title('Raven AI')

# text input with label
sequence = st.text_input('Enter Amino Acid Sequence')

model_type = st.radio(
    "Choose Linear Epitope Classifier",
    ('Linear T-Cells (MHC Class I Restriction)', 'Linear T-Cells (MHC Class II Restriction)', 'Linear B-Cell'))

# windows length slider
length = st.slider('Window Length', 1, 20, 10)
threshold = st.slider('Probability Threshold', 0.0, 1.0, 0.5)

model_checkpoint = "facebook/esm2_t6_8M_UR50D"

tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

if model_type == 'Linear T-Cells (MHC Class I Restriction)':
    model = TFAutoModelForSequenceClassification.from_pretrained('classifier')
elif model_type == 'Linear T-Cells (MHC Class II Restriction)':
    model = TFAutoModelForSequenceClassification.from_pretrained('classifier2')
elif model_type == 'Linear B-Cell':
    model = TFAutoModelForSequenceClassification.from_pretrained('bcell')
# submit button
if st.button('Submit'):
    # run model
    locations = []
    for i in range(len(sequence) - length):
        peptide_name = sequence[i:i+length]
        peptide = tokenizer(peptide_name, return_tensors="tf")
        output = model(peptide)
        locations.append([peptide_name, output.logits.numpy()[0][0]])
    
    locations = pd.DataFrame(locations, columns = ['Peptide', 'Probability'])

    # display table with sequence and probability as the headers
    def color_survived(x: float): # x between 0 and 1
        # red to green scale based on x
        # 0 -> red
        # 0.5 -> clear
        # 1 -> green

        # red
        if x < threshold:
            r = 179
            g = 40
            b = 2
        # green
        else:
            r = 18
            g = 150
            b = 6

        return f'background-color: rgb({r}, {g}, {b})'

    st.table(locations.style.applymap(color_survived, subset=['Probability']))