import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# ✅ ตั้งค่าโมเดล
MODEL_NAME = "deepseek-ai/deepseek-llm-7b-chat"
MODEL_DIR = "./deepseek_model"

# ✅ ตรวจสอบและสร้างโฟลเดอร์เก็บโมเดลหากยังไม่มี
if not os.path.exists(MODEL_DIR):
    os.makedirs(MODEL_DIR)

# ✅ โหลด Tokenizer และ Model พร้อมใช้ Half Precision (FP16)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=MODEL_DIR)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,  # Half Precision ลดการใช้ RAM
    device_map="auto",
    offload_folder=MODEL_DIR  # Offload โมเดลบางส่วนไปที่ disk ลดการใช้ RAM
)

def classify_esi(symptoms):
    """
    วิเคราะห์ระดับ ESI ตามอาการของผู้ป่วย
    """
    prompt = (
        "คุณเป็นแพทย์เวชศาสตร์ฉุกเฉินที่ต้องประเมินระดับ Emergency Severity Index (ESI) ตามอาการของผู้ป่วย:\n"
        "- ESI 1: ต้องช่วยชีวิตทันที เช่น หัวใจหยุดเต้น, หยุดหายใจ, ภาวะช็อก\n"
        "- ESI 2: มีภาวะเสี่ยงสูง เช่น สับสน หมดสติ อาการกำเริบรุนแรง\n"
        "- ESI 3: ต้องใช้ทรัพยากรหลายอย่าง เช่น ตรวจเลือด เอกซเรย์\n"
        "- ESI 4: ใช้ทรัพยากรทางการแพทย์เพียงอย่างเดียว เช่น ทำแผล เย็บแผล\n"
        "- ESI 5: ไม่ต้องใช้ทรัพยากรทางการแพทย์ เช่น เป็นหวัด แผลถลอก\n\n"
        f"อาการของผู้ป่วย: {symptoms}\n\n"
        "กรุณาประเมินระดับ ESI ที่เหมาะสมและตอบกลับเป็นตัวเลข 1-5 เท่านั้น หากข้อมูลไม่เพียงพอให้ตอบ 'ไม่สามารถประเมินได้'"
    )

    device = "cuda" if torch.cuda.is_available() else "cpu"
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=10)

    response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

    # ✅ ตรวจสอบการตอบกลับของโมเดล
    for esi in ["1", "2", "3", "4", "5"]:
        if esi in response:
            return int(esi)

    return "ไม่สามารถประเมินได้"  # หากโมเดลตอบกลับไม่ใช่ตัวเลข 1-5