# Import dependencies
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import nltk
from nltk.corpus import wordnet
from gensim.models import KeyedVectors
from nltk.tokenize import word_tokenize

# Download NLTK data (if not already downloaded)
nltk.download('punkt')
nltk.download('stopwords')
nltk.download('wordnet')  # Download WordNet

# Load Word2Vec model from Gensim
word_vectors = KeyedVectors.load_word2vec_format('path/to/GoogleNews-vectors-negative300.bin.gz', binary=True, limit=100000)  # Adjust path as needed

# Check for GPU and set the device accordingly
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load AI Detector model and tokenizer from Hugging Face (DistilBERT)
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english").to(device)

# Function to get synonyms using Gensim Word2Vec
def get_synonyms_gensim(word):
    try:
        synonyms = word_vectors.most_similar(positive=[word], topn=5)
        return [synonym[0] for synonym in synonyms]
    except KeyError:
        return []

# Paraphrasing function using Gensim for synonym replacement
def paraphrase_text(text):
    words = word_tokenize(text)
    paraphrased_words = []
    for word in words:
        synonyms = get_synonyms_gensim(word.lower())
        if synonyms:
            paraphrased_words.append(synonyms[0])
        else:
            paraphrased_words.append(word)
    return ' '.join(paraphrased_words)

# AI detection function using DistilBERT
def detect_ai_generated(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
    with torch.no_grad():
        outputs = model(**inputs)
        probabilities = torch.softmax(outputs.logits, dim=1)
    ai_probability = probabilities[0][1].item()  # Probability of being AI-generated
    return f"AI-Generated Content Probability: {ai_probability:.2f}%"

# Gradio interface definition
with gr.Blocks() as interface:
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(lines=5, label="Input Text")
            detect_button = gr.Button("AI Detection")
            paraphrase_button = gr.Button("Paraphrase Text")
        with gr.Column():
            output_text = gr.Textbox(label="Output")

    detect_button.click(detect_ai_generated, inputs=text_input, outputs=output_text)
    paraphrase_button.click(paraphrase_text, inputs=text_input, outputs=output_text)

# Launch the Gradio app
interface.launch(debug=False)