import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import re
import pandas as pd
import gradio as gr

MODEL_NAME = "dejanseo/ai-cop"

# --- Load model ---
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16 if (device.type=="cuda" and torch.cuda.is_bf16_supported()) else torch.float32
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, torch_dtype=dtype)
model.to(device).eval()

def sent_tokenize(text):
    return [s for s in re.split(r'(?<=[\.!?])\s+', text.strip()) if s]

# --- Inference function ---
def classify_text(text):
    if not text.strip():
        return "⚠️ Please enter some text.", None, None

    sentences = sent_tokenize(text)
    if not sentences:
        return "⚠️ No sentences detected.", None, None

    inputs = tokenizer(
        sentences,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=model.config.max_position_embeddings
    ).to(device)

    with torch.no_grad():
        logits = model(**inputs).logits
        probs = F.softmax(logits, dim=-1).cpu()
        preds = torch.argmax(probs, dim=-1).cpu()

    results = []
    highlighted_sentences = []
    for i, s in enumerate(sentences):
        p = preds[i].item()
        conf = probs[i, p].item()
        label = "AI" if p == 0 else "Human"

        results.append([s, label, f"{conf:.2f}"])
        if label == "AI":
            highlighted_sentences.append(f"<span style='color:red; font-weight:bold'>{s}</span>")
        else:
            highlighted_sentences.append(f"<span style='color:green; font-weight:bold'>{s}</span>")

    # Overall AI likelihood
    avg = torch.mean(probs, dim=0)
    model_ai = avg[0].item() * 100

    highlighted_text = " ".join(highlighted_sentences)
    df = pd.DataFrame(results, columns=["Sentence", "Classification", "Confidence"])
    return f"⚖️ AI Likelihood: {model_ai:.1f}%", highlighted_text, df

# --- Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 AI Article Detection by Writenix")

    with gr.Row():
        text_input = gr.Textbox(label="Enter text", lines=10, placeholder="Paste your text here…")

    classify_btn = gr.Button("Classify")

    ai_score = gr.Label(label="Overall AI Likelihood")
    highlighted = gr.HTML()
    table = gr.Dataframe(headers=["Sentence", "Classification", "Confidence"], wrap=True)

    classify_btn.click(classify_text, inputs=text_input, outputs=[ai_score, highlighted, table])

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