import os
import gradio as gr
import spacy
import nltk
from nltk.corpus import wordnet
from nltk.stem import WordNetLemmatizer
from collections import defaultdict

# Ensure necessary NLTK data is downloaded
nltk.download('wordnet')
nltk.download('averaged_perceptron_tagger')
nltk.download('punkt')

# Ensure the SpaCy model is installed for POS tagging
try:
    nlp = spacy.load("en_core_web_sm")
except OSError:
    subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])
    nlp = spacy.load("en_core_web_sm")

# Initialize lemmatizer
lemmatizer = WordNetLemmatizer()

# Helper function to map nltk POS tags to wordnet POS tags
def get_wordnet_pos(treebank_tag):
    if treebank_tag.startswith('J'):
        return wordnet.ADJ
    elif treebank_tag.startswith('V'):
        return wordnet.VERB
    elif treebank_tag.startswith('N'):
        return wordnet.NOUN
    elif treebank_tag.startswith('R'):
        return wordnet.ADV
    else:
        return None

# Function to correct tense, singular/plural, and verb forms
def grammar_correction(text):
    words = nltk.word_tokenize(text)
    tagged = nltk.pos_tag(words)
    
    corrected_text = []
    for word, tag in tagged:
        wordnet_pos = get_wordnet_pos(tag) or wordnet.NOUN
        lemma = lemmatizer.lemmatize(word, pos=wordnet_pos)
        
        # Apply basic rules for common errors
        if tag.startswith('VB') and word.lower() != lemma:  # Verb tense correction
            corrected_text.append(lemma)
        elif tag.startswith('NNS') and word.lower() == lemma:  # Singular/plural correction
            corrected_text.append(word + 's')
        else:
            corrected_text.append(word)
    
    return ' '.join(corrected_text)

# Gradio app setup
with gr.Blocks() as demo:
    with gr.Tab("Grammar Correction"):
        grammar_input = gr.Textbox(lines=5, label="Input Text")
        grammar_button = gr.Button("Correct Grammar")
        grammar_output = gr.Textbox(label="Corrected Text")

        # Connect the grammar correction function to the button
        grammar_button.click(grammar_correction, inputs=grammar_input, outputs=grammar_output)

# Launch the app
demo.launch()