import os
import gradio as gr
import pandas as pd
import numpy as np
import plotly.express as px
from sklearn.feature_extraction.text import TfidfVectorizer

# Try to import RAKE, fall back to basic scoring if it fails
try:
    from rake_nltk import Rake
    HAS_RAKE = True
except ImportError:
    HAS_RAKE = False

def load_data(file_obj):
    """Safely loads CSV, Excel, or TXT file into a Pandas DataFrame."""
    if file_obj is None:
        return None, gr.update(choices=[], visible=False), "Please upload a file."
    
    file_path = file_obj.name
    ext = os.path.splitext(file_path)[1].lower()
    
    try:
        if ext == '.csv':
            df = pd.read_csv(file_path)
        elif ext in ['.xls', '.xlsx']:
            df = pd.read_excel(file_path)
        elif ext == '.txt':
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            df = pd.DataFrame({'text': [content]})
        else:
            return None, gr.update(choices=[], visible=False), "Unsupported file format. Please upload .csv, .xlsx, or .txt."
        
        string_cols = [col for col in df.columns if df[col].dtype == 'object' or df[col].astype(str).str.len().mean() > 5]
        if not string_cols:
            string_cols = list(df.columns)
            
        return df, gr.update(choices=string_cols, value=string_cols[0], visible=True), f"Successfully loaded dataset with {len(df)} rows."
    except Exception as e:
        return None, gr.update(choices=[], visible=False), f"Error loading file: {str(e)}"

def run_tfidf(docs, top_n):
    """Extracts top keywords using TF-IDF."""
    vectorizer = TfidfVectorizer(stop_words='english', max_df=0.85, min_df=1)
    try:
        tfidf_matrix = vectorizer.fit_transform(docs)
        feature_names = vectorizer.get_feature_names_out()
        
        # Sum TF-IDF scores across all docs
        scores = tfidf_matrix.sum(axis=0).A1
        top_indices = scores.argsort()[::-1][:top_n]
        
        results = []
        for idx in top_indices:
            results.append({
                "Keyword": feature_names[idx],
                "Score": round(float(scores[idx]), 4),
                "Method": "TF-IDF"
            })
        return pd.DataFrame(results)
    except Exception as e:
        # Fallback to simple term frequency if vocabulary is too empty
        from collections import Counter
        import re
        words = []
        for doc in docs:
            words.extend(re.findall(r'\b[a-zA-Z]{3,}\b', doc.lower()))
        # Filter stopwords manually
        stopwords = {'the', 'and', 'for', 'that', 'with', 'this', 'have', 'from', 'your', 'will', 'not', 'are'}
        words = [w for w in words if w not in stopwords]
        counts = Counter(words).most_common(top_n)
        results = [{"Keyword": w, "Score": float(c), "Method": "Frequency Count"} for w, c in counts]
        return pd.DataFrame(results)

def run_rake(docs, top_n):
    """Extracts phrases and words using RAKE (or fallback)."""
    full_text = " ".join(docs)
    
    if HAS_RAKE:
        try:
            r = Rake()
            r.extract_keywords_from_text(full_text)
            ranked_phrases = r.get_ranked_phrases_with_scores()[:top_n]
            results = [{"Keyword": phrase, "Score": round(score, 4), "Method": "RAKE"} for score, phrase in ranked_phrases]
            return pd.DataFrame(results)
        except Exception:
            pass
            
    # Fallback to bigram/trigram phrase extraction if RAKE is unavailable or fails
    from sklearn.feature_extraction.text import CountVectorizer
    vectorizer = CountVectorizer(ngram_range=(2, 3), stop_words='english', min_df=1)
    try:
        dtm = vectorizer.fit_transform(docs)
        phrases = vectorizer.get_feature_names_out()
        counts = dtm.sum(axis=0).A1
        top_indices = counts.argsort()[::-1][:top_n]
        results = [{"Keyword": phrases[idx], "Score": float(counts[idx]), "Method": "N-gram Frequency"} for idx in top_indices]
        return pd.DataFrame(results)
    except Exception as e:
        return pd.DataFrame([{"Keyword": "No phrases found", "Score": 0.0, "Method": "Error"}])

def extract_keywords(text_input, file_obj, text_col, method, top_n):
    docs = []
    if file_obj is not None:
        df, _, _ = load_data(file_obj)
        if df is not None and text_col in df.columns:
            docs = df[text_col].astype(str).fillna("").tolist()
    elif text_input and text_input.strip():
        docs = [text_input]
        
    if not docs:
        return None, None, None, "Please enter text or upload a valid dataset first."
        
    try:
        if method == "TF-IDF (Word Frequency)":
            df_res = run_tfidf(docs, top_n)
        else:
            df_res = run_rake(docs, top_n)
            
        if df_res.empty:
            return None, None, None, "No keywords were successfully extracted."
            
        # Plotly Bar Chart
        fig = px.bar(
            df_res,
            x="Score",
            y="Keyword",
            orientation="h",
            color="Score",
            title=f"Top Extracted Keywords via {method}",
            template="plotly_dark",
            color_continuous_scale="Viridis"
        )
        fig.update_layout(yaxis={'categoryorder': 'total ascending'}, height=400, margin=dict(l=20, r=20, t=40, b=20))
        
        # Export CSV
        csv_path = "extracted_keywords.csv"
        df_res.to_csv(csv_path, index=False)
        
        return df_res, fig, csv_path, "Keywords extracted successfully!"
    except Exception as e:
        return None, None, None, f"Execution failed: {str(e)}"

custom_css = """
body {
    background-color: #0b0f19;
    color: #f3f4f6;
}
.gradio-container {
    font-family: 'Inter', sans-serif !important;
}
h1, h2 {
    color: #6366f1 !important;
}
"""

with gr.Blocks(theme=gr.themes.Default(primary_hue="indigo", secondary_hue="slate"), css=custom_css) as demo:
    df_state = gr.State()
    
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 2rem;">
        <h1 style="font-size: 2.5rem; font-weight: 700; margin-bottom: 0.5rem; background: linear-gradient(to right, #6366f1, #a855f7); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">Interactive Keyword Extractor</h1>
        <p style="font-size: 1.1rem; color: #94a3b8; max-width: 800px; margin: 0 auto;">
            Automatically extract the most significant words or phrases from text documents.
            Compare TF-IDF statistical scoring with RAKE phrase-level extraction.
        </p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. Choose Input")
            with gr.Tabs():
                with gr.TabItem("Paste Text"):
                    text_input = gr.Textbox(
                        label="Source Text",
                        placeholder="Paste your text here to analyze...",
                        lines=10
                    )
                with gr.TabItem("Upload Dataset File"):
                    file_input = gr.File(label="Upload (.csv, .xlsx, .txt)", file_types=[".csv", ".xlsx", ".txt"])
                    text_column_selector = gr.Dropdown(
                        label="Target Text Column", 
                        choices=[], 
                        visible=False,
                        interactive=True
                    )
                    status_text = gr.Markdown("No file uploaded yet.")
                    
            gr.Markdown("### 2. Configure Extraction")
            method_selector = gr.Radio(
                choices=["TF-IDF (Word Frequency)", "RAKE (Phrase Extraction)"],
                value="TF-IDF (Word Frequency)",
                label="Extraction Algorithm"
            )
            top_n = gr.Slider(minimum=3, maximum=40, value=15, step=1, label="Number of Keywords to Extract")
            
            run_btn = gr.Button("Extract Keywords", variant="primary")
            
        with gr.Column(scale=2):
            gr.Markdown("### 3. Results & Visualizations")
            status_markdown = gr.Markdown("Click 'Extract Keywords' to begin.")
            
            with gr.Tabs():
                with gr.TabItem("Keywords Score Plot"):
                    chart_output = gr.Plot(label="Keywords Score Plot")
                with gr.TabItem("Keywords Table"):
                    table_output = gr.Dataframe(
                        headers=["Keyword", "Score", "Method"],
                        datatype=["str", "number", "str"],
                        interactive=False,
                        wrap=True
                    )
                    
            gr.Markdown("### 4. Export")
            download_csv = gr.File(label="Download Keywords Report (CSV)")

    file_input.change(
        fn=load_data,
        inputs=file_input,
        outputs=[df_state, text_column_selector, status_text]
    )
    
    run_btn.click(
        fn=extract_keywords,
        inputs=[text_input, file_input, text_column_selector, method_selector, top_n],
        outputs=[table_output, chart_output, download_csv, status_markdown]
    )

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