from flask import Flask, render_template, request
import joblib
import requests
import re
import numpy as np
import torch

from sentence_transformers import SentenceTransformer, util
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from sklearn.metrics.pairwise import cosine_similarity


app = Flask(__name__)

model = joblib.load("hal_model.pkl")

SENTENCE_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
NLI_MODEL = "cross-encoder/nli-MiniLM2-L6-H768"
MAX_LENGTH = 256

TOPIC_SIMILARITY_THRESHOLD = 0.25
NLI_CONTRADICTION_THRESHOLD = 0.50

embedder = SentenceTransformer(SENTENCE_MODEL)

nli_tokenizer = AutoTokenizer.from_pretrained(NLI_MODEL)
nli_model = AutoModelForSequenceClassification.from_pretrained(NLI_MODEL).to("cpu")
nli_model.eval()

def clean_search_query(user_prompt):
    query = user_prompt.strip()

    patterns = [
        r"summarize the main facts about (.+)",
        r"summarize (.+)",
        r"tell me about (.+)",
        r"explain the relationship between (.+)",
        r"explain (.+)",
        r"describe (.+)",
        r"what is (.+)",
        r"who is (.+)",
        r"give me information about (.+)",
        r"write about (.+)"
    ]

    lowered = query.lower()

    for pattern in patterns:
        match = re.search(pattern, lowered)
        if match:
            query = match.group(1)
            break

    query = re.sub(r"[?.!,]+$", "", query)
    query = re.sub(r"^(the|a|an)\s+", "", query, flags=re.IGNORECASE)
    query = re.sub(r"\s+", " ", query).strip()

    return query if query else user_prompt


def get_wikipedia_reference(user_prompt):
    search_url = "https://en.wikipedia.org/w/api.php"

    headers = {
        "User-Agent": "HaluDetect/1.0 student-project (https://huggingface.co/spaces/jr-0/AI/tree/main)"
    }

    search_query = clean_search_query(user_prompt)

    search_params = {
        "action": "query",
        "generator": "search",
        "gsrsearch": search_query,
        "gsrlimit": 10,
        "prop": "extracts|pageprops",
        "exintro": True,
        "explaintext": True,
        "format": "json"
    }

    response = requests.get(
        search_url,
        params=search_params,
        headers=headers,
        timeout=10
    )

    if response.status_code != 200:
        return None, None, search_query

    try:
        data = response.json()
    except ValueError:
        return None, None, search_query

    pages = data.get("query", {}).get("pages", {})

    candidates = []

    for page in pages.values():
        title = page.get("title", "")
        text = page.get("extract", "")
        pageprops = page.get("pageprops", {})

        if not title or not text:
            continue

        if "disambiguation" in pageprops:
            continue

        if "(disambiguation)" in title.lower():
            continue

        candidates.append({
            "title": title,
            "text": text
        })

    if not candidates:
        return None, None, search_query

    normalized_query = search_query.lower().strip()

    for candidate in candidates:
        if candidate["title"].lower() == normalized_query:
            return candidate["title"], candidate["text"], search_query

    query_embedding = embedder.encode(search_query, convert_to_tensor=True)

    best_candidate = None
    best_score = -1

    for candidate in candidates:
        candidate_text = candidate["title"] + ". " + candidate["text"]
        candidate_embedding = embedder.encode(candidate_text, convert_to_tensor=True)

        score = util.cos_sim(query_embedding, candidate_embedding).item()

        if score > best_score:
            best_score = score
            best_candidate = candidate

    return best_candidate["title"], best_candidate["text"], search_query


def get_topic_similarity(reference_text, llm_response):
    reference_embedding = embedder.encode(reference_text, convert_to_tensor=True)
    response_embedding = embedder.encode(llm_response, convert_to_tensor=True)

    similarity = util.cos_sim(reference_embedding, response_embedding).item()

    return similarity


def get_nli_scores(reference_text, response_for_model):
    inputs = nli_tokenizer(
        [reference_text],
        [response_for_model],
        padding=True,
        truncation=True,
        max_length=MAX_LENGTH,
        return_tensors="pt"
    )

    with torch.no_grad():
        outputs = nli_model(**inputs)
        probs = torch.softmax(outputs.logits, dim=1)

    return probs.cpu().numpy()


def get_nli_label_scores(reference_text, llm_response):
    inputs = nli_tokenizer(
        [reference_text],
        [llm_response],
        padding=True,
        truncation=True,
        max_length=MAX_LENGTH,
        return_tensors="pt"
    )

    with torch.no_grad():
        outputs = nli_model(**inputs)
        probs = torch.softmax(outputs.logits, dim=1)[0]

    label_scores = {}

    for index, label in nli_model.config.id2label.items():
        label_scores[label.lower()] = probs[index].item()

    return label_scores


def build_features(reference_text, response_for_model):
    doc_vec = embedder.encode(
        [reference_text],
        convert_to_numpy=True,
        show_progress_bar=False
    )

    response_vec = embedder.encode(
        [response_for_model],
        convert_to_numpy=True,
        show_progress_bar=False
    )

    cosine = cosine_similarity(doc_vec, response_vec).reshape(1, 1)

    nli_scores = get_nli_scores(reference_text, response_for_model)

    features = np.concatenate([
        doc_vec,
        response_vec,
        np.abs(doc_vec - response_vec),
        cosine,
        nli_scores
    ], axis=1)

    return features

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

def get_sentence_nli_summary(reference_text, llm_response):
    sentences = split_into_sentences(llm_response)

    if not sentences:
        return {
            "max_contradiction": 0,
            "min_entailment": 0,
            "weakest_sentence": "",
            "suspicious_sentences": []
        }

    max_contradiction = 0
    min_entailment = 1
    weakest_sentence = ""
    suspicious_sentences = []

    for sentence in sentences:
        scores = get_nli_label_scores(reference_text, sentence)

        contradiction = scores.get("contradiction", 0)
        entailment = scores.get("entailment", 0)
        neutral = scores.get("neutral", 0)

        if contradiction > max_contradiction:
            max_contradiction = contradiction
            weakest_sentence = sentence

        if entailment < min_entailment:
            min_entailment = entailment

        if neutral > 0.65 and entailment < 0.35:
            weakest_sentence = sentence

        if neutral > 0.80 and entailment < 0.05 and contradiction < 0.40:
            suspicious_sentences.append({
                "sentence": sentence,
                "neutral": round(neutral, 3),
                "entailment": round(entailment, 3)
            })

    return {
        "max_contradiction": max_contradiction,
        "min_entailment": min_entailment,
        "weakest_sentence": weakest_sentence,
        "suspicious_sentences": suspicious_sentences
    }


@app.route("/", methods=["GET", "POST"])
def index():
    prediction = None
    confidence = None
    wiki_title = None
    reference_text = None
    search_query = None
    topic_similarity = None
    error = None
    suspicious_sentences = []

    if request.method == "POST":
        user_prompt = request.form.get("user_prompt", "").strip()
        llm_response = request.form.get("llm_response", "").strip()

        if not user_prompt or not llm_response:
            error = "Please enter both a user prompt and an LLM response."
        else:
            try:
                wiki_title, reference_text, search_query = get_wikipedia_reference(user_prompt)

                if not reference_text:
                    error = "Could not find a useful Wikipedia reference for this prompt."
                else:
                    response_for_model = (
                        "Question: " + user_prompt +
                        " Answer: " + llm_response
                    )

                    topic_similarity = get_topic_similarity(reference_text, llm_response)

                    nli_label_scores = get_nli_label_scores(reference_text, llm_response)
                    contradiction_score = nli_label_scores.get("contradiction", 0)
                    entailment_score = nli_label_scores.get("entailment", 0)
                    neutral_score = nli_label_scores.get("neutral", 0)

                    sentence_nli = get_sentence_nli_summary(reference_text, llm_response)

                    sentence_contradiction = sentence_nli["max_contradiction"]
                    sentence_min_entailment = sentence_nli["min_entailment"]
                    weakest_sentence = sentence_nli["weakest_sentence"]
                    suspicious_sentences = sentence_nli["suspicious_sentences"]

                    X = build_features(reference_text, response_for_model)
                    probabilities = model.predict_proba(X)[0]

                    not_hallucinated_prob = probabilities[0]
                    hallucinated_prob = probabilities[1]

                    confidence = round(max(probabilities) * 100, 2)

                    if topic_similarity < TOPIC_SIMILARITY_THRESHOLD:
                        prediction = "Hallucinated"
                        confidence = round((1 - topic_similarity) * 100, 2)

                    elif contradiction_score >= NLI_CONTRADICTION_THRESHOLD:
                        prediction = "Hallucinated"
                        confidence = round(contradiction_score * 100, 2)

                    elif contradiction_score <= 0.20 and topic_similarity >=0.70:
                        prediction = "Not Hallucinated"
                        confidence = round(((1 - contradiction_score) + topic_similarity) / 2 * 100, 2)

                    elif hallucinated_prob >= 0.85:
                        prediction = "Hallucinated"
                        confidence = round(hallucinated_prob * 100, 2)

                    elif not_hallucinated_prob >= 0.65:
                        prediction = "Not Hallucinated"
                        confidence = round(not_hallucinated_prob * 100, 2)

                    else:
                        prediction = "Uncertain"
                        confidence = round(max(probabilities) * 100, 2)

            except Exception as e:
                error = f"Something went wrong: {str(e)}"

    return render_template(
        "index.html",
        prediction=prediction,
        confidence=confidence,
        wiki_title=wiki_title,
        reference_text=reference_text,
        search_query=search_query,
        topic_similarity=topic_similarity,
        suspicious_sentences=suspicious_sentences,
        error=error
    )


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=False)