import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, T5Tokenizer, T5ForConditionalGeneration

# Model 1: Grammar & spelling correction
grammar_model = "vennify/t5-base-grammar-correction"
tokenizer = AutoTokenizer.from_pretrained(grammar_model)
model = AutoModelForSeq2SeqLM.from_pretrained(grammar_model)
corrector = pipeline("text2text-generation", model=model, tokenizer=tokenizer)

# Model 2: Explanation & suggestions
explain_model = "google/flan-t5-large"
exp_tokenizer = T5Tokenizer.from_pretrained(explain_model)
exp_model = T5ForConditionalGeneration.from_pretrained(explain_model)

def correct_text(text):
    input_text = "gec: " + text
    result = corrector(input_text, max_length=512, clean_up_tokenization_spaces=True)
    return result[0]['generated_text']

def explain_and_suggest(original, corrected):
    prompt = (
        f"Original: {original}\n"
        f"Corrected: {corrected}\n"
        "Explain the grammar, spelling, or style changes made. "
        "Mention any alternate better words. Then suggest how the writer can improve in the future."
    )
    inputs = exp_tokenizer(prompt, return_tensors="pt", truncation=True)
    outputs = exp_model.generate(**inputs, max_length=512)
    return exp_tokenizer.decode(outputs[0], skip_special_tokens=True)

# Streamlit UI
st.set_page_config(page_title="Grammar & Writing Coach", layout="centered")
st.title("📝 Grammar & Writing Coach")
st.markdown("Enter a sentence, paragraph, or essay. We'll correct it, explain the changes, and help you improve.")

text_input = st.text_area("✍️ Your text here:", height=200)

if st.button("🔍 Analyze & Improve"):
    if not text_input.strip():
        st.warning("Please enter some text first.")
    else:
        corrected = correct_text(text_input)
        explanation = explain_and_suggest(text_input, corrected)

        st.subheader("✅ Corrected Version:")
        st.write(corrected)

        st.subheader("📘 Explanation & Suggestions:")
        st.write(explanation)