import gradio as gr
from transformers import T5Tokenizer, T5ForConditionalGeneration
from sentence_transformers import SentenceTransformer, util
import torch

# Load models
tokenizer = T5Tokenizer.from_pretrained("Vamsi/T5_Paraphrase_Paws", use_fast=False)
model = T5ForConditionalGeneration.from_pretrained("Vamsi/T5_Paraphrase_Paws")
similarity_model = SentenceTransformer('all-MiniLM-L6-v2')

# Tone prompt variations
tone_prompts = {
    "Academic": "Rewrite this in a formal and academic way:",
    "Casual": "Rewrite this in a casual and relaxed way:",
    "Friendly": "Make this sound like a friendly human wrote it:",
    "Stealth (AI Detection Bypass)": "Reword this to avoid AI detection and sound natural:"
}

def generate_paraphrase(text, tone):
    prompt = tone_prompts.get(tone, "Paraphrase:")
    input_text = f"{prompt} {text.strip()}"
    input_ids = tokenizer.encode(input_text, return_tensors="pt", max_length=256, truncation=True)
    output_ids = model.generate(
        input_ids,
        max_length=80,
        num_return_sequences=1,
        do_sample=True,
        top_k=120,
        top_p=0.95
    )
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

def humanize_text(input_text, tone):
    if not input_text.strip():
        return "Please enter some text.", "", ""

    # Generate output
    output_text = generate_paraphrase(input_text, tone)

    # Compute semantic similarity
    emb1 = similarity_model.encode(input_text, convert_to_tensor=True)
    emb2 = similarity_model.encode(output_text, convert_to_tensor=True)
    similarity_score = util.pytorch_cos_sim(emb1, emb2).item()
    score_description = "✅ Very Human-Like" if similarity_score < 0.9 else "⚠️ May Still Sound AI-Generated"

    return output_text, f"{similarity_score:.2f}", score_description

# UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🧠 Taha's AI Humanizer Tool")
    gr.Markdown("*Rewriting AI-generated text to sound real, authentic, and undetectable — made by Taha.*")

    with gr.Row():
        input_text = gr.Textbox(lines=6, label="📝 Enter Your AI-Sounding Text")
        output_text = gr.Textbox(lines=6, label="✅ Humanized Output")

    tone = gr.Radio(["Academic", "Casual", "Friendly", "Stealth (AI Detection Bypass)"], label="🎯 Select Tone", value="Stealth (AI Detection Bypass)")

    with gr.Row():
        similarity = gr.Textbox(label="🔍 Semantic Similarity Score")
        score_label = gr.Textbox(label="🧠 Humanization Check")

    gr.Button("🚀 Humanize It").click(fn=humanize_text, inputs=[input_text, tone], outputs=[output_text, similarity, score_label])

demo.launch()