# app.py
from transformers import T5Tokenizer, T5ForConditionalGeneration
import gradio as gr
import nltk
from nltk.tokenize import sent_tokenize
import difflib

# Download the punkt tokenizer for sentence splitting
nltk.download('punkt')

# Load a pre-trained T5 model specifically fine-tuned for grammar correction
tokenizer = T5Tokenizer.from_pretrained("prithivida/grammar_error_correcter_v1")
model = T5ForConditionalGeneration.from_pretrained("prithivida/grammar_error_correcter_v1")

# Function to perform grammar correction
def grammar_check(text):
    # Split the text into sentences
    sentences = sent_tokenize(text)
    corrected_sentences = []
    original_sentences = []

    for sentence in sentences:
        original_sentences.append(sentence)
        input_text = f"gec: {sentence}"
        input_ids = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
        outputs = model.generate(input_ids, max_length=512, num_beams=4, early_stopping=True)
        corrected_sentence = tokenizer.decode(outputs[0], skip_special_tokens=True)
        corrected_sentences.append(corrected_sentence)

    # Function to underline and color revised parts
    def underline_and_color_revisions(original, corrected):
        diff = difflib.ndiff(original.split(), corrected.split())
        result = []
        for word in diff:
            if word.startswith("+ "):
                result.append(f"<u style='color:red;'>{word[2:]}</u>")
            elif word.startswith("- "):
                continue
            else:
                result.append(word[2:])
        return " ".join(result)

    # Join the corrected sentences back into a single string
    corrected_text = " ".join(
        underline_and_color_revisions(orig, corr) for orig, corr in zip(original_sentences, corrected_sentences)
    )
    return corrected_text

# Create Gradio interface with a writing prompt
interface = gr.Interface(
    fn=grammar_check,
    inputs="text",
    outputs="html",  # Change output type to HTML
    title="Grammar Checker",
    description=(
        "Enter text to check for grammar mistakes.\n\n"
        "Writing Prompt:\n"
        "In the story, Alex and his friends discovered an ancient treasure in Whispering Hollow and decided to donate the artifacts to the local museum.\n\n"
        "In the past, did you have a similar experience where you found something valuable or interesting? Tell the story. Describe what you found, what you did with it, and how you felt about your decision.\n\n"
        "Remember to use past tense in your writing.\n"
        "Sample text for testing: When I was 10, I find an old coin in my backyard. I kept it for a while and shows it to my friends. They was impressed and say it might be valuable. Later, I take it to a local antique shop, and the owner told me it was very old. I decided to give it to the museum in my town. The museum was happy and put it on display. I feel proud of my decision."
    )
)

# Launch the interface
interface.launch()