# app.py
from sentence_transformers import SentenceTransformer, util
import gradio as gr

# Load semantic similarity model
model = SentenceTransformer('distiluse-base-multilingual-cased-v2')  # Supports Thai + English

# Define bilingual Kratom knowledge base
kratom_contexts = {
    "urination": {
        "keywords": ["urination", "pee", "ปัสสาวะ", "ฉี่"],
        "response": {
            "en": "Kratom may affect urination due to its impact on the autonomic nervous system.",
            "th": "กระท่อมอาจส่งผลต่อการปัสสาวะเนื่องจากมีผลต่อระบบประสาทอัตโนมัติ"
        }
    },
    "skin": {
        "keywords": ["skin", "ผิวหนัง", "darkening", "dry skin", "ผิวแห้ง"],
        "response": {
            "en": "Long-term Kratom use may cause skin darkening and dryness.",
            "th": "การใช้กระท่อมระยะยาวอาจทำให้ผิวคล้ำและแห้ง"
        }
    },
    "behavior": {
        "keywords": ["behavior", "พฤติกรรม", "agitation", "sedation", "กระสับกระส่าย", "ง่วง"],
        "response": {
            "en": "Kratom can cause stimulant or sedative behaviors depending on dosage.",
            "th": "กระท่อมสามารถกระตุ้นหรือกดประสาทได้ขึ้นอยู่กับขนาดที่ใช้"
        }
    }
}

# Semantic matching function
def answer_question(user_input):
    user_embedding = model.encode(user_input, convert_to_tensor=True)
    best_match = None
    best_score = 0.0

    for topic, data in kratom_contexts.items():
        for keyword in data["keywords"]:
            keyword_embedding = model.encode(keyword, convert_to_tensor=True)
            score = util.cos_sim(user_embedding, keyword_embedding).item()
            if score > best_score:
                best_score = score
                best_match = topic

    if best_score > 0.5 and best_match:
        response = kratom_contexts[best_match]["response"]
        return f"🇬🇧 {response['en']}\n🇹🇭 {response['th']}"
    else:
        return "❓ Sorry, I couldn’t match your question. Try asking about urination, skin, or behavior (in Thai or English)."

# Gradio UI
iface = gr.Interface(
    fn=answer_question,
    inputs=gr.Textbox(label="ถามเกี่ยวกับผลของกระท่อม (Thai or English)"),
    outputs=gr.Textbox(label="คำตอบจากผู้ช่วยทางการแพทย์"),
    title="Kratom Semantic Tutor",
    description="Ask about Kratom effects in Thai or English. Uses semantic similarity to match flexible phrasing."
)

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