import gradio as gr
import pandas as pd
import warnings
import random
import re
import time
import os
import sys
warnings.filterwarnings('ignore')

# Pre-download NLTK data at startup
import nltk
print("Downloading NLTK data...")
nltk.download('stopwords', quiet=True)
nltk.download('punkt', quiet=True)
print("NLTK data downloaded.")

# Reliable model names and descriptions - replaced KeyBERT with working alternatives
KEYWORD_MODELS = {
    'yake_yake': 'YAKE - Yet Another Keyword Extractor (statistical)',
    'tfidf_cosine': 'TF-IDF - Term frequency-inverse document frequency (statistical)',
    'pke_topicrank': 'TopicRank - graph-based keyphrase ranking',
    'pke_multipartiterank': 'MultipartiteRank - graph-based, topic-aware ranking',
    'pke_positionrank': 'PositionRank - graph-based, position-weighted ranking',
    'pke_textrank': 'TextRank - graph-based ranking inspired by PageRank',
}

# Color palette for keywords based on scores
SCORE_COLORS = {
    'high': '#0E9F6E',      # Green - High relevance
    'medium': '#B45309',    # Amber/brown - Medium relevance (dark enough for white text)
    'low': '#DC2626'        # Red - Low relevance
}

# Additional colors for variety
KEYWORD_COLORS = [
    '#4ECDC4', '#45B7D1', '#6C5CE7', '#A0E7E5', '#FD79A8',
    '#8E8E93', '#55A3FF', '#E17055', '#DDA0DD', '#FF9F43',
    '#10AC84', '#EE5A24', '#0FBC89', '#5F27CD', '#FF3838'
]

class KeywordExtractionManager:
    def __init__(self):
        self.rake_extractor = None
        self.models_initialized = False
        self.initialize_models()
        
    def initialize_models(self):
        """Pre-initialize models to check availability"""
        print("Initializing models...")
        
        # Test YAKE
        try:
            import yake
            print("✓ YAKE available")
        except ImportError as e:
            print(f"✗ YAKE not available: {e}")
            
        # Test RAKE
        try:
            from rake_nltk import Rake
            print("✓ RAKE-NLTK available")
        except ImportError as e:
            print(f"✗ RAKE-NLTK not available: {e}")
            
        # Test sklearn for TF-IDF
        try:
            from sklearn.feature_extraction.text import TfidfVectorizer
            print("✓ Scikit-learn available for TF-IDF")
        except ImportError as e:
            print(f"✗ Scikit-learn not available: {e}")
            
        # Test networkx for TextRank
        try:
            import networkx
            print("✓ NetworkX available for TextRank")
        except ImportError as e:
            print(f"✗ NetworkX not available: {e}")
            
        self.models_initialized = True
        
    def load_rake_extractor(self):
        """Load RAKE extractor with better error handling"""
        if self.rake_extractor is None:
            try:
                from rake_nltk import Rake
                
                # Create RAKE instance
                self.rake_extractor = Rake()
                print("✓ RAKE extractor loaded successfully")
            except Exception as e:
                print(f"Error loading RAKE extractor: {str(e)}")
                print(f"Full error: {type(e).__name__}: {str(e)}")
                return None
        return self.rake_extractor
        
    def extract_keywords(self, text, model_name, num_keywords=10, ngram_range=(1, 3), progress=None):
        """Extract keywords using the specified model"""
        try:
            if progress:
                progress(0.3, desc="Loading model...")
            
            print(f"Attempting to extract keywords with {model_name}")
            
            # Handle different model types
            if model_name.startswith('yake_'):
                return self.extract_yake_keywords(text, num_keywords, ngram_range, progress)
            elif model_name.startswith('tfidf_'):
                return self.extract_tfidf_cosine_keywords(text, num_keywords, ngram_range, progress)
            elif model_name.startswith('rake_'):
                return self.extract_rake_keywords(text, num_keywords, progress)
            elif model_name.startswith('textrank'):
                return self.extract_textrank_keywords(text, num_keywords, ngram_range, progress)
            elif model_name.startswith('pke_'):
                method = {
                    'pke_topicrank': 'TopicRank',
                    'pke_multipartiterank': 'MultipartiteRank',
                    'pke_positionrank': 'PositionRank',
                    'pke_textrank': 'TextRank',
                }[model_name]
                return self.extract_pke_keywords(text, num_keywords, method, progress)
            else:
                raise ValueError(f"Unknown model: {model_name}")
                
        except Exception as e:
            print(f"Error with {model_name}: {str(e)}")
            print(f"Full error: {type(e).__name__}: {str(e)}")
            return self.fallback_keyword_extraction(text, num_keywords)
    
    def extract_yake_keywords(self, text, num_keywords, ngram_range, progress):
        """Extract keywords using YAKE"""
        try:
            import yake
            
            if progress:
                progress(0.5, desc="Processing with YAKE...")
            
            # Configure YAKE
            kw_extractor = yake.KeywordExtractor(
                lan="en",
                n=ngram_range[1],
                dedupLim=0.7,
                top=num_keywords
            )
            
            if progress:
                progress(0.7, desc="Extracting keywords...")
            
            keywords = kw_extractor.extract_keywords(text)
            
            # Format results (YAKE returns lower scores for better keywords)
            results = []
            for keyword, score in keywords:
                # Invert score for consistency (higher = better)
                inverted_score = 1.0 / (1.0 + score)
                results.append({
                    'keyword': keyword,
                    'score': inverted_score,
                    'model': 'YAKE'
                })
                
            print(f"YAKE extracted {len(results)} keywords")
            return results
            
        except Exception as e:
            print(f"YAKE extraction failed: {type(e).__name__}: {str(e)}")
            return self.fallback_keyword_extraction(text, num_keywords)
    
    def extract_tfidf_cosine_keywords(self, text, num_keywords, ngram_range, progress):
        """Extract keywords using TF-IDF with cosine similarity"""
        try:
            from sklearn.feature_extraction.text import TfidfVectorizer
            from sklearn.metrics.pairwise import cosine_similarity
            import numpy as np
            
            if progress:
                progress(0.5, desc="Processing with TF-IDF...")
            
            # Create TF-IDF vectorizer
            vectorizer = TfidfVectorizer(
                ngram_range=ngram_range,
                stop_words='english',
                max_features=5000,
                min_df=1,
                max_df=0.95
            )
            
            # Extract candidate keywords/phrases
            words = re.findall(r'\b[a-z]+\b', text.lower())
            candidates = []
            
            # Generate n-grams
            for n in range(ngram_range[0], ngram_range[1] + 1):
                for i in range(len(words) - n + 1):
                    candidate = ' '.join(words[i:i+n])
                    if len(candidate) > 2 and candidate not in candidates:
                        candidates.append(candidate)
            
            if not candidates:
                return self.fallback_keyword_extraction(text, num_keywords)
            
            # Limit candidates to prevent memory issues
            candidates = candidates[:300]
            
            if progress:
                progress(0.7, desc="Computing similarities...")
            
            try:
                # Create document embedding
                doc_embedding = vectorizer.fit_transform([text])
                
                # Create embeddings for candidates
                candidate_embeddings = vectorizer.transform(candidates)
                
                # Calculate similarities
                similarities = cosine_similarity(doc_embedding, candidate_embeddings)[0]
                
                # Get top keywords
                top_indices = similarities.argsort()[-num_keywords:][::-1]
                
                results = []
                for idx in top_indices:
                    if similarities[idx] > 0:
                        results.append({
                            'keyword': candidates[idx],
                            'score': float(similarities[idx]),
                            'model': 'TF-IDF-Cosine'
                        })
                
                if progress:
                    progress(0.8, desc="Formatting results...")
                
                print(f"TF-IDF extracted {len(results)} keywords")
                return results
                
            except Exception as e:
                print(f"TF-IDF approach failed: {e}")
                # Fall back to simple TF-IDF
                return self.simple_tfidf_extraction(text, num_keywords, ngram_range)
                
        except ImportError:
            print("scikit-learn not available for TF-IDF")
            return self.fallback_keyword_extraction(text, num_keywords)
        except Exception as e:
            print(f"TF-IDF extraction failed: {e}")
            return self.fallback_keyword_extraction(text, num_keywords)
    
    def extract_textrank_keywords(self, text, num_keywords, ngram_range, progress):
        """Extract keywords using TextRank algorithm"""
        try:
            import numpy as np
            from sklearn.feature_extraction.text import TfidfVectorizer
            from sklearn.metrics.pairwise import cosine_similarity
            import networkx as nx
            
            if progress:
                progress(0.5, desc="Processing with TextRank...")
            
            # Split text into sentences
            sentences = re.split(r'[.!?]+', text)
            sentences = [s.strip() for s in sentences if s.strip()]
            
            if len(sentences) < 2:
                # If text is too short, use simple extraction
                return self.simple_tfidf_extraction(text, num_keywords, ngram_range)
            
            # Create TF-IDF matrix
            vectorizer = TfidfVectorizer(
                ngram_range=(1, 1),  # Use unigrams for sentence similarity
                stop_words='english'
            )
            
            tfidf_matrix = vectorizer.fit_transform(sentences)
            
            # Calculate similarity matrix
            similarity_matrix = cosine_similarity(tfidf_matrix)
            
            if progress:
                progress(0.6, desc="Building graph...")
            
            # Build graph
            nx_graph = nx.from_numpy_array(similarity_matrix)
            
            # Calculate PageRank scores
            scores = nx.pagerank(nx_graph)
            
            # Extract keywords from top-ranked sentences
            top_sentence_indices = sorted(scores.keys(), key=lambda x: scores[x], reverse=True)[:3]
            
            # Extract keywords from top sentences
            keyword_vectorizer = TfidfVectorizer(
                ngram_range=ngram_range,
                stop_words='english',
                max_features=num_keywords * 2
            )
            
            top_sentences = [sentences[i] for i in top_sentence_indices]
            top_text = ' '.join(top_sentences)
            
            if progress:
                progress(0.7, desc="Extracting keywords...")
            
            tfidf_matrix = keyword_vectorizer.fit_transform([top_text])
            feature_names = keyword_vectorizer.get_feature_names_out()
            tfidf_scores = tfidf_matrix.toarray()[0]
            
            # Get top keywords
            top_indices = tfidf_scores.argsort()[-num_keywords:][::-1]
            
            results = []
            for idx in top_indices:
                if tfidf_scores[idx] > 0:
                    results.append({
                        'keyword': feature_names[idx],
                        'score': float(tfidf_scores[idx]),
                        'model': 'TextRank'
                    })
            
            print(f"TextRank extracted {len(results)} keywords")
            return results
            
        except ImportError as e:
            print(f"Required library not available for TextRank: {e}")
            return self.fallback_keyword_extraction(text, num_keywords)
        except Exception as e:
            print(f"TextRank extraction failed: {e}")
            return self.fallback_keyword_extraction(text, num_keywords)
    
    def simple_tfidf_extraction(self, text, num_keywords, ngram_range):
        """Simple TF-IDF extraction without cosine similarity"""
        try:
            from sklearn.feature_extraction.text import TfidfVectorizer
            
            vectorizer = TfidfVectorizer(
                ngram_range=ngram_range,
                stop_words='english',
                max_features=num_keywords * 2
            )
            
            # Fit and transform
            tfidf_matrix = vectorizer.fit_transform([text])
            
            # Get feature names and scores
            feature_names = vectorizer.get_feature_names_out()
            scores = tfidf_matrix.toarray()[0]
            
            # Get top keywords
            top_indices = scores.argsort()[-num_keywords:][::-1]
            
            results = []
            for idx in top_indices:
                if scores[idx] > 0:
                    results.append({
                        'keyword': feature_names[idx],
                        'score': float(scores[idx]),
                        'model': 'TF-IDF-Simple'
                    })
            
            return results
            
        except Exception as e:
            print(f"Simple TF-IDF failed: {e}")
            return self.fallback_keyword_extraction(text, num_keywords)
    
    def extract_rake_keywords(self, text, num_keywords, progress):
        """Extract keywords using RAKE"""
        try:
            if progress:
                progress(0.5, desc="Processing with RAKE...")
            
            rake_extractor = self.load_rake_extractor()
            if rake_extractor is None:
                print("RAKE extractor could not be loaded")
                return self.fallback_keyword_extraction(text, num_keywords)
            
            if progress:
                progress(0.7, desc="Extracting keywords...")
            
            # Extract keywords
            rake_extractor.extract_keywords_from_text(text)
            keywords_with_scores = rake_extractor.get_ranked_phrases_with_scores()
            
            # Normalize scores
            if keywords_with_scores:
                max_score = max(score for score, _ in keywords_with_scores)
                
                # Format results
                results = []
                for score, keyword in keywords_with_scores[:num_keywords]:
                    normalized_score = score / max_score if max_score > 0 else 0
                    results.append({
                        'keyword': keyword,
                        'score': normalized_score,
                        'model': 'RAKE-NLTK'
                    })
                
                print(f"RAKE extracted {len(results)} keywords")
                return results
            else:
                print("RAKE returned no keywords")
                return self.fallback_keyword_extraction(text, num_keywords)
            
        except Exception as e:
            print(f"RAKE extraction failed: {type(e).__name__}: {str(e)}")
            return self.fallback_keyword_extraction(text, num_keywords)

    def extract_pke_keywords(self, text, num_keywords, method, progress):
        """Extract keyphrases using the pke library (paper's methods)."""
        try:
            import pke

            if progress:
                progress(0.5, desc=f"Processing with {method}...")

            extractor_classes = {
                'TopicRank': pke.unsupervised.TopicRank,
                'MultipartiteRank': pke.unsupervised.MultipartiteRank,
                'PositionRank': pke.unsupervised.PositionRank,
                'TextRank': pke.unsupervised.TextRank,
            }
            if method not in extractor_classes:
                return self.fallback_keyword_extraction(text, num_keywords)

            extractor = extractor_classes[method]()
            extractor.load_document(input=text, language='en')
            extractor.candidate_selection()
            extractor.candidate_weighting()

            if progress:
                progress(0.7, desc="Extracting keywords...")

            keyphrases = extractor.get_n_best(n=num_keywords)

            results = []
            if keyphrases:
                max_score = max(score for _, score in keyphrases) or 1.0
                for phrase, score in keyphrases:
                    results.append({
                        'keyword': phrase,
                        'score': float(score) / float(max_score),
                        'model': method,
                    })
            print(f"{method} extracted {len(results)} keywords")
            return results if results else self.fallback_keyword_extraction(text, num_keywords)

        except Exception as e:
            print(f"{method} extraction failed: {type(e).__name__}: {str(e)}")
            return self.fallback_keyword_extraction(text, num_keywords)

    def fallback_keyword_extraction(self, text, num_keywords=10):
        """Simple fallback keyword extraction using basic statistics"""
        print("Using fallback keyword extraction")
        import re
        from collections import Counter
        
        # Simple tokenization and filtering
        words = re.findall(r'\b[a-z]+\b', text.lower())
        
        # Remove common stop words
        stop_words = {'the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for',
                     'of', 'with', 'by', 'from', 'as', 'is', 'was', 'are', 'were', 'been',
                     'be', 'have', 'has', 'had', 'do', 'does', 'did', 'will', 'would',
                     'could', 'should', 'may', 'might', 'must', 'can', 'this', 'that',
                     'these', 'those', 'i', 'you', 'he', 'she', 'it', 'we', 'they'}
        
        filtered_words = [w for w in words if w not in stop_words and len(w) > 3]
        
        # Count frequencies
        word_freq = Counter(filtered_words)
        
        # Get top keywords
        results = []
        for word, freq in word_freq.most_common(num_keywords):
            score = freq / len(filtered_words)  # Normalize by total words
            results.append({
                'keyword': word,
                'score': score,
                'model': 'Fallback-TFIDF'
            })
            
        return results

def get_score_color(score, max_score):
    """Get color based on score relative to max score"""
    if max_score == 0:
        return SCORE_COLORS['medium']
    
    relative_score = score / max_score
    if relative_score >= 0.7:
        return SCORE_COLORS['high']
    elif relative_score >= 0.4:
        return SCORE_COLORS['medium']
    else:
        return SCORE_COLORS['low']

def get_relevance_level(score, max_score):
    """Get relevance level name based on score"""
    if max_score == 0:
        return 'medium'
    
    relative_score = score / max_score
    if relative_score >= 0.7:
        return 'high'
    elif relative_score >= 0.4:
        return 'medium'
    else:
        return 'low'

def create_highlighted_html(text, keywords):
    """Create HTML with highlighted keywords in the text"""
    if not keywords:
        return f"<div style='padding: 15px; border: 1px solid #ddd; border-radius: 5px; background-color: #fafafa;'><p>{text}</p></div>"
    
    # Sort keywords by length (longest first) to avoid partial matches
    sorted_keywords = sorted(keywords, key=lambda x: len(x['keyword']), reverse=True)
    
    # Get max score for color scaling
    max_score = max(k['score'] for k in keywords) if keywords else 1
    
    # Create a modified text with highlights
    highlighted_text = text
    for i, kw_data in enumerate(sorted_keywords):
        keyword = kw_data['keyword']
        score = kw_data['score']
        color = get_score_color(score, max_score)
        
        # Create regex pattern for whole word matching (case-insensitive)
        pattern = r'\b' + re.escape(keyword) + r'\b'
        
        # Replace with highlighted version
        replacement = f'<span style="background-color: {color}; padding: 2px 4px; ' \
                     f'border-radius: 3px; margin: 0 1px; ' \
                     f'border: 1px solid {color}; color: white; font-weight: bold;" ' \
                     f'title="Score: {score:.3f}">{keyword}</span>'
        
        highlighted_text = re.sub(pattern, replacement, highlighted_text, flags=re.IGNORECASE)
    
    return f"""
    <div style='padding: 15px; border: 2px solid #ddd; border-radius: 8px; background-color: #fafafa; margin: 10px 0;'>
        <h4 style='margin: 0 0 15px 0; color: #333;'>📝 Text with Highlighted Keywords</h4>
        <div style='line-height: 1.8; font-size: 16px; background-color: white; padding: 15px; border-radius: 5px;'>{highlighted_text}</div>
    </div>
    """

def create_keyword_table_html(keywords):
    """Create HTML table for keywords"""
    if not keywords:
        return "<p style='text-align: center; padding: 20px;'>No keywords found.</p>"
    
    # Sort by score
    sorted_keywords = sorted(keywords, key=lambda x: x['score'], reverse=True)
    max_score = sorted_keywords[0]['score'] if sorted_keywords else 1
    
    table_html = """
    <div style='max-height: 600px; overflow-y: auto; border: 2px solid #ddd; border-radius: 8px; padding: 20px; background-color: #fafafa;'>
        <h3 style="margin: 0 0 20px 0;">🎯 Extracted Keywords</h3>
        <table style="width: 100%; border-collapse: collapse; border: 1px solid #ddd; background-color: white;">
            <thead>
                <tr style="background-color: #4ECDC4; color: white;">
                    <th style="padding: 12px; text-align: left; border: 1px solid #ddd;">Rank</th>
                    <th style="padding: 12px; text-align: left; border: 1px solid #ddd;">Keyword</th>
                    <th style="padding: 12px; text-align: left; border: 1px solid #ddd;">Score</th>
                    <th style="padding: 12px; text-align: left; border: 1px solid #ddd;">Relevance</th>
                    <th style="padding: 12px; text-align: left; border: 1px solid #ddd;">Model</th>
                </tr>
            </thead>
            <tbody>
    """
    
    for i, kw_data in enumerate(sorted_keywords):
        score = kw_data['score']
        color = get_score_color(score, max_score)
        
        # Create relevance bar
        bar_width = int((score / max_score) * 100) if max_score > 0 else 0
        relevance_bar = f"""
        <div style="width: 100%; background-color: #e0e0e0; border-radius: 10px; height: 20px;">
            <div style="width: {bar_width}%; background-color: {color}; height: 100%; border-radius: 10px;"></div>
        </div>
        """
        
        table_html += f"""
            <tr style="background-color: #fff;">
                <td style="padding: 10px; border: 1px solid #ddd; text-align: center; font-weight: bold;">#{i+1}</td>
                <td style="padding: 10px; border: 1px solid #ddd; font-weight: bold;">{kw_data['keyword']}</td>
                <td style="padding: 10px; border: 1px solid #ddd;">
                    <span style="color: {color}; font-weight: bold;">{score:.4f}</span>
                </td>
                <td style="padding: 10px; border: 1px solid #ddd;">{relevance_bar}</td>
                <td style="padding: 10px; border: 1px solid #ddd;">
                    <span style='background-color: #007bff; color: white; padding: 2px 6px; border-radius: 10px; font-size: 11px;'>
                        {kw_data['model']}
                    </span>
                </td>
            </tr>
        """
    
    table_html += """
            </tbody>
        </table>
    </div>
    """
    
    return table_html

def create_legend_html():
    """Create a legend showing score colors"""
    html = """
    <div style='margin: 15px 0; padding: 15px; background-color: #f8f9fa; border-radius: 8px;'>
        <h4 style='margin: 0 0 15px 0;'>🎨 Relevance Score Legend</h4>
        <div style='display: flex; flex-wrap: wrap; gap: 15px;'>
            <span style='background-color: #0E9F6E; padding: 4px 12px; border-radius: 15px; color: white; font-weight: bold;'>
                High Relevance (70%+)
            </span>
            <span style='background-color: #B45309; padding: 4px 12px; border-radius: 15px; color: white; font-weight: bold;'>
                Medium Relevance (40-70%)
            </span>
            <span style='background-color: #DC2626; padding: 4px 12px; border-radius: 15px; color: white; font-weight: bold;'>
                Low Relevance (&lt;40%)
            </span>
        </div>
    </div>
    """
    return html

# Initialize the keyword extraction manager
print("Initializing keyword extraction manager...")
keyword_manager = KeywordExtractionManager()

def process_text(text, selected_model, num_keywords, ngram_min, ngram_max, progress=gr.Progress()):
    """Main processing function for Gradio interface with progress tracking"""
    
    if not text.strip():
        return "❌ Please enter some text to analyse", "", ""
    
    progress(0.1, desc="Initialising...")
    
    # Extract keywords
    progress(0.2, desc="Extracting keywords...")
    keywords = keyword_manager.extract_keywords(
        text, 
        selected_model, 
        num_keywords=num_keywords,
        ngram_range=(ngram_min, ngram_max),
        progress=progress
    )
    
    if not keywords:
        return "❌ No keywords found. Try adjusting the parameters.", "", ""
    
    progress(0.8, desc="Processing results...")
    
    # Create outputs
    legend_html = create_legend_html()
    highlighted_html = create_highlighted_html(text, keywords)
    results_html = create_keyword_table_html(keywords)
    
    progress(0.9, desc="Creating summary...")
    
    # Create summary
    avg_score = sum(k['score'] for k in keywords) / len(keywords)
    model_display = KEYWORD_MODELS.get(selected_model, selected_model).split(' - ')[0]
    summary = f"""
    ## 📊 Analysis Summary
    - **Keywords extracted:** {len(keywords)}
    - **Model used:** {model_display}
    - **Average relevance score:** {avg_score:.4f}
    - **N-gram range:** {ngram_min}-{ngram_max} words
    """
    
    progress(1.0, desc="Complete!")
    
    return summary, legend_html + highlighted_html, results_html

# Create Gradio interface
def create_interface():
    with gr.Blocks(title="Keyword Extraction Tool", theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # Keyword Extraction Explorer Tool
        
        Extract the most important keywords and phrases from your text using various algorithms! This tool compares statistical methods (YAKE, TF-IDF) with graph-based methods (TopicRank, MultipartiteRank, PositionRank, TextRank) so you can see how each approaches the same text.
          
        ### How to use:
        1. **📝 Enter your text** in the text area below
        2. **🎯 Select a model** from the dropdown for keyword extraction
        3. **⚙️ Adjust parameters** (number of keywords, n-gram range)
        4. **🔍 Click "Extract Keywords"** to see results with organized output
        """)
        
        # Add tip box
        gr.HTML("""
        <div style="background-color: #fff3cd; border: 1px solid #ffeaa7; border-radius: 8px; padding: 12px; margin: 15px 0;">
            <strong style="color: #856404;">💡 Top tip:</strong> Different models excel at different types of texts - experiment to find the best one for your content!
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=2):
                text_input = gr.Textbox(
                    label="📝 Text to Analyse",
                    placeholder="Enter your text here...",
                    lines=18,
                    max_lines=20
                )
            
            with gr.Column(scale=1):
                # Model selector
                model_dropdown = gr.Dropdown(
                    choices=list(KEYWORD_MODELS.keys()),
                    value='yake_yake',
                    label="🎯 Select Keyword Extraction Model"
                )
                
                # Parameters
                num_keywords = gr.Slider(
                    minimum=5,
                    maximum=30,
                    value=10,
                    step=1,
                    label="📊 Number of Keywords"
                )
                
                with gr.Row():
                    ngram_min = gr.Slider(
                        minimum=1,
                        maximum=3,
                        value=1,
                        step=1,
                        label="Min N-gram"
                    )
                    ngram_max = gr.Slider(
                        minimum=1,
                        maximum=4,
                        value=3,
                        step=1,
                        label="Max N-gram"
                    )
                
                # Add N-gram tip box
                gr.HTML("""
                <div style="background-color: #fff3cd; border: 1px solid #ffeaa7; border-radius: 8px; padding: 10px; margin: 10px 0;">
                    <strong style="color: #856404;">💡 Top tip:</strong> N-grams are sequences of words. Set Min=1, Max=3 to extract single words, phrases of 2 words, and phrases of 3 words. Higher values capture longer phrases but may reduce precision.
                </div>
                """)
        
        # Add model descriptions
        gr.HTML("""
        <details style="margin: 20px 0; padding: 10px; background-color: #f8f9fa; border-radius: 8px; border: 1px solid #ddd;">
            <summary style="cursor: pointer; font-weight: bold; padding: 5px; color: #1976d2;">
                ℹ️ Model Descriptions
            </summary>
            <div style="margin-top: 10px; padding: 10px;">
                <p style="font-size: 14px; margin: 0 0 12px 0;">All six methods here are <strong>statistical or graph-based</strong>. They need no training data and work directly from the words in your text. They differ in how they decide which words matter.</p>
                <dl style="margin: 0; font-size: 14px; line-height: 1.5;">
                    <div style="margin-bottom: 10px;">
                        <dt style="font-weight: bold; display: inline; color: #0E9F6E;">YAKE</dt><span style="font-size: 12px; color: #555;"> (statistical):</span>
                        <dd style="display: inline; margin-left: 5px;">Scores each word using simple text features &mdash; how often it appears, where it sits, whether it's capitalised, and what surrounds it. Good all-rounder, works on short texts and many languages.</dd>
                    </div>
                    <div style="margin-bottom: 10px;">
                        <dt style="font-weight: bold; display: inline; color: #0E9F6E;">TF-IDF</dt><span style="font-size: 12px; color: #555;"> (statistical):</span>
                        <dd style="display: inline; margin-left: 5px;">Picks out words that are common in <em>your</em> text but rare in general &mdash; the ones that make it distinctive. A long-standing, dependable baseline.</dd>
                    </div>
                    <div style="margin-bottom: 10px;">
                        <dt style="font-weight: bold; display: inline; color: #1976d2;">TopicRank</dt><span style="font-size: 12px; color: #555;"> (graph-based):</span>
                        <dd style="display: inline; margin-left: 5px;">Groups similar words into topics first, then ranks the topics and picks one phrase to represent each. Reduces near-duplicates, so results cover a broader spread of themes.</dd>
                    </div>
                    <div style="margin-bottom: 10px;">
                        <dt style="font-weight: bold; display: inline; color: #1976d2;">MultipartiteRank</dt><span style="font-size: 12px; color: #555;"> (graph-based):</span>
                        <dd style="display: inline; margin-left: 5px;">A refinement of TopicRank that keeps individual phrases as the unit while still respecting topics. Tends to give the most balanced, precise keyphrases (the strongest performer in our tests).</dd>
                    </div>
                    <div style="margin-bottom: 10px;">
                        <dt style="font-weight: bold; display: inline; color: #1976d2;">PositionRank</dt><span style="font-size: 12px; color: #555;"> (graph-based):</span>
                        <dd style="display: inline; margin-left: 5px;">Gives extra weight to words that appear early in the text, on the idea that important terms are often introduced up front. Useful for structured writing like articles and abstracts.</dd>
                    </div>
                    <div style="margin-bottom: 10px;">
                        <dt style="font-weight: bold; display: inline; color: #1976d2;">TextRank</dt><span style="font-size: 12px; color: #555;"> (graph-based):</span>
                        <dd style="display: inline; margin-left: 5px;">Builds a network of words that appear near each other and ranks them using the same idea as Google's PageRank &mdash; a word is important if other important words connect to it.</dd>
                    </div>
                </dl>
                <p style="font-size: 13px; margin: 12px 0 0 0; color: #555;"><em>TopicRank, MultipartiteRank, PositionRank and TextRank are graph-based methods provided by the <a href="https://github.com/boudinfl/pke" target="_blank" style="color: #1976d2;">pke</a> library.</em></p>
            </div>
        </details>
        """)
        
        extract_btn = gr.Button("🔍 Extract Keywords", variant="primary", size="lg")
        
        # Output sections
        with gr.Row():
            summary_output = gr.Markdown(label="Summary")
        
        with gr.Row():
            highlighted_output = gr.HTML(label="Highlighted Text")
        
        # Results section
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 📋 Detailed Results")
                results_output = gr.HTML(label="Keyword Results")
        
        # Connect the button to the processing function
        extract_btn.click(
            fn=process_text,
            inputs=[
                text_input,
                model_dropdown,
                num_keywords,
                ngram_min,
                ngram_max
            ],
            outputs=[summary_output, highlighted_output, results_output]
        )
        
        gr.Examples(
            examples=[
                [
                    "On June 6, 1944, Allied forces launched Operation Overlord, the invasion of Normandy. General Dwight D. Eisenhower commanded the operation, while Field Marshal Bernard Montgomery led ground forces. The BBC broadcast coded messages to the French Resistance, including the famous line 'The long sobs of autumn violins.'",
                    "yake_yake",
                    10,
                    1,
                    3
                ],
                [
                    "In Jane Austen's 'Pride and Prejudice', Elizabeth Bennet first meets Mr. Darcy at the Meryton assembly. The novel, published in 1813, explores themes of marriage and social class in Regency England. Austen wrote to her sister Cassandra about the manuscript while staying at Chawton Cottage.",
                    "tfidf_cosine",
                    10,
                    1,
                    3
                ],
                [
                    "Charles Darwin arrived at the Galápagos Islands aboard HMS Beagle in September 1835. During his five-week visit, Darwin collected specimens of finches, tortoises, and mockingbirds. His observations of these species' variations across different islands later contributed to his theory of evolution by natural selection, published in 'On the Origin of Species' in 1859.",
                    "pke_multipartiterank",
                    10,
                    1,
                    3
                ]
            ],
            inputs=[
                text_input,
                model_dropdown,
                num_keywords,
                ngram_min,
                ngram_max
            ]
        )
        
        # Add model information links
        gr.HTML("""
        <hr style="margin-top: 40px; margin-bottom: 20px;">
        <div style="background-color: #f8f9fa; padding: 20px; border-radius: 8px; margin-top: 20px;">
            <h4 style="margin-top: 0;">📚 Model Information & Documentation</h4>
            <p style="font-size: 14px; margin-bottom: 15px;">Learn more about the algorithms used in this tool:</p>
            <ul style="font-size: 14px; line-height: 1.8;">
                <li><strong>YAKE:</strong> 
                    <a href="https://github.com/LIAAD/yake" target="_blank" style="color: #1976d2;">
                        Yet Another Keyword Extractor ↗
                    </a>
                </li>
                <li><strong>TF-IDF:</strong> 
                    <a href="https://scikit-learn.org/stable/modules/feature_extraction.html#tfidf-term-weighting" target="_blank" style="color: #1976d2;">
                        Term Frequency-Inverse Document Frequency ↗
                    </a>
                </li>
                <li><strong>TopicRank:</strong> 
                    <a href="https://aclanthology.org/I13-1062/" target="_blank" style="color: #1976d2;">
                        Bougouin et al. (2013) ↗
                    </a>
                </li>
                <li><strong>MultipartiteRank:</strong> 
                    <a href="https://aclanthology.org/N18-2105/" target="_blank" style="color: #1976d2;">
                        Boudin (2018) ↗
                    </a>
                </li>
                <li><strong>PositionRank:</strong> 
                    <a href="https://aclanthology.org/P17-1102/" target="_blank" style="color: #1976d2;">
                        Florescu &amp; Caragea (2017) ↗
                    </a>
                </li>
                <li><strong>TextRank:</strong> 
                    <a href="https://aclanthology.org/W04-3252/" target="_blank" style="color: #1976d2;">
                        Mihalcea &amp; Tarau (2004) ↗
                    </a>
                </li>
                <li><strong>pke library</strong> (TopicRank, MultipartiteRank, PositionRank, TextRank): 
                    <a href="https://github.com/boudinfl/pke" target="_blank" style="color: #1976d2;">
                        Python Keyphrase Extraction ↗
                    </a>
                </li>
            </ul>
        </div>

        <br>
        <hr style="margin-top: 40px; margin-bottom: 20px;">
        <div style="background-color: #f8f9fa; padding: 20px; border-radius: 8px; margin-top: 20px; text-align: center;">
            <p style="font-size: 14px; line-height: 1.8; margin: 0;">
                This <strong>Keyword Extraction Explorer Tool</strong> was created as part of the
                <a href="https://digitalscholarship.web.ox.ac.uk/" target="_blank" style="color: #1976d2;">Digital Scholarship at Oxford (DiSc)</a>
                funded research project: <em>Extracting Keywords from Crowdsourced Collections</em>.
            </p>
            <p style="font-size: 14px; line-height: 1.8; margin: 20px 0 0 0;">
                <strong>See also</strong> the main EKCC project repository:
                <a href="https://github.com/Digital-Scholarship-Oxford/crowdsourced-data-tools" target="_blank" style="color: #1976d2;">Crowdsourced Data Tools ↗</a>
            </p>
            <p style="font-size: 14px; line-height: 1.8; margin: 8px 0 0 0;">
                <strong>and</strong> Arana-Catania, M., Conisbee, C., &amp; Kidd, M. (2026).
                <em>Letting the Data Speak: Extracting Keywords from Crowdsourced Collections with AI.</em>
                <a href="https://doi.org/10.48550/arXiv.2607.09324" target="_blank" style="color: #1976d2;">https://doi.org/10.48550/arXiv.2607.09324 ↗</a>
            </p>
            <p style="font-size: 14px; line-height: 1.8; margin: 20px 0 0 0;">
                The code for this tool was built with the aid of Claude Opus 4 and updated with Opus 4.8.
            </p>
        </div>
        """)

    return demo

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