import os
import streamlit as st
from rapidfuzz import process

# with st.spinner("Initializing the environment... This may take up to 10 minutes at the start of each session."):
#     # Create a temporary placeholder for the message
#     loading_placeholder = st.empty()

#     # Show the info message only while the spinner is active
#     loading_placeholder.info("""
#     **Note:** This initialization is required at the start of each session.  
#     Once the app is ready, you can run multiple predictions without re-initializing by clicking the **Reset** button in the sidebar.
#     """)

#     # Run setup script if not already executed
#     if not os.path.exists(".setup_done"):
#         start_time = time.time()
#         os.system("bash setup.sh")
#         end_time = time.time()
#         print(f"Environment prepared in {end_time - start_time:.2f} seconds")
#         with open(".setup_done", "w") as f:
#             f.write("done")

# # ❌ Remove the info message after initialization is complete
# loading_placeholder.empty()


from run_prothgt_app import *

def convert_df(df):
   return df.to_csv(index=False).encode('utf-8')

# Initialize session state variables
if 'predictions_df' not in st.session_state:
    st.session_state.predictions_df = None
if 'submitted' not in st.session_state:
    st.session_state.submitted = False

with st.expander("🚀 Upcoming Features"):
    st.info("""
    We are actively working on enhancing ProtHGT application with new capabilities:
    
    - **Real-time data retrieval for new proteins**: Currently, ProtHGT can only generate predictions for proteins that already exist in our knowledge graph. We are developing a new feature that will allow users to **predict functions for entirely new proteins starting from their sequences**. This will work by **retrieving relevant relationship data in real time from external source databases** (e.g., UniProt, STRING, and other biological repositories). The system will dynamically construct a knowledge graph for the query protein, incorporating its interactions, domains, pathways, and other biological associations before running function prediction. This approach will enable ProtHGT to analyze newly discovered or less-studied proteins even if they are not pre-annotated in our dataset.  
    - **Expanded embedding options**: Currently, this application represents proteins using **TAPE embeddings**, which serve as the initial numerical representations of protein sequences before being processed in the heterogeneous graph model. We are working on integrating **ProtT5** and **ESM-2** as alternative initial embeddings, allowing users to choose different sequence representations that may enhance performance for specific tasks. A detailed comparison of how these embeddings influence function prediction accuracy will be included in our upcoming publication.
    - **Knowledge graph visualization for interpretability**: To improve model explainability, we are developing an interactive **knowledge graph visualization** feature. This will allow users to explore the biological relationships that contributed to ProtHGT’s predictions for a given protein. Users will be able to inspect **protein interactions, GO annotations, domains, pathways, and other key connections** in a structured graphical format, making it easier to interpret and validate predictions.

    Stay tuned for updates and future publications!
    """)

with st.sidebar:
    st.markdown("""
        <style>
        .title {
            font-size: 35px;
            font-weight: bold;
            color: #424242;
            margin-bottom: 0px;
        }
        .subtitle {
            font-size: 20px;
            color: #424242;
            margin-bottom: 20px;
            line-height: 1.5;
        }
        .badges {
            margin-top: 10px;
            margin-bottom: 20px;
        }
        </style>
        
        <div class="title">ProtHGT</div>
        <div class="subtitle">Heterogeneous Graph Transformers for Automated Protein Function Prediction Using Knowledge Graphs and Language Models</div>
        <div class="badges">
            <a href="https://github.com/HUBioDataLab/ProtHGT">
                <img src="https://img.shields.io/badge/GitHub-black?logo=github" alt="github-repository">
            </a>
        </div>
    """, unsafe_allow_html=True)

    available_proteins = get_available_proteins() 

    selected_proteins = []
    
    # Add protein selection methods
    selection_method = st.radio(
        "Choose input method:",
        ["Search proteins", "Upload protein ID file"]
    )

    if selection_method == "Search proteins":
        # User enters search term
        search_query = st.text_input("Start typing a protein ID (at least 3 characters)", "")

        # Apply fuzzy search only if query length is >= 3
        filtered_proteins = []
        if len(search_query) >= 3:
            filtered_proteins = [match[0] for match in process.extract(search_query, available_proteins, limit=50)]  # Show top 50 matches

        # Multi-select for filtered results
        selected_proteins = st.multiselect(
            "Select proteins from search results",
            options=filtered_proteins,
            placeholder="Start typing to search...",
            max_selections=1000
        )

        if selected_proteins:
            st.write(f"Selected {len(selected_proteins)} proteins")
                
    else:
        uploaded_file = st.file_uploader(
            "Upload a text file with UniProt IDs (one per line, max 1000)*",
            type=['txt']
        )        

        if uploaded_file:
            protein_list = [line.strip() for line in uploaded_file.read().decode('utf-8').splitlines()]

            # Remove empty lines and duplicates
            protein_list = list(filter(None, protein_list))
            protein_list = list(dict.fromkeys(protein_list))
            
            # Check for proteins not in available_proteins
            proteins_not_found = [p for p in protein_list if p not in available_proteins]
            # Filter to keep only available proteins
            protein_list = [p for p in protein_list if p in available_proteins]

            if len(protein_list) > 1000:
                st.error("Please upload a file with maximum 1000 protein IDs.")
                selected_proteins = []
            else:
                selected_proteins = protein_list
                st.write(f"Loaded {len(selected_proteins)} proteins")

                if proteins_not_found:
                    st.warning(f"""
                    The following proteins were not found in our input knowledge graph and have been discarded: 
                    """)
                    with st.expander("View Discarded Proteins"):
                        # Create scrollable container with fixed height
                        st.markdown(
                            f"""
                            <div style="
                                height: 150px; 
                                overflow-y: scroll;
                                border: 1px solid #ccc;
                                border-radius: 4px;
                                padding: 8px;
                                margin-bottom: 16px;
                                background-color: white;">
                                {'<br>'.join(proteins_not_found)}
                            </div>
                            """, 
                            unsafe_allow_html=True
                        )
                        
                    st.warning(f"""
                    Currently, our system can only generate predictions for proteins that are already included in our knowledge graph. **Real-time retrieval of relationship data from external source databases is not yet supported.**  
                    We are actively working on integrating this capability in future updates. Stay tuned!
                    """)
    if selected_proteins:
        with st.expander("View Selected Proteins"):
            st.write(f"Total proteins selected: {len(selected_proteins)}")
            
            # Create scrollable container with fixed height
            st.markdown(
                f"""
                <div style="
                    height: 150px; 
                    overflow-y: scroll;
                    border: 1px solid #ccc;
                    border-radius: 4px;
                    padding: 8px;
                    background-color: white;">
                    {'<br>'.join(selected_proteins)}
                </div>
                """, 
                unsafe_allow_html=True
            )

            st.markdown("<div style='padding-top: 10px;'></div>", unsafe_allow_html=True)

            # Add download button
            proteins_text = '\n'.join(selected_proteins)
            st.download_button(
                label="Download List",
                data=proteins_text,
                file_name="selected_proteins.txt",
                mime="text/plain",
                key="download_button"
            )

        # Add GO category selection
        go_category_options = {
            'All Categories': None,
            'Molecular Function': 'GO_term_F',
            'Biological Process': 'GO_term_P',
            'Cellular Component': 'GO_term_C'
        }
        selected_go_category = st.selectbox(
            "Select GO Category for predictions",
            options=list(go_category_options.keys()),
            help="Choose which GO category to generate predictions for. Selecting 'All Categories' will generate predictions for all three categories."
        )

    st.warning("⚠️ Due to memory and computational constraints, the maximum number of proteins that can be processed at once is limited to 1000 proteins. For larger datasets, please consider running the model locally using our GitHub repository.")

    if selected_proteins and selected_go_category:
        # Add a button to trigger predictions
        if st.button("Generate Predictions"):
            st.session_state.submitted = True

if st.session_state.submitted:
    with st.spinner("Generating predictions..."):
        # Generate predictions only if not already in session state
        if st.session_state.predictions_df is None:

            # Load model config from JSON file
            import json
            import os

            # Define data directory path
            data_dir = "data"
            models_dir = os.path.join(data_dir, "models")

            # Load model configuration
            model_config_paths = {
                'GO_term_F': os.path.join(models_dir, "prothgt-config-molecular-function.yaml"),
                'GO_term_P': os.path.join(models_dir, "prothgt-config-biological-process.yaml"),
                'GO_term_C': os.path.join(models_dir, "prothgt-config-cellular-component.yaml")
            }

            # Paths for model and data
            model_paths = {
                'GO_term_F': os.path.join(models_dir, "prothgt-model-molecular-function.pt"),
                'GO_term_P': os.path.join(models_dir, "prothgt-model-biological-process.pt"),
                'GO_term_C': os.path.join(models_dir, "prothgt-model-cellular-component.pt")
            }

            # Get the selected GO category
            go_category = go_category_options[selected_go_category]

            # If a specific category is selected, use that model path
            if go_category:
                model_config_paths = [model_config_paths[go_category]]
                model_paths = [model_paths[go_category]]
                go_categories = [go_category]
            else:
                model_config_paths = [model_config_paths[cat] for cat in ['GO_term_F', 'GO_term_P', 'GO_term_C']]
                model_paths = [model_paths[cat] for cat in ['GO_term_F', 'GO_term_P', 'GO_term_C']]
                go_categories = ['GO_term_F', 'GO_term_P', 'GO_term_C']

            # Generate predictions
            predictions_df = generate_prediction_df(
                protein_ids=selected_proteins,
                model_paths=model_paths,
                model_config_paths=model_config_paths,
                go_category=go_categories
            )

            st.session_state.predictions_df = predictions_df
        
        # Display and filter predictions
        st.success("Predictions generated successfully!")
        st.markdown("### Filter and View Predictions")
        
        # Create filters
        st.markdown("### Filter Predictions")
        col1, col2, col3, col4 = st.columns(4)  # Changed to 4 columns
        
        with col1:
            # Protein filter
            selected_protein = st.selectbox(
                "Filter by Protein",
                options=['All'] + sorted(st.session_state.predictions_df['Protein'].unique().tolist())
            )
            
        with col2:
            # GO category filter
            selected_category = st.selectbox(
                "Filter by GO Category",
                options=['All'] + sorted(st.session_state.predictions_df['GO_category'].unique().tolist())
            )

        with col3:
            # GO term filter
            go_term_filter = st.text_input(
                "Filter by GO Term ID",
                placeholder="e.g., GO:0003674",
                help="Enter a GO term ID to filter results"
            ).strip()
            
        with col4:
            # Probability threshold
            min_probability_threshold = st.slider(
                "Minimum Probability",
                min_value=0.0,
                max_value=1.0,
                value=0.5,
                step=0.05
            )

            max_probability_threshold = st.slider(
                "Maximum Probability",
                min_value=0.0,
                max_value=1.0,
                value=1.0,
                step=0.05
            )

        # Filter the dataframe using session state data
        filtered_df = st.session_state.predictions_df.copy()

        if selected_protein != 'All':
            filtered_df = filtered_df[filtered_df['Protein'] == selected_protein]
            
        if selected_category != 'All':
            filtered_df = filtered_df[filtered_df['GO_category'] == selected_category]

        if go_term_filter:
            filtered_df = filtered_df[filtered_df['GO_term'].str.contains(go_term_filter, case=False, na=False)]
            
        filtered_df = filtered_df[(filtered_df['Probability'] >= min_probability_threshold) & 
                                (filtered_df['Probability'] <= max_probability_threshold)]

        # Custom CSS to increase table width and improve layout
        st.markdown("""
            <style>
                .stDataFrame {
                    width: 100%;
                }
                .stDataFrame > div {
                    width: 100%;
                }
                .stDataFrame [data-testid="stDataFrameResizable"] {
                    width: 100%;
                    min-width: 100%;
                }
                .pagination-info {
                    font-size: 14px;
                    color: #666;
                    padding: 10px 0;
                }
                .page-controls {
                    display: flex;
                    align-items: center;
                    justify-content: center;
                    gap: 20px;
                    padding: 10px 0;
                }
            </style>
        """, unsafe_allow_html=True)

        # Add pagination controls
        col1, col2, col3 = st.columns([2, 1, 2])
        with col2:
            rows_per_page = st.selectbox("Rows per page", [50, 100, 200, 500], index=1)

        total_rows = len(filtered_df)
        total_pages = (total_rows + rows_per_page - 1) // rows_per_page

        # Initialize page number in session state
        if "page_number" not in st.session_state:
            st.session_state.page_number = 0

        # Calculate start and end indices for current page
        start_idx = st.session_state.page_number * rows_per_page
        end_idx = min(start_idx + rows_per_page, total_rows)

        # Display the paginated dataframe with increased width
        st.dataframe(
            filtered_df.iloc[start_idx:end_idx],
            hide_index=True,
            use_container_width=True,  # This makes the table use full width
            column_config={
                "Probability": st.column_config.ProgressColumn(
                    "Probability",
                    format="%.2f",
                    min_value=0,
                    max_value=1,
                ),
                "Protein": st.column_config.TextColumn(
                    "Protein",
                    help="UniProt ID",
                ),
                "GO_category": st.column_config.TextColumn(
                    "GO Category",
                    help="Gene Ontology Category",
                ),
                "GO_term": st.column_config.TextColumn(
                    "GO Term",
                    help="Gene Ontology Term ID",
                ),
            }
        )

        # Pagination controls with better layout
        col1, col2, col3 = st.columns([1, 3, 1])
        with col1:
            if st.button("⬅️ Previous", disabled=st.session_state.page_number == 0):
                st.session_state.page_number -= 1
                st.rerun()

        with col2:
            st.markdown(f"""
                <div class="pagination-info" style="text-align: center">
                    Page {st.session_state.page_number + 1} of {total_pages}<br>
                    Showing rows {start_idx + 1} to {end_idx} of {total_rows}
                </div>
            """, unsafe_allow_html=True)

        with col3:
            if st.button("Next ➡️", disabled=st.session_state.page_number >= total_pages - 1):
                st.session_state.page_number += 1
                st.rerun()


        # Download filtered results
        st.download_button(
            label="Download Filtered Results",
            data=convert_df(filtered_df),
            file_name="filtered_predictions.csv",
            mime="text/csv",
            key="download_filtered_predictions"
        )

# Add a reset button in the sidebar
with st.sidebar:
    if st.session_state.submitted:
        if st.button("Reset"):
            st.session_state.predictions_df = None
            st.session_state.submitted = False
            st.rerun()