# app.py
import os
from fastai.text.all import load_learner
import gradio as gr

# 모델 파일 경로 설정 (환경 변수 또는 기본값 사용)
MODEL_PATH = os.getenv('MODEL_PATH', 'model.pkl')

try:
    # 모델 로드
    learn = load_learner(MODEL_PATH)
except Exception as e:
    raise RuntimeError(f"모델 로드 오류: {e}")

# 한국어 질병 목록 (모델 출력 순서에 맞춰 설정)
categories = [
    '알레르기', '빈혈', '기관지염', '당뇨병',
    '설사', '피로', '독감', '말라리아', '스트레스'
]

def classify_korean(txt):
    """
    입력된 한국어 증상을 받아 각 질병별 확률을 반환합니다.
    텍스트가 빈 문자열일 경우 확률 0.0을 반환합니다.
    """
    if not txt.strip():
        return {cat: 0.0 for cat in categories}
    # 예측 수행
    _, _, probs = learn.predict(txt)
    # 결과를 딕셔너리 형태로 반환
    return {cat: float(probs[i]) for i, cat in enumerate(categories)}

# Gradio 인터페이스 설정
def create_interface():
    return gr.Interface(
        fn=classify_korean,
        inputs=gr.Textbox(lines=3, label='증상을 자세히 입력하세요'),
        outputs=gr.Label(label='질병별 확률'),
        examples=[
            '목이 아프고 기침이 나요',
            '열이 나고 몸살 기운이 있어요',
            '속이 불편하고 설사를 합니다'
        ],
        title='한국어 의료 진단 확률 예측기',
        description='증상을 입력하면 해당 질병별 확률을 한국어로 보여줍니다.'
    )

if __name__ == '__main__':
    # 인터페이스 실행 (외부 접속 필요 시 share=True 설정)
    create_interface().launch()