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 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()

# Optimize for CPU
torch.set_num_threads(4)  # Adjust based on available CPU cores

# 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 suggest_alternatives(text, top_k=5, similarity_threshold=0.3):
    if not text.strip():
        return "Please enter some text."

    words = text.split()
    
    # Get POS tags for all words
    doc = nlp(text)
    word_pos = {token.text: token.pos_ for token in doc}
    
    # Filter valid words and create masked versions
    valid_word_indices = []
    masked_texts = []
    
    for i, word in enumerate(words):
        if not word.isalpha() or len(word) < 2:
            continue
        
        original_pos = word_pos.get(word)
        if not original_pos:
            continue
        
        # Create masked version
        masked_words = words.copy()
        masked_words[i] = "[MASK]"
        masked_text = " ".join(masked_words)
        
        valid_word_indices.append(i)
        masked_texts.append(masked_text)
    
    if not masked_texts:
        return "No valid words found to analyze."
    
    # **BATCHED INFERENCE** - Process all masked sentences at once
    inputs = tokenizer(masked_texts, return_tensors="pt", padding=True, truncation=True)
    
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits  # Shape: [batch_size, seq_len, vocab_size]
    
    # Process results for each word
    suggestions = {}
    
    for batch_idx, word_idx in enumerate(valid_word_indices):
        word = words[word_idx]
        original_pos = word_pos.get(word)
        
        # Find mask position in this batch item
        mask_positions = (inputs["input_ids"][batch_idx] == tokenizer.mask_token_id).nonzero(as_tuple=True)[0]
        
        if len(mask_positions) == 0:
            continue
        
        mask_pos = mask_positions[0]
        
        # Get predictions for this masked position
        probs = torch.nn.functional.softmax(logits[batch_idx, mask_pos, :], dim=-1)
        top_indices = torch.topk(probs, top_k * 10).indices.tolist()
        
        # Decode and filter candidates
        candidates = []
        for idx in top_indices:
            decoded = tokenizer.decode([idx]).strip()
            
            # Basic filters
            if not decoded.isalpha() or len(decoded) < 2 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
            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, float(similarity)))
            
            if len(candidates) >= top_k:
                break
        
        if candidates:
            candidates.sort(key=lambda x: x[1], reverse=True)
            suggestions[word] = candidates[:top_k]
    
    if not suggestions:
        return "No suggestions found. Try lowering the similarity threshold."
    
    # Format output
    output = ""
    for word, alts in suggestions.items():
        pos_tag = word_pos.get(word, 'unknown')
        output += f"**{word}** ({pos_tag}) → "
        formatted_alts = [f"{alt} `({score:.2f})`" for alt, score in alts]
        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"),
        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)")
    ],
    outputs=gr.Markdown(label="Word Alternatives with Similarity Scores"),
    title="BERT Word Alternatives (Optimized)",
    description="🚀 **Batched processing** for faster results! Shows contextually relevant word replacements using BERT-Large with POS filtering and semantic similarity scores.",
    examples=[
        ["Artificial intelligence is transforming global commerce.", 5, 0.3],
        ["The quick brown fox jumps over the lazy dog.", 5, 0.4],
        ["Machine learning algorithms process data efficiently.", 7, 0.2]
    ]
)

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