import os
# --- ANTI-CRASH ENVIRONMENT VARIABLES ---
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import streamlit as st
import pandas as pd
import warnings
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

from bertopic import BERTopic
from bertopic.representation import MaximalMarginalRelevance, KeyBERTInspired
from sentence_transformers import SentenceTransformer, models
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.decomposition import PCA
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
from umap import UMAP
from hdbscan import HDBSCAN
import gensim.corpora as corpora
from gensim.models.coherencemodel import CoherenceModel

warnings.filterwarnings("ignore")

# ==========================================
# 1. PAGE CONFIGURATION & MAPPINGS
# ==========================================
st.set_page_config(page_title="Topic Modeling Pipeline", layout="wide", initial_sidebar_state="collapsed")

EMBEDDING_MAP = {
    "MiniLM (Fast & Lightweight)": "sentence-transformers/all-MiniLM-L6-v2",
    "MPNet (High Accuracy)": "sentence-transformers/all-mpnet-base-v2",
    "Specter2 (Scientific/Academic)": "allenai/specter2_base"
}

POOLING_MAP = {
    "Mean (Smooth context)": "mean",
    "Max (Sharp keywords)": "max",
    "CLS (Classification)": "cls",
    "Mean-Max (Combined)": "mean-max"
}

# --- CACHE THE NEURAL NETWORK ---
@st.cache_resource
def load_embedder(model_name, pool_strat):
    word_emb = models.Transformer(model_name)
    pool_model = models.Pooling(
        word_emb.get_word_embedding_dimension(),
        pooling_mode_mean_tokens=("mean" in pool_strat),
        pooling_mode_max_tokens=("max" in pool_strat),
        pooling_mode_cls_token=("cls" in pool_strat)
    )
    return SentenceTransformer(modules=[word_emb, pool_model])

# ==========================================
# 2. THE GUIDED UI (MAIN PAGE)
# ==========================================
st.title("🧠 BERTopic Topic Modeling Pipeline")

try:
    st.image("pipeline.png", use_container_width=True)
except FileNotFoundError:
    pass

st.divider()

# --- STEP 0: DATA SETTINGS ---
st.header("📥 Step 0: Input Data & Core Settings")
data_source = st.radio("Choose Data Source:", ["Use Sample ACM Abstract", "Paste Text"], horizontal=True)

sample_abstract = """
Students who registered for the Mapping with Google massive open online course (MOOC) 
were asked several questions during the registration process to identify prior 
experience with eleven skills as well as their goals for registering for the course. 
At the end of the course, we compared students' self reports of goal achievement 
with behavioral click-stream analysis. In addition, we assessed how well prior 
skill in a subject predicts a student's course completion and found no correlation. 
Our research shows that students who completed course activities were more likely 
to earn certificates of completion than peers who did not.
"""

raw_data = st.text_area("Text Data:", value=sample_abstract if data_source == "Use Sample ACM Abstract" else "", height=150)

col_a, col_b = st.columns(2)
with col_a:
    n_themes = st.slider("Target Number of Themes", 2, 20, 3)
with col_b:
    words_per_theme = st.slider("Words to Output per Theme", 3, 10, 5)

# --- THE VERTICAL CONFIGURATION WIZARD ---
st.header("⚙️ Model Configuration")

with st.expander("1️⃣ Semantic Layer (Embeddings & Pooling)", expanded=True):
    ui_embedding = st.selectbox("Embedding Model", list(EMBEDDING_MAP.keys()))
    ui_pooling = st.selectbox("Pooling Strategy", list(POOLING_MAP.keys()))

with st.expander("2️⃣ Geometry Layer (Dimensionality Reduction)", expanded=True):
    ui_algo = st.selectbox("Algorithm", ["UMAP (Complex geometry)", "PCA (Fast/Deterministic)"])
    if "UMAP" in ui_algo:
        ui_metric = st.selectbox("Distance Metric", ["cosine", "euclidean", "manhattan"])
    else:
        ui_metric = "euclidean"
        st.info("PCA inherently uses Variance (Euclidean math), so distance metrics are bypassed.")

with st.expander("3️⃣ Clustering Layer (Grouping)", expanded=True):
    st.markdown("""
    *Clustering mathematically draws boundaries around similar sentences.*
    * **Primary Engine (HDBSCAN):** Runs on datasets $\ge$ 15 sentences. Automatically filters outliers and finds dense semantic clouds. 
      *(Defaults: min_cluster_size=10, cluster_selection_method='eom', metric='euclidean')*
    * **Fallback Engine (K-Means):** Runs on datasets $<$ 15 sentences. Forces all sentences into buckets to prevent math crashes on tiny text samples. 
      *(Defaults: n_clusters = Target Themes, random_state=42)*
    """)

with st.expander("4️⃣ Vocabulary Layer (Vectorization)", expanded=True):
    ngram_range = st.slider("N-Gram Range", 1, 3, (1, 2), help="1=Unigrams, 2=Bigrams (e.g., 'machine learning')")
    # Added the explanation here!
    auto_noise = st.checkbox(
        "Auto-Remove Redundant Noise (max_df)", 
        value=True,
        help="Mathematically deletes words that appear in more than 85% of the documents."
    )
    st.caption("Deletes overly common words (like 'paper' or 'study') that appear everywhere, preventing generic filler from dominating your themes.")

with st.expander("5️⃣ Extraction Layer (Representation)", expanded=True):
    ui_extraction = st.selectbox("Strategy", ["c-TF-IDF (Word frequency)", "KeyBERTInspired (Semantic cosine)", "MMR (Reduce redundancy)"])
    if "MMR" in ui_extraction:
        mmr_diversity = st.slider("MMR Diversity Penalty", 0.0, 1.0, 0.3)
    else:
        mmr_diversity = None

# --- EVALUATION METRICS ---
st.header("📊 Evaluation Metrics")
eval_metrics = st.multiselect(
    "Select KPIs to generate a final report card:",
    ["Topic Diversity", "NPMI Coherence", "Silhouette Score"],
    default=["Topic Diversity", "NPMI Coherence", "Silhouette Score"]
)

st.divider()

# ==========================================
# 3. ENGINE EXECUTION
# ==========================================
if st.button("🚀 Run Topic Modeling Pipeline", type="primary", use_container_width=True):
    
    if not raw_data or len(raw_data) < 20:
        st.error("Please provide more text data!")
        st.stop()

    # --- MATH EXECUTION (Inside Spinner) ---
    with st.spinner("Processing Semantic Pipeline... (Models are cached to prevent crashes)"):
        
        sentences = [s.strip() for s in raw_data.split('.') if len(s.strip()) > 10]
        dataset_size = len(sentences)
        
        academic_noise = ['students', 'course', 'research', 'paper', 'found', 'likely', 'did']
        from sklearn.feature_extraction import text
        stop_w = list(text.ENGLISH_STOP_WORDS.union(academic_noise))
        
        vectorizer_model = CountVectorizer(stop_words=stop_w, ngram_range=ngram_range, max_df=0.85 if auto_noise and dataset_size > 10 else 1.0)
        
        custom_embedder = load_embedder(EMBEDDING_MAP[ui_embedding], POOLING_MAP[ui_pooling])
        embeddings = custom_embedder.encode(sentences)
        
        # Fallback Logic (Step 3 representation in code)
        is_fallback = False
        if dataset_size < 15 or "PCA" in ui_algo:
            safe_n_themes = min(n_themes, dataset_size)
            dim_model = PCA(n_components=2, random_state=42)
            cluster_model = KMeans(n_clusters=safe_n_themes, random_state=42)
            reduce_topics = None
            is_fallback = True
            algo_used = "PCA"
            cluster_algo = "K-Means"
        else:
            dim_model = UMAP(n_neighbors=15, n_components=5, metric=ui_metric, random_state=42)
            clustering_model = HDBSCAN(min_cluster_size=10, metric='euclidean', cluster_selection_method='eom')
            reduce_topics = n_themes
            algo_used = "UMAP"
            cluster_algo = "HDBSCAN"

        # Representation
        if "MMR" in ui_extraction:
            rep_model = MaximalMarginalRelevance(diversity=mmr_diversity, top_n_words=words_per_theme)
        elif "KeyBERT" in ui_extraction:
            rep_model = KeyBERTInspired(top_n_words=words_per_theme)
        else:
            rep_model = None

        topic_model = BERTopic(
            embedding_model=custom_embedder,
            umap_model=dim_model,
            hdbscan_model=cluster_model,
            vectorizer_model=vectorizer_model,
            representation_model=rep_model,
            nr_topics=reduce_topics,
            top_n_words=words_per_theme,
            language="english"
        )
        topics, _ = topic_model.fit_transform(sentences)

    # ==========================================
    # 4. UI DISPLAY & METRICS (Outside Spinner)
    # ==========================================
    st.success("Analysis Complete!")
    
    if is_fallback:
        if safe_n_themes < n_themes:
            st.warning(f"⚠️ **Reduced requested themes from {n_themes} to {safe_n_themes}.**\n\n"
                       f"*The Math Explanation:* BERTopic clusters complete sentences to preserve context. "
                       f"You cannot sort {dataset_size} sentences into {n_themes} buckets without leaving empty buckets, "
                       f"which mathematically breaks the clustering algorithms!")
        else:
            st.info(f"ℹ️ Auto-switched to PCA/K-Means due to small dataset size ({dataset_size} sentences).")
            
    st.markdown("### 🏆 Discovered Themes")
    topic_info = topic_model.get_topic_info()
    all_words = []
    
    cols = st.columns(3)
    col_idx = 0
    for t_id in topic_info['Topic']:
        if t_id == -1: continue
        theme_w = [w[0] for w in topic_model.get_topic(t_id)]
        all_words.append(theme_w)
        with cols[col_idx % 3]:
            st.info(f"**Theme {t_id + 1}**\n\n" + "\n".join([f"🔹 {w}" for w in theme_w]))
        col_idx += 1

    # --- METRICS CALCULATIONS ---
    div_val, npmi_val, sil_val = 0.0, 0.0, 0.0
    
    if len(eval_metrics) > 0:
        st.markdown("### 📊 Key Performance Indicators (KPI)")
        
        with st.spinner("Calculating mathematical metrics... (NPMI requires building a dictionary and takes a moment)"):
            
            for metric in eval_metrics:
                if "Diversity" in metric:
                    if len(all_words) > 0:
                        u_words = set([w for t in all_words for w in t])
                        t_words = sum([len(t) for t in all_words])
                        div_val = len(u_words) / t_words if t_words > 0 else 0
                        st.metric("Topic Diversity (Target: 1.0)", f"{div_val:.2f}")
                    else:
                        st.metric("Topic Diversity", "Skipped")
                        
                elif "NPMI" in metric:
                    try:
                        tokenized = [vectorizer_model.build_analyzer()(s) for s in sentences]
                        dictionary = corpora.Dictionary(tokenized)
                        cm = CoherenceModel(topics=all_words, texts=tokenized, dictionary=dictionary, coherence='c_npmi')
                        temp_npmi = cm.get_coherence()
                        if np.isnan(temp_npmi):
                            st.metric("NPMI Coherence", "N/A (Too few words)")
                        else:
                            npmi_val = float(temp_npmi)
                            st.metric("NPMI Coherence (Target: >0.1)", f"{npmi_val:.2f}")
                    except Exception:
                        st.metric("NPMI Coherence", "Skipped (Data too small)")
                        
                elif "Silhouette" in metric:
                    valid_idx = [i for i, t in enumerate(topics) if t != -1]
                    unique_topics = set([topics[i] for i in valid_idx])
                    
                    if 1 < len(unique_topics) < len(valid_idx):
                        sil_val = float(silhouette_score(
                            np.array([embeddings[i] for i in valid_idx]), 
                            [topics[i] for i in valid_idx], 
                            metric='cosine'
                        ))
                        st.metric("Silhouette Score (Target: >0.0)", f"{sil_val:.2f}")
                    else:
                        st.metric("Silhouette Score", "Skipped (Themes need ≥2 sentences each)")

    # ==========================================
    # 5. XAI VISUALIZATION GRAPH
    # ==========================================
    st.markdown("### 📈 Explainable AI (XAI) Architecture Map")
    
    with st.spinner("Rendering Explainable AI Dashboard..."):
        sns.set_theme(style="whitegrid")
        fig = plt.figure(figsize=(16, 14)) 
        fig.suptitle(f"Topic Modeling Pipeline Analytics\n(Pooling: {ui_pooling.split()[0]} | Rep: {ui_extraction.split()[0]})", fontsize=20, fontweight='bold', y=0.98)
        box_style = dict(boxstyle="round,pad=0.4", facecolor='lightyellow', edgecolor='orange', alpha=0.9)

        # 1. Embeddings
        ax1 = plt.subplot(3, 2, 1)
        sns.heatmap(embeddings[:, :50], cmap="viridis", cbar=False, ax=ax1)
        ax1.set_title("STEP 1: Embeddings & Pooling", fontsize=13, fontweight='bold')
        ax1.set_ylabel("Sentences")
        ax1.set_xlabel("Vector Dimensions (First 50 shown)")
        ax1.text(0.5, -0.25, f"Math: {ui_embedding.split()[0]} encodes text into 384D.\nPooling '{ui_pooling.split()[0]}' squashes word vectors into 1 sentence vector.", 
                 fontsize=10, ha='center', va='top', transform=ax1.transAxes, bbox=box_style)

        # 2. Geometry
        ax2 = plt.subplot(3, 2, 2)
        reduced_embeddings = topic_model.umap_model.transform(embeddings)
        ax2.scatter(reduced_embeddings[:, 0], reduced_embeddings[:, 1], c='grey', s=100, alpha=0.6, edgecolor='k')
        ax2.set_title(f"STEP 2: Geometry ({algo_used})", fontsize=13, fontweight='bold')
        ax2.text(0.5, -0.25, f"Math: {algo_used} reduces 384D vectors into a 2D map.\nPlaces similar sentences close together.", 
                 fontsize=10, ha='center', va='top', transform=ax2.transAxes, bbox=box_style)

        # 3. Clustering
        ax3 = plt.subplot(3, 2, 3)
        ax3.scatter(reduced_embeddings[:, 0], reduced_embeddings[:, 1], c=topics, cmap='tab10', s=150, edgecolor='k')
        ax3.set_title(f"STEP 3: Clustering ({cluster_algo})", fontsize=13, fontweight='bold')
        ax3.text(0.5, -0.25, f"Math: {cluster_algo} scans the 2D space to draw boundaries.\nColors represent assigned semantic clusters.", 
                 fontsize=10, ha='center', va='top', transform=ax3.transAxes, bbox=box_style)

        # 4. Representation
        ax4 = plt.subplot(3, 2, 4)
        theme_1_data = topic_model.get_topic(0)
        if theme_1_data:
            words = [x[0] for x in theme_1_data][::-1] 
            scores = [x[1] for x in theme_1_data][::-1]
            ax4.barh(words, scores, color='coral', edgecolor='black')
            ax4.set_title(f"STEP 4: Topic Representation ({ui_extraction.split()[0]})", fontsize=13, fontweight='bold')
            ax4.text(0.5, -0.25, f"Math: Applies {ui_extraction.split()[0]} to rank vocabulary.\nLonger bars = higher semantic relevance.", 
                     fontsize=10, ha='center', va='top', transform=ax4.transAxes, bbox=box_style)
        else:
            ax4.text(0.5, 0.5, "Theme not found", ha='center', transform=ax4.transAxes)

        # 5. KPI Dashboard
        ax5 = plt.subplot(3, 2, 5)
        ax5.axis('off') 
        ax5.set_title("STEP 5: Key Performance Indicators (KPI)", fontsize=13, fontweight='bold', y=0.9)
        
        div_str = f"{div_val:.2f}" if div_val > 0 else "Skipped"
        npmi_str = f"{npmi_val:.2f}" if npmi_val != 0.0 else "Skipped"
        sil_str = f"{sil_val:.2f}" if sil_val != 0.0 else "Skipped"

        kpi_text = (
            f"📊 Topic Diversity:  {div_str} (Target: 1.0)\n\n"
            f"🧠 NPMI Coherence:  {npmi_str} (Target: >0.1)\n\n"
            f"📐 Silhouette Score:  {sil_str} (Target: >0.0)"
        )
        ax5.text(0.5, 0.4, kpi_text, fontsize=12, va='center', ha='center', 
                 bbox=dict(boxstyle="square,pad=1.5", facecolor='#e6f2ff', edgecolor='#377eb8', lw=2))
        ax5.text(0.5, -0.15, "Math: Since pipeline algorithms don't use 'Training Loss',\nthese KPIs provide the absolute mathematical grade of the topics.", 
                 fontsize=10, ha='center', va='top', transform=ax5.transAxes, bbox=box_style)

        # 6. Summary Panel
        ax6 = plt.subplot(3, 2, 6)
        ax6.axis('off') 
        summary_text = (
            "=== PIPELINE ARCHITECTURE ===\n\n"
            f"1. Embeddings: {ui_embedding.split()[0]}\n"
            f"2. Pooling: {ui_pooling.split()[0]}\n"
            f"3. N-Grams: {ngram_range}\n"
            f"4. Geometry: {algo_used}\n"
            f"5. Clustering: {cluster_algo}\n"
            f"6. Representation: {ui_extraction.split()[0]}\n\n"
            "This modular pipeline successfully transforms unstructured text\n"
            "into mathematically validated semantic domains."
        )
        ax6.text(0.1, 0.5, summary_text, fontsize=12, va='center', ha='left', 
                 bbox=dict(boxstyle="square,pad=1", facecolor='#f0f0f0', edgecolor='grey', lw=2))

        plt.subplots_adjust(hspace=0.6, wspace=0.3)
        st.pyplot(fig)