import streamlit as st
from transformers import T5ForConditionalGeneration, T5Tokenizer
from spellchecker import SpellChecker
import re
import torch

# Load model and tokenizer
@st.cache_resource
def load_model():
    model_name = "vennify/t5-base-grammar-correction"
    tokenizer = T5Tokenizer.from_pretrained(model_name)
    model = T5ForConditionalGeneration.from_pretrained(model_name)
    return tokenizer, model

tokenizer, model = load_model()

# Step 1: Fix typos using pyspellchecker
def correct_spelling(text):
    spell = SpellChecker()
    words = re.findall(r'\b\w+\b|\S', text)
    corrected_words = []

    for word in words:
        # Remove non-alphanumeric characters for spellcheck
        clean_word = re.sub(r'[^\w\s]', '', word)
        if clean_word.isalpha():
            corrected_word = spell.correction(clean_word.lower()) or clean_word
            # Restore punctuation
            trailing = ''.join(re.findall(r'[^\w\s]', word))
            corrected_words.append(corrected_word + trailing)
        else:
            corrected_words.append(word)

    return ' '.join(corrected_words)

# Step 2: Grammar Correction using the model
def correct_grammar(text):
    input_text = "gec: " + text
    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 = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return corrected

# UI
st.set_page_config(page_title="Grammar & Spelling Assistant", page_icon="🧠")
st.title("🧠 Grammar & Spelling Correction Assistant")
st.write("Fixes your typos and grammar without changing your meaning.")

user_input = st.text_area("✍️ Enter your sentence below:", height=150)

if st.button("Correct & Explain"):
    if not user_input.strip():
        st.warning("Please enter a sentence.")
    else:
        # Step 1: Spelling correction
        spelling_fixed = correct_spelling(user_input)

        # Step 2: Grammar correction
        final_output = correct_grammar(spelling_fixed)

        st.markdown("### ✅ Final Correction:")
        st.success(final_output)

        st.markdown("### 🔍 Explanation of Changes:")
        st.info(f"""
**Original Sentence:**  
{user_input}

**After Spelling Correction:**  
{spelling_fixed}

**After Grammar Correction:**  
{final_output}

**Explanation:**  
- Typos like `ober` → `over`, `dogz#` → `dogs` were corrected.  
- Grammar and punctuation (like capitalization) were fixed.  
- Your original message and word choices were preserved.
""")