import json
import gradio as gr
import spaces
import wbgtopic
from topic_translator import translate_topics  # 새로 만들 파일입니다

SAMPLE_TEXT = """A growing literature attributes gender inequality in labor market outcomes in part to the reduction in female labor supply after childbirth, the child penalty..."""

clf = wbgtopic.WBGDocTopic()

def process_results(results):
    # 결과를 보기 좋게 가공
    if not results or not results[0]:
        return []
    
    # 첫 번째 문서의 결과만 사용 (단일 문서 기준)
    topics = results[0]
    
    # 상위 5개 주제만 선택
    top_topics = sorted(topics, key=lambda x: x['score_mean'], reverse=True)[:5]
    
    # 각 주제의 점수를 퍼센트로 변환하고 한글로 번역
    formatted_topics = []
    for topic in top_topics:
        formatted_topic = {
            'label': translate_topics.get(topic['label'], topic['label']),
            'score': round(topic['score_mean'] * 100, 1),
            'confidence': round((1 - topic['score_std']) * 100, 1)
        }
        formatted_topics.append(formatted_topic)
    
    return formatted_topics

@spaces.GPU(enable_queue=True, duration=50)
def fn(inputs):
    raw_results = clf.suggest_topics(inputs)
    return process_results(raw_results)

# Gradio 인터페이스 생성
with gr.Blocks(title="문서 주제 분석기") as demo:
    gr.Markdown("## 📚 문서 주제 분석기")
    gr.Markdown("문서를 입력하면 관련된 주제들을 분석하여 보여줍니다.")
    
    with gr.Row():
        text = gr.Textbox(
            value=SAMPLE_TEXT,
            label="분석할 텍스트",
            placeholder="여기에 분석할 텍스트를 입력하세요",
            lines=5
        )
    
    with gr.Row():
        submit_btn = gr.Button("분석 시작", variant="primary")
    
    with gr.Row():
        output = gr.JSON(label="분석 결과")
    
    # 이벤트 연결
    submit_btn.click(
        fn=fn,
        inputs=[text],
        outputs=output
    )

demo.launch(debug=True)