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"

# Global variables (loaded lazily on first request)
tokenizer = None
model = None

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

    print("⏳ ZeroGPU allocated! Loading model into VRAM...")
    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
def humanize_paragraph(paragraph):
    paragraph = paragraph.strip()
    if not paragraph:
        return "Please enter a valid academic paragraph."

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

    sft_instruction = f'''You are a human researcher hastily writing a draft for an Essay.
- Tone: Write with natural human imperfection, varied sentence lengths (mixing very short and long sentences), and stylistic asymmetry.
- Rule 1: Output MUST be between {int(word_count*0.85)} and {int(word_count*1.15)} words.
- Rule 2: Force clause inversion. Actively rewrite the paragraph structures by shifting the order of clauses (e.g., start with the consequence, context, or ending arguments rather than the original opening clauses).
- Rule 3: Use highly varied sentence structures and alternative vocabulary choices to avoid contiguous 3-to-5 word matches with the source, preventing plagiarism detection.
- Rule 4: Keep all inline citations (e.g. [1] or (Smith, 2023)) 100% intact.
- Rule 5: Maintain the precise narrative tense of the original text (do not change past tense to present or present tense to past).
- Rule 6: Output ONLY the rewritten text.'''

    messages = [
        {"role": "system", "content": sft_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_paragraph,
    inputs=gr.Textbox(lines=10, placeholder="Enter AI-style academic paragraph here..."),
    outputs=gr.Textbox(lines=10, label="Humanized Paragraph (Low Plagiarism & Low AI Score)"),
    title="Turnitin-Calibrated Essay Humanizer",
    description="Input a paragraph to rewrite it using humanized structures, clause inversion, and exact tense retention."
)

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