import streamlit as st
from utils import validate_sequence, predict
from model import models
import pandas as pd

def main():
    st.set_page_config(layout="wide")  # Keep the wide layout for overall flexibility
    st.title("AA Property Inference Demo", anchor=None)

    # Apply monospace font to the entire app
    st.markdown("""
        <style>
        .reportview-container {
            font-family: 'Courier New', monospace;
        }
        .css-1d391kg {
            padding: 0.25rem 1rem;  # Adjust padding if necessary for better spacing
        }
        </style>
        """, unsafe_allow_html=True)
    
    # Create columns for input controls
    col1, col2 = st.columns([1, 3])  # Adjust column width ratios as needed
    
    with col1:
        sequence = st.text_input("Enter your amino acid sequence:")
        uploaded_file = st.file_uploader("Or upload a CSV file with amino acid sequences", type="csv")

    if st.button("Analyze Sequence"):
        sequences = [sequence] if sequence else []
        if uploaded_file:
            df = pd.read_csv(uploaded_file)
            sequences.extend(df['sequence'].tolist())

        results = []
        for seq in sequences:
            if validate_sequence(seq):
                model_results = {}
                for model_name, model in models.items():
                    prediction, confidence = predict(model, seq)
                    model_results[f"{model_name}_prediction"] = prediction
                    model_results[f"{model_name}_confidence"] = round(confidence, 3)
                results.append({"Sequence": seq, **model_results})
            else:
                st.error(f"Invalid sequence: {seq}")

        if results:
            st.write("### Results")
            results_df = pd.DataFrame(results)
            # Use full width for DataFrame display only
            st.dataframe(results_df.style.format(precision=3), width=None, height=None)

if __name__ == "__main__":
    main()