from agi_web import python_ui as st
import pandas as pd
import altair as alt
import text_core
import os

# Streamlit 1.64: no set_page_config, use width='stretch' where applicable
st.title("Text atlas")
st.caption("Source: INRIA scikit-learn MOOC, dimred_text.ipynb (CC BY 4.0). Corpus: Wikinews (CC BY 2.5).")

# Initialize session state for results
if "analysis_results" not in st.session_state:
    st.session_state.analysis_results = None
    st.session_state.analysis_params = None

# Bordered form with sliders
with st.form("analysis_form"):
    col1, col2 = st.columns(2)
    with col1:
        n_clusters = st.slider("Number of clusters", min_value=2, max_value=8, value=5, step=1, key="n_clusters_slider")
    with col2:
        min_df = st.slider("Min document frequency", min_value=2, max_value=12, value=5, step=1, key="min_df_slider")
        max_df = st.slider("Max document frequency", min_value=0.60, max_value=0.95, value=0.80, step=0.01, key="max_df_slider")
    
    seed = st.number_input("Random seed", min_value=0, max_value=10000, value=42, step=1, key="seed_input")
    
    submitted = st.form_submit_button("Run analysis")

if submitted:
    try:
        with st.spinner("Running analysis..."):
            # Use cached analysis
            results = st.cache_data(text_core.analyze, max_entries=8)(min_df=min_df, max_df=max_df, n_clusters=n_clusters, seed=int(seed))
            st.session_state.analysis_results = results
            st.session_state.analysis_params = {
                "min_df": min_df,
                "max_df": max_df,
                "n_clusters": n_clusters,
                "seed": int(seed)
            }
    except Exception as e:
        st.error(f"Analysis failed: {e}")

# Display Results
if st.session_state.analysis_results is not None:
    results = st.session_state.analysis_results
    params = st.session_state.analysis_params
    
    st.subheader(f"Results for parameters: {params}")
    
    # Metrics
    col1, col2, col3, col4 = st.columns(4)
    col1.metric("Articles", len(results['coordinates']))
    col2.metric("Vocabulary Size", len(results['vocabulary']))
    col3.metric("Silhouette Score", f"{results['silhouette']:.3f}")
    col4.metric("2D Variance", f"{results['displayed_variance']:.3f}")
    
    st.markdown("""
    **Interpretation:**
    - Nearby points reflect similar vocabulary usage.
    - Clusters are not validated topics; they are unsupervised groupings.
    - 2D projection drops significant detail from the high-dimensional space.
    - Silhouette score is an internal diagnostic of cluster separation.
    """)
    
    # Load corpus for display
    corpus = text_core.load_corpus()
    
    # Prepare DataFrame for plotting
    df_plot = pd.DataFrame({
        "PC1": results["coordinates"][:, 0],
        "PC2": results["coordinates"][:, 1],
        "Cluster": results["labels"],
        "Category": corpus["category"],
        "Text": corpus["text"].str[:100] + "..." # Short preview for tooltip
    })
    
    # Color selector
    color_option = st.radio("Color by", ["Cluster", "Category"], horizontal=True, key="color_selector")
    
    if color_option == "Cluster":
        color_field = "Cluster"
    else:
        color_field = "Category"
        
    # Scatter Plot
    chart = alt.Chart(df_plot).mark_circle(size=60, opacity=0.7).encode(
        x=alt.X("PC1:Q", title="PC1"),
        y=alt.Y("PC2:Q", title="PC2"),
        color=alt.Color(color_field, legend=alt.Legend(title=color_field)),
        tooltip=[
            alt.Tooltip("Text:N", title="Preview"),
            alt.Tooltip("Cluster:N", title="Cluster"),
            alt.Tooltip("Category:N", title="Category")
        ]
    ).properties(
        width="container",
        height=400
    )
    
    st.altair_chart(chart, use_container_width=True)
    
    # Top Terms Table
    st.subheader("Cluster Top Terms")
    top_terms_data = []
    for cluster_id, terms in results["top_terms"].items():
        top_terms_data.append({
            "Cluster": cluster_id,
            "Top Terms": ", ".join(terms)
        })
    st.dataframe(pd.DataFrame(top_terms_data))
    
    # Article Selector
    st.subheader("Inspect Articles")
    article_indices = list(range(len(corpus)))
    selected_idx = st.selectbox("Select article index", article_indices, key="article_selector")
    
    if selected_idx is not None:
        st.text_area("Article Text", corpus.iloc[selected_idx]["text"], height=200, disabled=True)