import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import numpy as np

# Multiple AI text detection models
MODELS = {    
    "DeBERTa Detector": "distilbert-base-uncased-finetuned-sst-2-english",
    "MonkeyDAnh":"MonkeyDAnh/my-awesome-ai-detector-roberta-base-v4-human-vs-machine-finetune",
    "Andreas122001":"andreas122001/roberta-academic-detector",
    "mrm8488":"mrm8488/deberta-v3-small-finetuned-mnli",
      "rdp99":"rdp99/deberta-v3-small-finetuned-mnli",
      "cross_encoder":"cross-encoder/nli-deberta-v3-small"
}
    
def load_model(model_id):
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForSequenceClassification.from_pretrained(model_id)
    return tokenizer, model

def predict(text, tokenizer, model):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    return probs[0].numpy()  # [human_prob, ai_prob]


def verdict(ai_prob):
    """Return a human-readable verdict based on AI probability"""
    if ai_prob < 20:
        return "Most likely human-written."
    elif 20 <= ai_prob < 40:
        return "Possibly human-written with minimal AI assistance."
    elif 40 <= ai_prob < 60:
        return "Unclear – could be either human or AI-assisted."
    elif 60 <= ai_prob < 80:
        return "Possibly AI-generated, or a human using AI assistance."
    else:  # ai_prob >= 80
        return "Likely AI-generated or heavily AI-assisted."

def detect_text(text):
    results = {}
    ai_scores = []

    for name, model_id in MODELS.items():
        try:
            tokenizer, model = load_model(model_id)
            probs = predict(text, tokenizer, model)
            human_score, ai_score = probs
            results[name] = {
                "Human Probability": round(float(human_score) * 100, 2),
                "AI Probability": round(float(ai_score) * 100, 2),
            }
            ai_scores.append(ai_score)
        except Exception as e:
            results[name] = {"error": str(e)}

    # ------------------ Final Score (Average) ------------------
    try:
        ai_scores, human_scores = [], []

        for r in results.values():
            if isinstance(r, dict) and "AI Probability" in r and "Human Probability" in r:
                ai_scores.append(r["AI Probability"])
                human_scores.append(r["Human Probability"])

        if ai_scores and human_scores:
            avg_ai = sum(ai_scores) / len(ai_scores)
            avg_human = sum(human_scores) / len(human_scores)

            results["Final Score"] = {
                # "Human Probability (average)": float(round(avg_human, 2)),
                # "AI Probability (average)": float(round(avg_ai, 2))
                # "Verdict": verdict(avg_ai)
                verdict(avg_ai)
            }
    except Exception as e:
        results["Final Score"] = {"error": str(e)}

    return results


if __name__ == "__main__":
    text = input("Enter text to analyze:\n")
    output = detect_text(text)
    print("\n--- Detection Results ---")
    for model, scores in output.items():
        print(f"\n[{model}]")
        for k, v in scores.items():
            if isinstance(v, (int, float)):  # only add % for numeric values
                print(f"{k}: {v}%")
            else:
                print(f"{k}: {v}")