import gradio as gr
from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch
import spacy
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
import re
import subprocess

# Download and load spaCy model
try:
    nlp = spacy.load("en_core_web_sm")
except OSError:
    print("Downloading spaCy model...")
    subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])
    nlp = spacy.load("en_core_web_sm")

# Load models
MODEL_NAME = "bert-large-uncased"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForMaskedLM.from_pretrained(MODEL_NAME)
model.eval()

# Load sentence transformer for semantic similarity
similarity_model = SentenceTransformer('all-MiniLM-L6-v2')

# POS tag mapping
POS_MAP = {
    'NOUN': ['NOUN', 'PROPN'],
    'VERB': ['VERB'],
    'ADJ': ['ADJ'],
    'ADV': ['ADV'],
    'PROPN': ['NOUN', 'PROPN']
}

def get_pos_tag(word):
    """Get POS tag for a word"""
    doc = nlp(word)
    if len(doc) > 0:
        return doc[0].pos_
    return None

def extract_ngrams(text, max_n=3):
    """Extract n-grams (phrases) from text"""
    doc = nlp(text)
    ngrams = []
    
    # Extract noun chunks (natural phrases)
    for chunk in doc.noun_chunks:
        if len(chunk.text.split()) > 1:
            ngrams.append({
                'text': chunk.text,
                'start': chunk.start,
                'end': chunk.end,
                'type': 'noun_chunk'
            })
    
    # Extract compound words and hyphenated terms
    tokens = text.split()
    for i, token in enumerate(tokens):
        # Handle hyphenated words as single units
        if '-' in token and len(token) > 2:
            ngrams.append({
                'text': token,
                'start': i,
                'end': i + 1,
                'type': 'hyphenated'
            })
    
    return ngrams

def suggest_alternatives_for_phrase(phrase, context, top_k=5, similarity_threshold=0.3):
    """Get alternatives for multi-word phrases"""
    # For now, we'll mask the entire phrase and get alternatives
    masked_text = context.replace(phrase, "[MASK]")
    
    inputs = tokenizer(masked_text, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
    
    mask_indices = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]
    
    if len(mask_indices) == 0:
        return []
    
    mask_idx = mask_indices[0]
    probs = torch.nn.functional.softmax(logits[0, mask_idx, :], dim=-1)
    top_indices = torch.topk(probs, top_k * 10).indices.tolist()
    
    candidates = []
    phrase_lower = phrase.lower()
    
    for idx in top_indices:
        decoded = tokenizer.decode([idx]).strip()
        
        if not decoded.isalpha() or len(decoded) < 2 or decoded.lower() == phrase_lower:
            continue
        
        # Semantic similarity
        phrase_embedding = similarity_model.encode([phrase_lower])
        candidate_embedding = similarity_model.encode([decoded.lower()])
        similarity = cosine_similarity(phrase_embedding, candidate_embedding)[0][0]
        
        if similarity >= similarity_threshold:
            candidates.append((decoded, similarity))
        
        if len(candidates) >= top_k:
            break
    
    candidates.sort(key=lambda x: x[1], reverse=True)
    return candidates[:top_k]

def suggest_alternatives(text, top_k=5, similarity_threshold=0.3, analyze_phrases=True):
    if not text.strip():
        return "Please enter some text."

    suggestions = {}
    
    # Get POS tags for all words
    doc = nlp(text)
    word_pos = {token.text: token.pos_ for token in doc}
    
    # Extract and analyze phrases first (if enabled)
    if analyze_phrases:
        ngrams = extract_ngrams(text)
        for ngram in ngrams:
            phrase = ngram['text']
            phrase_alts = suggest_alternatives_for_phrase(phrase, text, top_k, similarity_threshold)
            if phrase_alts:
                suggestions[phrase] = {
                    'alternatives': phrase_alts,
                    'pos': 'PHRASE',
                    'type': ngram['type']
                }
    
    # Now analyze individual words
    # Split by whitespace but preserve original tokens
    words = text.split()
    
    for i, word in enumerate(words):
        # Clean the word but preserve hyphenated words
        cleaned_word = word.strip('.,!?;:"\'')
        
        # Skip if already analyzed as part of phrase
        if cleaned_word in suggestions:
            continue
        
        # Skip pure punctuation or very short words (but include "I", "a")
        if not re.search(r'[a-zA-Z]', cleaned_word) or len(cleaned_word) < 1:
            continue
        
        # Convert to lowercase for analysis (but preserve original for display)
        word_lower = cleaned_word.lower()
        
        # Get POS tag
        original_pos = word_pos.get(cleaned_word)
        if not original_pos:
            # Try lowercase version
            original_pos = get_pos_tag(word_lower)
        
        if not original_pos:
            continue

        # Create masked sentence
        masked_words = words.copy()
        masked_words[i] = "[MASK]"
        masked_text = " ".join(masked_words)
        
        # Tokenize
        inputs = tokenizer(masked_text, return_tensors="pt")
        
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits

        # Find mask position
        mask_indices = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]
        
        if len(mask_indices) == 0:
            continue

        mask_idx = mask_indices[0]
        probs = torch.nn.functional.softmax(logits[0, mask_idx, :], dim=-1)
        top_indices = torch.topk(probs, top_k * 10).indices.tolist()

        # Decode and filter
        candidates = []
        for idx in top_indices:
            decoded = tokenizer.decode([idx]).strip()
            
            # Basic filters (use lowercase comparison)
            if not decoded.isalpha() or len(decoded) < 1 or decoded.lower() == word_lower:
                continue
            
            # POS filter
            candidate_pos = get_pos_tag(decoded)
            if candidate_pos:
                allowed_pos = POS_MAP.get(original_pos, [original_pos])
                if candidate_pos not in allowed_pos:
                    continue
            
            # Semantic similarity filter (compare lowercase versions)
            word_embedding = similarity_model.encode([word_lower])
            candidate_embedding = similarity_model.encode([decoded.lower()])
            similarity = cosine_similarity(word_embedding, candidate_embedding)[0][0]
            
            if similarity >= similarity_threshold:
                candidates.append((decoded, similarity))
            
            if len(candidates) >= top_k:
                break

        if candidates:
            # Sort by similarity score (descending)
            candidates.sort(key=lambda x: x[1], reverse=True)
            suggestions[cleaned_word] = {
                'alternatives': candidates[:top_k],
                'pos': original_pos,
                'type': 'word'
            }

    if not suggestions:
        return "No suggestions found. Try lowering the similarity threshold."

    # Format output
    output = "## 📝 Word & Phrase Alternatives\n\n"
    
    # Separate phrases and words
    phrases = {k: v for k, v in suggestions.items() if v['type'] in ['noun_chunk', 'hyphenated']}
    words_dict = {k: v for k, v in suggestions.items() if v['type'] == 'word'}
    
    if phrases:
        output += "### 🔗 Phrases\n\n"
        for phrase, data in phrases.items():
            output += f"**{phrase}** ({data['type']}) → "
            formatted_alts = [f"{alt} `({score:.2f})`" for alt, score in data['alternatives']]
            output += ", ".join(formatted_alts)
            output += "\n\n"
    
    if words_dict:
        output += "### 📖 Individual Words\n\n"
        for word, data in words_dict.items():
            output += f"**{word}** ({data['pos']}) → "
            formatted_alts = [f"{alt} `({score:.2f})`" for alt, score in data['alternatives']]
            output += ", ".join(formatted_alts)
            output += "\n\n"
    
    return output

demo = gr.Interface(
    fn=suggest_alternatives,
    inputs=[
        gr.Textbox(
            label="Input text", 
            placeholder="Enter a sentence to analyze (all words will be analyzed)",
            lines=3
        ),
        gr.Slider(3, 10, value=5, step=1, label="Number of suggestions per word"),
        gr.Slider(0.0, 1.0, value=0.3, step=0.05, label="Similarity threshold (higher = more similar)"),
        gr.Checkbox(label="Analyze phrases (2-3 word combinations)", value=True)
    ],
    outputs=gr.Markdown(label="Word & Phrase Alternatives"),
    title="🔍 BERT Synonym Tool V2 - Advanced",
    description="""
    **Enhanced features:**
    - ✅ Analyzes ALL words (including "the", "a", "and", etc.)
    - ✅ Handles hyphenated words ("cutting-edge", "data-driven")
    - ✅ Detects multi-word phrases ("machine learning", "digital transformation")
    - ✅ Case-insensitive analysis with preserved capitalization
    - ✅ Semantic similarity scores (0.0-1.0)
    """,
    examples=[
        ["Leveraging cutting-edge artificial intelligence and machine learning paradigms, organizations can harness data-driven insights to drive transformative digital innovation and optimize operational workflows.", 5, 0.3, True],
        ["The quick brown fox jumps over the lazy dog.", 5, 0.4, True],
        ["Machine learning algorithms process data efficiently.", 7, 0.2, True]
    ]
)

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