main.py

import os
import subprocess
import traceback

import pandas as pd
import streamlit as st
from streamlit_mermaid import st_mermaid

from src.code_executor import execute_code
from src.data_handler import get_data_summary, load_data, validate_data
from src.llm_interface import (
    gen_viz_code,
    generate_explanation,
    generate_mermaid_flowchart,
    get_visualization_prompt,
    initialize_llm_client,
    strip_markdown_codeblock,
)

st.set_page_config(
    page_title="Bioinformatics Analysis Studio",
    layout="wide",
    initial_sidebar_state="expanded",
)

# Initialize session state
if "data" not in st.session_state:
    st.session_state.data = []
    # Add default data loading here
    default_file_path = (
        "brca_umich_proteomics_imputed.csv"  # Default data file in project root
    )
    if os.path.exists(default_file_path):
        try:
            data = pd.read_csv(default_file_path)  # or pd.read_excel() for xlsx
            data_description = get_data_summary(data)
            if validate_data(data):
                data_entry = {
                    "file_path": os.path.abspath(default_file_path),
                    "data_description": data_description,
                }
                st.session_state.data.append(data_entry)
                st.toast(
                    f"✅ Loaded default file: {os.path.basename(default_file_path)}"
                )
        except Exception as e:
            st.error(f"Error loading default file: {str(e)}")
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "current_code" not in st.session_state:
    st.session_state.current_code = None
if "llm_provider" not in st.session_state:
    st.session_state.llm_provider = "anthropic"
if "llm_client" not in st.session_state:
    st.session_state.llm_client = None
if "current_output_tab" not in st.session_state:
    st.session_state.current_output_tab = "Generated Prompt"
if "current_error" not in st.session_state:
    st.session_state.current_error = None
if "current_explanation" not in st.session_state:
    st.session_state.current_explanation = None


@st.dialog("Debug State", width="large")
def show_state():
    st.write(st.session_state)


# Define a function to switch to a specific tab
def switch_to_tab(tab_name):
    js_code = f"""
    <script>
        // Select the tab container
        console.log("Switching to tab: {tab_name}")
        var tabContainer = window.parent.document.querySelector('.stTabs');
        
        // Select the tab buttons
        var tabButtons = tabContainer.querySelectorAll('[role="tab"]');
        
        // Find the button for the target tab and click it
        tabButtons.forEach(function(button) {{
            if (button.innerText.trim() === "{tab_name}") {{
                button.click();
            }}
        }});
    </script>
    """
    # Execute the JavaScript code
    st.components.v1.html(js_code, height=0, width=0)


def git_info_component():
    """Helper function to get the current git hash and message."""
    try:
        git_hash = (
            subprocess.check_output(["git", "rev-parse", "--short", "HEAD"])
            .decode("ascii")
            .strip()
        )
        git_msg = (
            subprocess.check_output(["git", "log", "-1", "--pretty=%B"])
            .decode("ascii")
            .strip()
        )
    except Exception as e:
        return None, str(e)

    if git_hash:
        return st.caption(f"Git Hash: {git_hash} - {git_msg}")
    else:
        return st.caption(f"Not in git repository: {git_msg}")


def main():
    st.title("🧬 Bioinformatics Analysis Studio")
    # Display current git hash
    git_info_component()

    # Sidebar for file upload and data info
    with st.sidebar:
        # Initialize LLM client if not already initialized
        if st.session_state.llm_client is None:
            try:
                st.session_state.llm_client = initialize_llm_client(
                    st.session_state.llm_provider
                )
            except Exception as e:
                st.error(
                    f"Failed to initialize {st.session_state.llm_provider}: {str(e)}"
                )
                st.session_state.llm_client = None

        # LLM Provider selection with status indicator
        st.header("LLM Model")

        status_indicator = "✅" if st.session_state.llm_client is not None else "❌"
        selected_provider = st.selectbox(
            f"{status_indicator} Select Provider",
            options=["openai", "anthropic"],
            index=0 if st.session_state.llm_provider == "openai" else 1,
        )

        if selected_provider != st.session_state.llm_provider:
            st.session_state.llm_provider = selected_provider
            try:
                st.session_state.llm_client = initialize_llm_client(selected_provider)
            except Exception as e:
                st.error(f"Failed to initialize {selected_provider}: {str(e)}")
                st.session_state.llm_client = None

        st.header("Data Files")

        uploaded_files = st.file_uploader(
            "Choose data files", type=["csv", "xlsx"], accept_multiple_files=True
        )

        if uploaded_files:
            try:
                st.info("Loading data files...")
                st.session_state.data = []
                for file in uploaded_files:
                    data = load_data(file)
                    data_description = get_data_summary(data)
                    if validate_data(data):
                        data_entry = {
                            "file_path": os.path.abspath(file.name),
                            "data_description": data_description,
                        }
                        st.session_state.data.append(data_entry)
                st.toast(f"✅ Loaded {len(st.session_state.data)} data files")

            except Exception as e:
                st.error(f"Error loading files: {str(e)}")

        if st.session_state.data:
            # Display data files
            st.header("Available Data Files")
            for data_entry in st.session_state.data:
                with st.expander(f"📄 {data_entry['file_path']}"):
                    summary = data_entry["data_description"]
                    st.write(summary)

        if st.button("Session state (debugging)"):
            show_state()

    if len(st.session_state.data) > 0:
        # Main content area with two columns
        col1, col2 = st.columns([4, 5])

        with col1:
            st.header("📝 Instructions")
            user_input = st.text_area(
                "Enter your visualization request:",
                value="""열: 샘플  
행: 단백질

clustermap의 계층적 클러스터링을 사용하여 클러스터 정의

각 클러스터별 상위 100개의 표준 편차가 가장 높은 유전자 선택

변동성이 큰 유전자 상위 목록의 clustermap 표시; 샘플 클러스터를 col_colors 인수로 표시, 딕셔너리를 전달할 수 있음

ticklabels 자동으로 표시

각 클러스터에 고유한 유전자마다 클러스터와 동일한 색깔로 행에 표시 

행/열 둘다 색갈 표시되는지 확인

각 클러스터에 해당하는 유전자 print""",
                height=300,
                key="viz_input",
                on_change=generate_visualization,
            )

            if st.button("Go!"):
                if user_input:
                    generate_visualization()
                else:
                    st.warning("Please enter a visualization request first")

            st.checkbox(
                "Reuse previous code",
                value=False,
                # disabled=not st.session_state.current_code,
                key="reuse_code",
            )

            st.checkbox(
                "Try to autofix error",
                value=False,
                # disabled=not st.session_state.current_error,
                key="autofix_error",
            )

            st.slider(
                "Temperature",
                min_value=0.0,
                max_value=1.0,
                value=0.3,
                step=0.1,
                format="%.1f",
                key="temperature",
            )

        with col2:
            st.header("📊 Analysis Results")
            status = st.status(
                "Waiting for instructions", expanded=False, state="error"
            )

            if "current_prompt" in st.session_state:
                tab1, tab2, tab3 = st.tabs(["Code", "Visualization", "Explanation"])

                with tab1:
                    if st.session_state.current_code:
                        status.write("Generated code")
                        st.code(
                            st.session_state.current_code,
                            language="python",
                        )
                        switch_to_tab("Code")

                with tab2:
                    if st.session_state.current_code:
                        try:
                            execute_code(
                                st.session_state.current_code,
                                st.session_state.data,
                                status,
                            )
                            # Generate explanation after successful visualization
                            st.session_state.current_error = None
                        except Exception as e:
                            st.session_state.current_error = (
                                str(e) + "\n" + str(traceback.format_exc())
                            )
                            st.error(st.session_state.current_error)
                        switch_to_tab("Visualization")

                with tab3:
                    if st.button("Generate Explanation"):
                        try:
                            explanation = generate_explanation(
                                st.session_state.current_prompt,
                                st.session_state.current_code,
                                st.session_state.llm_client,
                                st.session_state.temperature,
                                status,
                            )
                            st.markdown(explanation)

                            # Add Mermaid diagram
                            st.divider()
                            st.subheader("Analysis Workflow")
                            mermaid_code = generate_mermaid_flowchart(
                                st.session_state.current_prompt,
                                st.session_state.current_code,
                                st.session_state.llm_client,
                                st.session_state.temperature,
                                status,
                                explanation,
                            )
                            st_mermaid(
                                mermaid_code, pan=False, zoom=False, show_controls=False
                            )

                        except Exception as e:
                            st.error(f"Error generating explanation: {str(e)}")


def generate_visualization():
    """Helper function to generate visualization from prompt"""
    if not st.session_state.llm_client:
        st.error("Please ensure LLM provider is properly configured")
        return

    user_prompt = st.session_state.viz_input
    current_code = (
        st.session_state.current_code if st.session_state.reuse_code else None
    )
    current_error = st.session_state.current_error
    autofix_error = st.session_state.autofix_error
    if autofix_error:
        user_prompt += str(current_error)
    temperature = st.session_state.temperature

    try:
        st.session_state.chat_history.append(("user", user_prompt))
        prompt = get_visualization_prompt(
            st.session_state.data, user_prompt, current_code
        )
        # Store the prompt in session state
        st.session_state.current_prompt = prompt

        code = strip_markdown_codeblock(
            gen_viz_code(
                prompt,
                st.session_state.llm_client,
                temperature,
            )
        )
        st.session_state.current_code = code
        st.session_state.chat_history.append(
            ("assistant", "Generated visualization code:")
        )
    except Exception as e:
        if "Rate limit exceeded" in str(e):
            st.error(
                "⚠️ Rate limit exceeded. Please try a simpler request or wait a moment before trying again."
            )
        else:
            st.error(f"Error generating visualization: {str(e)}")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        st.error(f"Application error: {str(e)}")
        st.exception(e)