import spaces
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

base_model_id = "unsloth/llama-3-8b-Instruct-bnb-4bit"
adapter_model_id = "mutij8/academic-humanizer-llama3"

tokenizer = None
model = None

def load_model_once():
    global tokenizer, model
    if model is not None:
        return tokenizer, model

    print("⏳ ZeroGPU allocated! Loading base model and fine-tuned adapter...")
    tokenizer = AutoTokenizer.from_pretrained(adapter_model_id)
    
    base_model = AutoModelForCausalLM.from_pretrained(
        base_model_id,
        torch_dtype=torch.float16,
        device_map="cuda"
    )
    model = PeftModel.from_pretrained(base_model, adapter_model_id)
    print("✅ Model loaded successfully on GPU!")
    return tokenizer, model

@spaces.GPU(duration=120)
def humanize_and_correct(paragraph):
    paragraph = paragraph.strip()
    if not paragraph:
        return "Please enter a valid academic paragraph."

    tok, mod = load_model_once()
    word_count = len(paragraph.split())

    dual_instruction = f"""You are an expert academic editor and researcher. Your task is to simultaneously fix all grammatical errors, spelling mistakes, and awkward phrasing while humanizing the text.
- Tone: Write with natural human imperfection, varied sentence lengths, and stylistic asymmetry.
- Rule 1: Output MUST be between {int(word_count*0.85)} and {int(word_count*1.15)} words.
- Rule 2: Fix any grammar, punctuation, tense inconsistencies, or syntax errors present in the input text.
- Rule 3: Force clause inversion and structural variety to eliminate AI detection signatures and ensure low plagiarism.
- Rule 4: Keep all inline citations (e.g. [1] or (Smith, 2023)) 100% intact and correctly placed.
- Rule 5: Output ONLY the fully corrected and humanized text."""

    messages = [
        {"role": "system", "content": dual_instruction},
        {"role": "user", "content": paragraph}
    ]

    prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tok([prompt], return_tensors="pt").to("cuda")

    outputs = mod.generate(
        **inputs,
        max_new_tokens=int(word_count * 1.5) + 30,
        do_sample=True,
        temperature=0.85,
        top_p=0.95,
        top_k=50
    )

    input_len = inputs["input_ids"].shape[1]
    generated_text = tok.decode(outputs[0][input_len:], skip_special_tokens=True).strip()
    return generated_text

interface = gr.Interface(
    fn=humanize_and_correct,
    inputs=gr.Textbox(lines=8, placeholder="Paste your text with grammar errors or AI style here..."),
    outputs=gr.Textbox(lines=8, label="Grammar-Corrected & Humanized Output"),
    title="Turnitin-Calibrated Dual-Action Academic Humanizer",
    description="Simultaneously fixes grammatical errors and transforms AI or flawed text into Turnitin-verified human prose."
)

if __name__ == "__main__":
    interface.launch()