import streamlit as st
from transformers import pipeline
import re

# Load grammar correction model
@st.cache_resource
def load_corrector():
    return pipeline("text2text-generation", model="vennify/t5-base-grammar-correction")

# Load explanation model
@st.cache_resource
def load_explainer():
    return pipeline("text2text-generation", model="google/flan-t5-large")

grammar_corrector = load_corrector()
explainer = load_explainer()

# App UI
st.title("🧠 Advanced Grammar Correction Assistant")
st.write("Enter a sentence. I’ll fix it and explain the correction in detail.")

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

if st.button("Correct & Explain"):
    if user_input.strip() == "":
        st.warning("Please enter a sentence.")
    else:
        # Step 1: Grammar correction
        corrected = grammar_corrector(user_input, max_length=128, do_sample=False)[0]['generated_text']
        
        # Step 2: Explanation prompt
        explanation_prompt = f"Original: {user_input}\nCorrected: {corrected}\nExplain the changes simply with grammar rules and examples."
        explanation = explainer(explanation_prompt, max_length=200, do_sample=False)[0]['generated_text']
        
        # Output
        st.markdown(f"**Correction:** {corrected}")
        st.markdown(f"**Explanation:** {explanation}")