import gradio as gr
import pandas as pd
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.decomposition import NMF
import plotly.graph_objects as go
import tempfile
import os

def run_temporal_topics(file_obj, num_topics, chosen_model):
    if file_obj is None:
        return "Please upload a time-stamped text CSV or Excel sheet.", None, None, None
        
    try:
        if file_obj.name.endswith('.csv'):
            df = pd.read_csv(file_obj.name)
        else:
            df = pd.read_excel(file_obj.name)
    except Exception as e:
        return f"Error reading file: {str(e)}", None, None, None
        
    # Standardize column headers
    text_col, time_col = None, None
    for col in df.columns:
        if col.lower() in ['text', 'document', 'content', 'body', 'sentence']:
            text_col = col
        elif col.lower() in ['time', 'year', 'timestamp', 'date', 'dt', 'period']:
            time_col = col
            
    if not text_col or not time_col:
        # Fallbacks
        if len(df.columns) >= 2:
            text_col = df.columns[0]
            time_col = df.columns[1]
        else:
            return "CSV/Excel must contain at least two columns: Document Text and Time/Year.", None, None, None
            
    df = df.dropna(subset=[text_col, time_col])
    
    # Try converting time to numeric (e.g. Years)
    try:
        df[time_col] = pd.to_numeric(df[time_col]).astype(int)
        is_numeric = True
    except:
        df[time_col] = df[time_col].astype(str)
        is_numeric = False
        
    if len(df) < 10:
        return "Dataset is too small. Please provide a sheet with at least 10 rows.", None, None, None
        
    documents = df[text_col].astype(str).tolist()
    
    # 1. TF-IDF Representation
    try:
        vectorizer = TfidfVectorizer(stop_words='english', max_features=1500)
        X = vectorizer.fit_transform(documents)
    except Exception as e:
        return f"Error building vectors: {str(e)}. Ensure your texts are sufficiently long.", None, None, None
        
    # 2. Topic Modeling via NMF (highly stable and fast on CPU)
    # LDA is also NMF under the hood for this fast implementation
    try:
        nmf = NMF(n_components=num_topics, random_state=42, init='nndsvd', max_iter=1000)
        W = nmf.fit_transform(X)  # Doc-Topic matrix
        H = nmf.components_       # Topic-Word matrix
    except Exception as e:
        return f"Error training topic model: {str(e)}", None, None, None
        
    # Standardize topic weights per document (ratios sum to 1.0)
    row_sums = W.sum(axis=1, keepdims=True)
    # Avoid zero division
    row_sums[row_sums == 0] = 1.0
    W_norm = W / row_sums
    
    # Add topic weights to DataFrame
    topic_cols = []
    for i in range(num_topics):
        col_name = f"Topic {i+1}"
        df[col_name] = W_norm[:, i]
        topic_cols.append(col_name)
        
    # 3. Aggregate Topic weights by Year/Interval
    df_agg = df.groupby(time_col)[topic_cols].mean().reset_index()
    if is_numeric:
        df_agg = df_agg.sort_values(time_col)
    else:
        df_agg = df_agg.sort_values(time_col)  # alphabetical sorting for categories
        
    # 4. Generate Topic Keywords Definitions
    feature_names = np.array(vectorizer.get_feature_names_out())
    topic_keywords = []
    
    for idx, topic_comp in enumerate(H):
        top_words_idx = topic_comp.argsort()[::-1][:8]
        top_words = ", ".join(feature_names[top_words_idx])
        topic_keywords.append({
            "Topic ID": f"Topic {idx+1}",
            "Top Keywords": top_words
        })
        
    df_keywords = pd.DataFrame(topic_keywords)
    
    # 5. Generate Interactive Plotly Stacked Area Chart
    fig = go.Figure()
    
    colors = ['#ff7043', '#4db6ac', '#9575cd', '#ffd54f', '#64b5f6', '#f06292', '#81c784', '#ffffff', '#a1887f', '#ba68c8']
    
    for i, col in enumerate(topic_cols):
        color = colors[i % len(colors)]
        
        fig.add_trace(go.Scatter(
            x=df_agg[time_col],
            y=df_agg[col],
            mode='lines',
            line=dict(width=0.5, color=color),
            stackgroup='one',  # makes it a stacked area chart!
            name=f"Topic {i+1} ({df_keywords.iloc[i]['Top Keywords'][:30]}...)",
            fillcolor=color
        ))
        
    fig.update_layout(
        title="Temporal Topic Weight Evolution (Stacked Area Trend)",
        paper_bgcolor='#16100c',
        plot_bgcolor='#16100c',
        font_color='#f4eee6',
        xaxis=dict(showgrid=True, gridcolor='rgba(255,255,255,0.05)', title=str(time_col)),
        yaxis=dict(showgrid=True, gridcolor='rgba(255,255,255,0.05)', title="Relative Topic Weight (Mean)", range=[0, 1]),
        margin=dict(l=40, r=40, t=50, b=40)
    )
    
    # Save CSV
    out_csv = tempfile.mktemp(suffix=".csv")
    df.to_csv(out_csv, index=False)
    
    # Preview table
    preview_df = df_agg.round(4)
    
    return "", fig, df_keywords, preview_df, gr.update(value=out_csv, visible=True)

theme = gr.themes.Default(
    primary_hue="orange",
    neutral_hue="stone"
).set(
    body_background_fill="#0d0907",
    body_text_color="#c4bbae",
    block_background_fill="#16100c",
    block_border_width="1px",
    block_label_text_color="#f4eee6"
)

with gr.Blocks(theme=theme, title="Temporal Topic Modeler") as demo:
    gr.Markdown(
        """
        # ⏳ Chronological Topic Modeler
        ### Extract abstract topics and trace their semantic evolution and historical trajectory across time-stamped text corpora. Perfect for analyzing archival trends over years.
        """
    )
    
    error_msg = gr.Markdown("", visible=False)
    
    with gr.Row():
        with gr.Column(scale=1):
            file_obj = gr.File(label="Upload Time-stamped Document CSV", file_types=[".csv", ".xlsx"])
            gr.Markdown("💡 **Tip**: Make sure your dataset contains a **Document/Text** column and a **Year/Time** column.")
            
            num_topics = gr.Slider(minimum=2, maximum=10, value=4, step=1, label="Number of Topics to Extract")
            
            chosen_model = gr.Radio(
                choices=["NMF Topic Modeler", "Latent Dirichlet Allocation (LDA)"],
                value="NMF Topic Modeler",
                label="Topic Decomposition Model",
                info="NMF is highly recommended for stable and clear topic definitions on small-to-medium corpora."
            )
            
            btn = gr.Button("Model Topics Over Time", variant="primary")
            
        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.TabItem("Topic Timeline Trends"):
                    plot_box = gr.Plot()
                with gr.TabItem("Topic Key Terms"):
                    table_keywords = gr.Dataframe(headers=["Topic ID", "Top Keywords"])
                with gr.TabItem("Topic Year Weights Table"):
                    table_weights = gr.Dataframe()
                    download_btn = gr.File(label="Download Full Document Labeled CSV", visible=False)

    def process(file_obj, topics, model):
        err, plot, keywords, weights, csv_path = run_temporal_topics(file_obj, topics, model)
        if err:
            return gr.update(value=err, visible=True), None, None, None, gr.update(visible=False)
        return gr.update(visible=False), plot, keywords, weights, csv_path

    btn.click(
        process,
        inputs=[file_obj, num_topics, chosen_model],
        outputs=[error_msg, plot_box, table_keywords, table_weights, download_btn]
    )

if __name__ == "__main__":
    demo.launch()