chatbot_updated.py

import json
import os
import uuid
import tempfile
import numpy as np
from PIL import Image
from deep_translator import GoogleTranslator
from ultralytics import YOLO
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from gtts import gTTS
import whisper as whisper_lib
import re
from rapidfuzz import fuzz

# ==============================
# Load Models
# ==============================
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np

embed_model = SentenceTransformer("all-MiniLM-L6-v2")
yolo_model = YOLO("best_egypt.pt")
whisper_model = whisper_lib.load_model("small")

_model     = None
_tokenizer = None

def load_llm():
    global _model, _tokenizer
    if _model is None:
        print("Loading LLM...")
        model_name = "google/flan-t5-base"
        _tokenizer = AutoTokenizer.from_pretrained(model_name)
        _model     = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        print("LLM loaded.")

# ==============================
# Load Artifacts JSON
# ==============================
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(BASE_DIR, "artifacts.json")

with open(file_path, "r", encoding="utf-8") as f:
    data = json.load(f)

# ==============================
# Build Embedding Index
# ==============================
artifact_texts = []
artifact_docs = []

for doc in data:
    name = doc.get("name", "")
    keywords = " ".join(doc.get("keywords", []))
    text = f"{name} {keywords}"
    artifact_texts.append(text)
    artifact_docs.append(doc)

embeddings = embed_model.encode(artifact_texts, convert_to_numpy=True)
index = faiss.IndexFlatL2(embeddings.shape[1])
index.add(np.array(embeddings))

# ==============================
# Helper Functions
# ==============================
def is_arabic(text):
    return any('\u0600' <= c <= '\u06FF' for c in text)

def translate_to_en(text):
    try:
        return GoogleTranslator(source='auto', target='en').translate(text)
    except:
        return text

def translate_to_ar(text):
    try:
        return GoogleTranslator(source='auto', target='ar').translate(text)
    except:
        return text

# ==============================
# Intent Detection
# ==============================
INTENT_KEYWORDS = {
    "creator":          ["who built", "who made", "creator", "made by",
                         "من بناه", "من صنعه", "المنشئ", "من بنى"],
    "built_year":       ["when was it built", "year", "date",
                         "متى بني", "سنة", "تاريخ"],
    "type":             ["type", "what kind", "نوع", "ما نوع"],
    "era":              ["era", "period", "dynasty",
                         "العصر", "الفترة", "الحقبة", "عصر"],
    "material":         ["material", "made of",
                         "مم صنع", "مصنوع من", "المادة"],
    "description":      ["describe", "appearance", "look like",
                         "وصف", "كيف يبدو", "شكل"],
    "importance":       ["importance", "significance", "why important",
                         "الأهمية", "أهميته", "ليه مهم"],
    "location":         ["where is it now and where was it found", "where ", "what is location", "location", "all locations",
                         "location history", "مكانه القديم والحالي", "فين كان وفين دلوقتي", "مكانه فين دلوقتي واتلاقى فين",
                         "اين ", "مكان"],
    "location_found":   ["where was it found", "discovered",
                         "اين وجد", "مكان اكتشافه", "اكتشف"],
    "current_location": ["where is", "current location", "located",
                         "اين يوجد", "يقع", "مكانه"],
    "summary":          ["tell me about", "overview", "summary", "brief",
                         "what is", "who is", "information",
                         "احكيلي", "أهم المعلومات", "نبذة", "معلومات"],
}

def detect_intents(q_en, q_ar=""):
    q_en = (q_en or "").lower()
    q_ar = (q_ar or "").lower()
    detected_intents = []

    for intent, keywords in INTENT_KEYWORDS.items():
        if intent == "summary":
            continue

        for kw in keywords:
            if kw in q_en or kw in q_ar:
                detected_intents.append(intent)
                break

    if not detected_intents:
        detected_intents.append("summary")

    return detected_intents

# ==============================
# Natural Language Response
# ==============================
def generate_intent_response(artifact, intent, user_lang="en"):
    name_en = artifact.get("name", "This artifact").replace("-", " ").replace("_", " ")
    name = translate_to_ar(name_en) if user_lang == "ar" else name_en

    value = artifact.get(intent, "Unknown")

    if user_lang == "ar":
        creator = translate_to_ar(str(artifact.get('creator', '')))
        era = translate_to_ar(str(artifact.get('era', '')))
        location = translate_to_ar(str(artifact.get('current_location', '')))
        description = translate_to_ar(str(artifact.get('description', '')))
        material = translate_to_ar(str(artifact.get('material', '')))
        value = translate_to_ar(str(value))

        templates = {
            "creator": f"تم إنشاء {name} بواسطة {creator}.",
            "built_year": f"تم بناء {name} في عام {value}.",
            "type": f"{name} هو {value}.",
            "era": f"يرجع {name} إلى عصر {era}.",
            "material": f"{name} مصنوع من {material}.",
            "description": (
                f"يتميز {name} بأنه {description}، "
                f"ويعكس أهمية كبيرة في تاريخ وحضارة مصر القديمة."
            ),
            "importance": f"تكمن أهمية {name} في أنه {value}.",
            "location_found": f"تم اكتشاف {name} في {value}.",
            "location": (
                f"تم اكتشاف {name} في {translate_to_ar(str(artifact.get('location_found', '')))}, "
                f"وهو موجود حاليًا في {translate_to_ar(str(artifact.get('current_location', '')))}."
            ),
            "current_location": f"يوجد {name} حاليًا في {location}.",
            "summary": (
                f"يُعد {name} من أبرز المعالم الأثرية في مصر القديمة، "
                f"حيث يتميز بأنه {description}. "
                f"تم إنشاؤه بواسطة {creator}، ويرجع تاريخه إلى عصر {era}. "
                f"ويقع حاليًا في {location}."
            ),
        }
    else:
        templates = {
            "creator": f"{name} was created by {value}.",
            "built_year": f"{name} was built around {value}.",
            "type": f"{name} is a {value}.",
            "era": f"{name} dates back to the {value}.",
            "material": f"{name} is made of {value}.",
            "description": (
                f"{name} is characterized by {value}, and it represents "
                f"an important part of ancient Egyptian heritage and civilization."
            ),
            "importance": f"The importance of {name}: {value}.",
            "location": (
                f"{name} was discovered in {artifact.get('location_found', '')}, "
                f"and it is currently located in {artifact.get('current_location', '')}."
            ),
            "location_found": f"{name} was discovered in {value}.",
            "current_location": f"{name} is currently located in {value}.",
            "summary": (
                f"{name} is one of the most significant monuments of ancient Egypt. "
                f"It is characterized by {artifact.get('description', '')}. "
                f"It was created by {artifact.get('creator', '')} and dates back to the "
                f"{artifact.get('era', '')} era. "
                f"It is currently located in {artifact.get('current_location', '')}."
            ),
        }

    return templates.get(intent, f"{name}: {value}")


def merge_responses(responses, intents, artifact, lang="en"):
    unique_responses = []
    for resp in responses:
        if resp and resp not in unique_responses:
            unique_responses.append(resp.strip())

    name = artifact.get("name", "This artifact")
    name = name.replace("_", " ").replace("-", " ")
    if lang == "ar":
        name = translate_to_ar(name)

    priority_order = ["current_location", "material", "era", "creator", "type"]

    ordered_responses = []
    for intent in priority_order:
        if intent in intents:
            idx = intents.index(intent)
            if idx < len(unique_responses):
                ordered_responses.append(unique_responses[idx])

    for resp in unique_responses:
        if resp not in ordered_responses:
            ordered_responses.append(resp)

    if not ordered_responses:
        return name

    if lang == "ar":
        connectors = ["كما أنه", "بالإضافة إلى ذلك", "أيضًا"]
        intro = f"يُعد {name} من أبرز المعالم الأثرية في مصر القديمة، حيث "
    else:
        connectors = ["Additionally,", "Moreover,", "Also,"]
        intro = f"{name} is one of the most significant Egyptian monuments. "

    cleaned_responses = []
    for resp in ordered_responses:
        if lang == "ar":
            resp = resp.replace(f"{name} ", "").replace(f"هو {name}", "").strip()
        else:
            resp = resp.replace(f"{name} ", "").strip()
        cleaned_responses.append(resp)

    merged = intro + cleaned_responses[0]
    for i, resp in enumerate(cleaned_responses[1:], start=1):
        connector = connectors[(i - 1) % len(connectors)]
        merged += f" {connector} {resp}"

    if not merged.endswith(("۔", ".", "؟", "?")):
        merged += "."

    return merged

# ==============================
# Smart Artifact Matching (Scoring)
# ==============================
GENERIC_KEYWORDS = {
    "pyramid", "temple", "statue", "tomb", "monument",
    "king", "queen", "pharaoh", "museum", "ancient",
    "تمثال", "معبد", "هرم", "ملك", "ملكة", "فرعون",
}

def normalize(text):
    if not text:
        return ""
    text = text.lower()
    text = text.replace("_", " ").replace("-", " ")
    text = text.replace("of", " ")
    text = re.sub(r"[^\w\s]", "", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text

# ==============================
# Synonyms Dictionary
# ==============================
SYNONYMS = {
    "ramesseum": ["ramessum", "ramessesium", "temple of ramesses", "ramesses temple"],
    "sphinx": ["abu al hol", "great sphinx", "abu el hol"],
    "khafre": ["chephren", "khefren"],
    "khafre pyramid": ["pyramid of khafre", "khafre pyramid", "khafre-pyramid"],
    "hatshepsut": ["mortuary temple of hatshepsut", "hatshepsut temple"],
    "djoser": [
        "joser", "josa", "zoser", "zozer", "djosar",
        "djoser pyramid", "step pyramid", "step pyramid of djoser"
    ]
}

def expand_synonyms(text):
    if not text:
        return ""
    text = normalize(text)
    expanded = text
    for key, values in SYNONYMS.items():
        if key in expanded:
            for v in values:
                expanded += " " + v
        for v in values:
            if v in expanded:
                expanded += " " + key
    return expanded

def find_artifact(q_ar, q_en):
    q_ar_n = normalize(q_ar)
    q_en_n = normalize(q_en)
    query = expand_synonyms(q_en_n + " " + q_ar_n)

    q_vec = embed_model.encode([query])
    D, I = index.search(np.array(q_vec), k=5)
    candidates = [artifact_docs[i] for i in I[0]]

    best_doc = None
    best_score = 0

    for doc in candidates:
        name = normalize(doc.get("name", ""))
        keywords = " ".join(doc.get("keywords", []))
        full = name + " " + normalize(keywords)

        score = 0

        if name in query:
            score += 300

        score += fuzz.ratio(name, query)

        for token in name.split():
            if token in query:
                score += 50

        for kw in doc.get("keywords", []):
            if normalize(kw) in query:
                score += 80

        if score > best_score:
            best_score = score
            best_doc = doc

    if best_score < 70:
        print(f"SEARCH: low confidence ({best_score})")
        return None

    print(f"SEARCH: '{best_doc['name']}' score={best_score}")
    return best_doc

# ==============================
# YOLO
# ==============================
def detect_artifact(image):
    results = yolo_model(image)
    result  = results[0]
    annotated = Image.fromarray(result.plot())
    if result.boxes is None or len(result.boxes) == 0:
        return None, annotated
    best_idx  = int(np.argmax(result.boxes.conf.cpu().numpy()))
    class_id  = int(result.boxes.cls[best_idx])
    return result.names[class_id], annotated

def is_related_to_image(question, detected_name):
    if not detected_name:
        return False
    q = normalize(question)
    pronouns_ar = [
        "هذا", "هذه", "ذلك", "تلك",
        "هذا الاثر", "هذه القطعة", "في الصورة"
    ]
    pronouns_en = [
        "this", "that", "this artifact",
        "this monument", "in the image"
    ]
    for word in pronouns_ar + pronouns_en:
        if word in q:
            return True
    if normalize(detected_name) in q:
        return True
    return False

# ==============================
# STT
# ==============================
AR_PROMPT = (
    "أسماء الآثار المصرية: أبو الهول، رمسيس الثاني، توت عنخ آمون، "
    "خفرع، نفرتيتي، حتشبسوت، أخناتون، الهرم المدرج، هرم خفرع، "
    "هرم منكاورع، معبد الكرنك، معبد فيلة، كوم أمبو، الرامسيوم."
)
EN_PROMPT = (
    "Egyptian artifacts: Sphinx, Ramesses II, Tutankhamun, Khafre, "
    "Nefertiti, Hatshepsut, Akhenaten, Step Pyramid, Pyramid of Khafre, "
    "Pyramid of Menkaure, Karnak Temple, Philae Temple, Kom Ombo, Ramesseum."
)
STT_FIXES = {
    "apple hall": "sphinx", "apple hole": "sphinx",
    "his phoenix": "sphinx", "the phoenix": "sphinx",
    "ابل هول": "أبو الهول", "أبل هول": "أبو الهول",
    "رمسيز": "رمسيس", "خفره": "خفرع", "خفري": "خفرع",
    "توتنخامون": "توت عنخ آمون",
    "josa": "djoser",
    "jose": "djoser",
    "joseph": "djoser",
    "doser": "djoser",
    "dozer": "djoser",
    "joser": "djoser",
    "zoser": "djoser",
    "pyramid josa": "djoser",
    "josa pyramid": "djoser",
}

def correct_stt(text):
    t = text.lower()
    for wrong, right in STT_FIXES.items():
        if wrong.lower() in t:
            text = t.replace(wrong.lower(), right)
            print(f"STT FIX: '{wrong}' -> '{right}'")
            break
    return text.strip()

def speech_to_text(audio_bytes):
    if not audio_bytes:
        return "", "en"
    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
            f.write(audio_bytes)
            tmp_path = f.name

        audio     = whisper_lib.load_audio(tmp_path)
        clip      = whisper_lib.pad_or_trim(audio)
        mel       = whisper_lib.log_mel_spectrogram(clip).to(whisper_model.device)
        _, probs  = whisper_model.detect_language(mel)
        det_lang  = max(probs, key=probs.get)
        conf      = probs[det_lang]
        print(f"STT detected: {det_lang} ({conf:.2f})")

        langs = [det_lang] if conf >= 0.70 else ["ar", "en"]

        res_map = {}
        for lg in langs:
            prompt = AR_PROMPT if lg == "ar" else EN_PROMPT
            r = whisper_model.transcribe(tmp_path, language=lg, fp16=False, initial_prompt=prompt)
            res_map[lg] = r.get("text", "").strip()
            print(f"STT [{lg}]: '{res_map[lg]}'")

        ar_t = res_map.get("ar", "")
        en_t = res_map.get("en", "")

        if ar_t and en_t:
            ar_ratio = sum(1 for c in ar_t if '\u0600' <= c <= '\u06FF') / max(len(ar_t), 1)
            best_text, best_lang = (ar_t, "ar") if ar_ratio > 0.30 else (en_t, "en")
        else:
            best_text, best_lang = (ar_t, "ar") if ar_t else (en_t, "en")

        best_text = correct_stt(best_text)
        print(f"STT RESULT: '{best_text}' lang={best_lang}")
        return best_text, best_lang

    except Exception as e:
        print(f"STT ERROR: {e}")
        return "", "en"
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.remove(tmp_path)

# ==============================
# TTS
# ==============================
def text_to_speech(text):
    try:
        path = os.path.join(tempfile.gettempdir(), f"tts_{uuid.uuid4().hex}.mp3")
        gTTS(text=text, lang="ar" if is_arabic(text) else "en").save(path)
        return path
    except Exception as e:
        print(f"TTS ERROR: {e}")
        return None

def cleanup_audio_file(filepath):
    try:
        if filepath and os.path.exists(filepath):
            os.remove(filepath)
    except:
        pass

# ==============================
# Chatbot Core
# ==============================
def chatbot_updated(question, image=None):
    """
    Main chatbot function that handles:
    - Text or speech input
    - Image artifact detection
    - Multi-intent detection
    - Response generation (template or LLM)
    """

    # =========================
    # Step 1: Handle Text / Speech
    # =========================
    if isinstance(question, bytes):
        text, lang = speech_to_text(question)
    else:
        text = (question or "").strip()
        lang = "ar" if is_arabic(text) else "en"

    if not text:
        return (
            "Please provide a question."
            if lang == "en"
            else "من فضلك اكتب سؤالك."
        )

    q_en = translate_to_en(text) if lang == "ar" else text

    # =========================
    # Step 2: Detect Artifact from Image
    # =========================
    artifact_from_image = None
    detected_name = None

    if image is not None:
        detected_name, _ = detect_artifact(image)
        if detected_name:
            artifact_from_image = find_artifact("", detected_name)
            print(f"IMAGE artifact: {detected_name}")

    # =========================
    # Step 3: Detect Artifact from Text
    # =========================
    artifact_from_text = find_artifact(text, q_en)
    if artifact_from_text:
        print(f"TEXT artifact: {artifact_from_text['name']}")

    # =========================
    # Step 4: Decide Which Artifact to Use
    # =========================
    if artifact_from_image:
        if is_related_to_image(text, detected_name):
            artifact = artifact_from_image
            print("PRIORITY: image (pronoun reference)")
        elif artifact_from_text:
            artifact = artifact_from_text
            print("PRIORITY: text (explicit artifact)")
        else:
            artifact = artifact_from_image
            print("PRIORITY: image (no text artifact)")
    else:
        artifact = artifact_from_text
        print("PRIORITY: text only")

    # =========================
    # Step 5: Fallback if No Artifact Found
    # =========================
    if not artifact:
        return (
            "لم يتم التعرف على الأثر. حاول ذكر اسمه بوضوح."
            if lang == "ar"
            else "Artifact not found. Please mention the artifact name clearly."
        )

    # =========================
    # Step 6: Detect Multiple Intents
    # =========================
    intents = detect_intents(q_en, text)
    print(f"INTENTS: {intents}")

    # =========================
    # Step 7: Generate Responses
    # ✅ FIX: Translate LLM response immediately if lang == "ar"
    # =========================
    responses = []

    for intent in intents:
        if intent in ["description"]:
            load_llm()

            prompt = f"""
You are an expert in Egyptian artifacts.
Use ONLY the following data to answer clearly and concisely.

Name: {artifact.get('name', '')}
Type: {artifact.get('type', '')}
Creator: {artifact.get('creator', '')}
Built Year: {artifact.get('built_year', '')}
Era: {artifact.get('era', '')}
Material: {artifact.get('material', '')}
Description: {artifact.get('description', '')}
Importance: {artifact.get('importance', '')}
Location Found: {artifact.get('location_found', '')}
Current Location: {artifact.get('current_location', '')}
Location: {artifact.get('location', '')}

Question: {q_en}
Answer:
"""
            inputs = _tokenizer(
                prompt,
                return_tensors="pt",
                truncation=True,
                max_length=512
            )

            outputs = _model.generate(
                **inputs,
                max_new_tokens=150,
                do_sample=False,
                no_repeat_ngram_size=3
            )

            resp = _tokenizer.decode(outputs[0], skip_special_tokens=True)

            if "Answer:" in resp:
                resp = resp.split("Answer:")[-1].strip()

            # ✅ الحل: ترجم فوراً لو المستخدم بيتكلم عربي
            if lang == "ar":
                resp = translate_to_ar(resp)

            responses.append(resp)

        else:
            responses.append(
                generate_intent_response(artifact, intent, lang)
            )

    # =========================
    # Step 8: Merge Responses
    # ✅ FIX: Remove the translate_to_ar call here — already handled above
    # =========================
    final_response = merge_responses(responses, intents, artifact, lang)

    return final_response