web_app.py

import streamlit as st
from streamlit_option_menu import option_menu
import os
import sys
import base64
from pathlib import Path

# Import the page modules.
from attribution_analysis.attribution_analysis_page import show_attribution_analysis
from function_vectors.function_vectors_page import show_function_vectors_page
from circuit_analysis.circuit_trace_page import show_circuit_trace_page
from utilities.welcome_page import show_welcome_page
from utilities.utils import set_seed
from utilities.localization import initialize_localization, tr, language_selector
from utilities.feedback_survey import get_next_participant_id

# Import functions with persisted cache to clear them when needed.
from attribution_analysis.attribution_analysis_page import (
    get_influential_docs, 
    _cached_explain_heatmap as attr_explain_heatmap, 
    generate_all_attribution_analyses
)
from circuit_analysis.circuit_trace_page import explain_circuit_visualization
from function_vectors.function_vectors_page import (
    _perform_analysis as fv_perform_analysis, 
    _explain_with_llm as fv_explain_llm
)


# Set the page configuration.
st.set_page_config(
    page_title="LLM Analysis Suite",
    page_icon="🧠",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Set TOKENIZERS_PARALLELISM to false to avoid warnings.
os.environ["TOKENIZERS_PARALLELISM"] = "false"

# Suppress a harmless error on macOS.
os.environ["OBJC_DISABLE_INITIALIZE_FORK_SAFETY"] = "YES"


# Initialize theme in session state
if 'theme_mode' not in st.session_state:
    st.session_state.theme_mode = 'dark'

_is_light = st.session_state.theme_mode == 'light'

# ── Shared CSS (both themes) ──────────────────────────────────────────────
st.markdown("""
<style>
    [data-testid="stToolbar"] { visibility: hidden; }
    .main-header {
        font-size: 3rem;
        text-align: center;
        margin-bottom: 2rem;
    }
    .stButton > button {
        background-color: #2f3f70;
        color: #f5f7fb;
        border-radius: 20px;
        border: none;
        padding: 0.5rem 2rem;
        font-weight: bold;
        box-shadow: 0 10px 20px rgba(47, 63, 112, 0.25);
    }
    .stButton > button:hover {
        background-color: #3a4c86;
        color: #ffffff;
    }
    .stTextArea > div > div > textarea { border-radius: 10px; }
    .attribution-info {
        background-color: rgba(47, 63, 112, 0.82);
        color: #f5f7fb;
        padding: 1rem;
        border-radius: 10px;
        margin: 1rem 0;
        border-left: 4px solid #dcae36;
    }
</style>
""", unsafe_allow_html=True)

# ── Light-mode CSS overrides ──────────────────────────────────────────────
if _is_light:
    st.markdown("""
<style>
    /* ── Streamlit chrome ────────────────────────────────────── */
    [data-testid="stAppViewContainer"],
    [data-testid="stApp"],
    .main .block-container {
        background-color: #ffffff !important;
        color: #1b1b2f !important;
    }
    [data-testid="stSidebar"],
    [data-testid="stSidebar"] > div:first-child {
        background-color: #f0f2f6 !important;
        color: #1b1b2f !important;
    }
    [data-testid="stHeader"] {
        background-color: #ffffff !important;
    }

    /* ── Global text color reset ─────────────────────────────── */
    .main-header { color: #1b1b2f !important; }

    h1, h2, h3, h4, h5, h6 { color: #1b1b2f !important; }

    [data-testid="stMarkdownContainer"],
    [data-testid="stMarkdownContainer"] p,
    [data-testid="stMarkdownContainer"] li,
    [data-testid="stMarkdownContainer"] strong,
    [data-testid="stMarkdownContainer"] em,
    [data-testid="stMetricValue"],
    [data-testid="stMetricLabel"],
    [data-testid="stMetricDelta"],
    label, .stSelectbox label, .stRadio label {
        color: #1b1b2f !important;
    }

    /* ── Inputs & widgets ────────────────────────────────────── */
    .stTextInput > div > div > input,
    .stTextArea > div > div > textarea,
    [data-baseweb="select"],
    [data-baseweb="select"] > div,
    .stSelectbox > div > div,
    .stMultiSelect > div > div {
        background-color: #ffffff !important;
        color: #1b1b2f !important;
        border-color: #c8cdd5 !important;
    }
    [data-baseweb="select"] span,
    [data-baseweb="select"] div {
        color: #1b1b2f !important;
    }
    /* dropdown lists */
    [data-baseweb="popover"],
    [data-baseweb="menu"],
    ul[role="listbox"],
    ul[role="listbox"] li {
        background-color: #ffffff !important;
        color: #1b1b2f !important;
    }
    ul[role="listbox"] li:hover {
        background-color: #e8eaee !important;
    }

    /* ── Tabs ────────────────────────────────────────────────── */
    .stTabs [data-baseweb="tab-list"] {
        background-color: transparent !important;
    }
    .stTabs [data-baseweb="tab"] {
        color: #1b1b2f !important;
    }

    /* ── Expanders ───────────────────────────────────────────── */
    [data-testid="stExpander"],
    [data-testid="stExpander"] summary,
    [data-testid="stExpander"] > div {
        background-color: #f5f6fa !important;
        color: #1b1b2f !important;
    }

    /* ── option_menu sidebar override ────────────────────────── */
    .nav-link { color: #1b1b2f !important; }
    .nav-link.active, .nav-link:hover {
        background-color: #2f3f70 !important;
        color: #ffffff !important;
    }

    /* ── Plotly charts ────────────────────────────────────────── */
    .js-plotly-plot .main-svg {
        background-color: transparent !important;
    }

    /* ── Info & attribution boxes ─────────────────────────────── */
    .attribution-info {
        background-color: rgba(47, 63, 112, 0.08) !important;
        color: #1b1b2f !important;
    }

    /* ── Chat messages ────────────────────────────────────────── */
    [data-testid="stChatMessage"] {
        background-color: #f5f6fa !important;
        color: #1b1b2f !important;
    }

    /* ── Buttons (light override) ────────────────────────────── */
    .stButton > button {
        background-color: #e8ecf6 !important;
        color: #2f3f70 !important;
        border: 1px solid #c8cdd5 !important;
        box-shadow: 0 2px 6px rgba(0,0,0,0.06) !important;
    }
    .stButton > button:hover {
        background-color: #d8dff0 !important;
        color: #2f3f70 !important;
    }

    /* ── Download buttons ────────────────────────────────────── */
    .stDownloadButton > button {
        background-color: #e8ecf6 !important;
        color: #2f3f70 !important;
        border: 1px solid #c8cdd5 !important;
        box-shadow: 0 2px 6px rgba(0,0,0,0.06) !important;
    }
    .stDownloadButton > button:hover {
        background-color: #d8dff0 !important;
        color: #2f3f70 !important;
    }

    /* ── Chat input bar & bottom container ─────────────────── */
    [data-testid="stBottom"],
    [data-testid="stBottom"] > div,
    [data-testid="stChatFloatingInputContainer"],
    .stChatFloatingInputContainer {
        background-color: #ffffff !important;
    }
    [data-testid="stChatInput"],
    [data-testid="stChatInput"] textarea,
    [data-testid="stChatInput"] > div {
        background-color: #f5f6fa !important;
        color: #1b1b2f !important;
        border-color: #c8cdd5 !important;
    }

    /* ── Toggle switches ──────────────────────────────────────── */
    [data-testid="stToggle"] label,
    [data-testid="stToggle"] span {
        color: #1b1b2f !important;
    }
    [data-testid="stToggle"] [role="checkbox"] {
        border: 2px solid #2f3f70 !important;
    }

    /* ── Sidebar collapse button ──────────────────────────────── */
    [data-testid="stSidebar"] button[kind="header"],
    [data-testid="collapsedControl"] button,
    button[data-testid="stBaseButton-headerNoPadding"] {
        color: #1b1b2f !important;
        background-color: #e8ecf6 !important;
        border: 1px solid #c8cdd5 !important;
        border-radius: 8px !important;
    }

    /* ── Slider thumb & track ─────────────────────────────────── */
    [data-testid="stSlider"] label {
        color: #1b1b2f !important;
    }

    /* ── Checkbox ─────────────────────────────────────────────── */
    .stCheckbox label span {
        color: #1b1b2f !important;
    }

    /* ── Option menu (sidebar nav) ────────────────────────────── */
    [data-testid="stSidebar"] iframe + div,
    [data-testid="stSidebar"] .nav {
        background-color: #f0f2f6 !important;
    }
</style>
""", unsafe_allow_html=True)

def main():
    # Main function to run the app.
    set_seed()
    initialize_localization()

    # The language selector is now on the welcome page.
    # We don't need it in the sidebar of the main app.

    # Check if the user has submitted the welcome form.
    if "user_info" not in st.session_state or not st.session_state.user_info.get("form_submitted"):
        show_welcome_page()
    else:
        # If the form is submitted, show the main application.
        
        # Assign a participant ID if one doesn't exist.
        if 'participant_id' not in st.session_state:
            st.session_state.participant_id = get_next_participant_id()
            
        # Initialize session state for feedback forms.
        if 'attr_feedback_submitted' not in st.session_state:
            st.session_state.attr_feedback_submitted = False
        if 'fv_feedback_submitted' not in st.session_state:
            st.session_state.fv_feedback_submitted = False

        logo_path = Path(__file__).parent / "LOGO" / "Logo.png"
        if logo_path.exists():
            with open(logo_path, "rb") as logo_file:
                logo_base64 = base64.b64encode(logo_file.read()).decode("utf-8")
            st.markdown(
                f"""
                <div style="text-align: center; margin-bottom: 2rem;">
                    <img src="data:image/png;base64,{logo_base64}" alt="{tr('llm_analysis_suite')}" style="max-width: 320px; width: 60%; min-width: 200px;" />
                </div>
                """,
                unsafe_allow_html=True
            )
        else:
            st.markdown(f"<h1 class='main-header'>{tr('llm_analysis_suite')}</h1>", unsafe_allow_html=True)
        
        with st.sidebar:
            # Theme toggle switch at the top of sidebar
            light_on = st.toggle(
                tr('light_mode'),
                value=(st.session_state.theme_mode == 'light'),
                key="theme_toggle"
            )
            new_mode = 'light' if light_on else 'dark'
            if new_mode != st.session_state.theme_mode:
                st.session_state.theme_mode = new_mode
                st.rerun()
            
            _menu_styles = {
                "container": {"background-color": "#f0f2f6"},
                "icon": {"color": "#4a5568"},
                "nav-link": {"color": "#1b1b2f", "--hover-color": "#e2e4ea"},
                "nav-link-selected": {"background-color": "#2f3f70", "color": "#ffffff"},
                "menu-title": {"color": "#1b1b2f"},
                "menu-icon": {"color": "#4a5568"},
            } if _is_light else {}

            selected_page = option_menu(
                menu_title=tr('main_menu'),
                options=[tr('attribution_analysis'), tr('function_vectors'), tr('circuit_tracing')],
                icons=['search', 'cpu', 'diagram-3'],
                menu_icon='cast',
                default_index=0,
                styles=_menu_styles
            )

            

        if selected_page == tr('attribution_analysis'):
            show_attribution_analysis()
        elif selected_page == tr('function_vectors'):
            show_function_vectors_page()
        elif selected_page == tr('circuit_tracing'):
            show_circuit_trace_page()

if __name__ == "__main__":
    main()