import gradio as gr
from gradio.themes.utils import colors
from t5 import T5
from koalpaca import KoAlpaca

LOCAL_TEST = False
MODEL_STRS = ['T5', 'KoAlpaca']
MODELS = []
cur_index = 0

def prepare_theme():
    theme = gr.themes.Default(primary_hue=colors.gray, 
                            secondary_hue=colors.emerald,
                            neutral_hue=colors.emerald).set(
        body_background_fill="*primary_800",
        body_background_fill_dark="*primary_800",
        
        block_background_fill="*primary_700",
        block_background_fill_dark="*primary_700",
        
        border_color_primary="*secondary_300",
        border_color_primary_dark="*secondary_300",
        block_border_width="3px",
        input_border_width="2px",
        
        input_background_fill="*primary_700",
        input_background_fill_dark="*primary_700",
        
        background_fill_secondary="*primary_700",
        background_fill_secondary_dark="*primary_700",
        
        body_text_color="white",
        body_text_color_dark="white",
        
        block_label_text_color="*secondary_300",
        block_label_text_color_dark="*secondary_300",
        block_label_background_fill="*primary_800",
        block_label_background_fill_dark="*primary_800",
        
        color_accent_soft="*primary_600",
        color_accent_soft_dark="*primary_600",
    )
    return theme

def chat(message, chat_history):
    response = MODELS[cur_index].generate(message)
    chat_history.append((message, response))
    return "", gr.Chatbot(chat_history, label=MODEL_STRS[cur_index], bubble_full_width=False)

def change_model_index(idx, chatbot):
    global cur_index
    cur_index = idx
    # print(cur_index)
    return gr.Chatbot(chatbot, label=MODEL_STRS[cur_index], bubble_full_width=False)

if __name__=='__main__':
    theme = prepare_theme()
    
    MODELS.append(T5())
    if not LOCAL_TEST:
        MODELS.append(KoAlpaca())

    with gr.Blocks(theme=theme) as demo:
        with gr.Tab("소개"):
            gr.Markdown(
        """
        # KOMUCHAT : Korean community-style relationship counseling chabot
        KOMUChat: 온라인 커뮤니티 스타일의 연애상담챗봇
        ### 개요
        사람들의 실제 대화를 반영하기 위해 한국 커뮤니티 에펨코리아 + 인스티즈에서 수집한 약 3만건의 질문-답변 대화 데이터셋으로 학습시킨 챗봇<br>
        이별, 데이트코스, 선물 등 연애와 관련된 질문에 대해 온라인 커뮤니티에서 상담하는 것 같이 응답합니다.<br>
        <b>paust-t5-large</b>모델(평균 응답시간 n초)과 <b>polyglot-ko-12.8b-safetensors</b> 모델(평균 응답시간 n초)을 파인튜닝하여 제공합니다.<br> 
        데모에서 선택하여 이용해보실 수 있습니다.

        ### 응답예시
        | Question | Answer |
        | --- | --- |
        | 헤어졌읍니다 | 힘내 더 좋은 인연 나타날거야 |
        | 드라마만 주구장창 보는중 | 뭐 보고 있으신가요 |
        | 결혼까지 생각하는 상태일수록 더 깐깐하게 따지게 돼? | 아무래도 그런 거 같아 |
        | ISTJ 여자가 남자한테 관심있으면 어떻게 해? | 연락 자주 하고 취향 많이 물어보고 자주는 아니어도 꾸준히 만나려 함 |
        | #청자#들아 연애는 갑자기 시작한다고 생각해? | 나는 진자 신기한 게 연애에 대한 기대가 떨어지면 하고 있더랔ㅋㅋㅋ |

        ### 사용기술
        <img src="https://img.shields.io/badge/Python-3776AB?style=for-the-badge&logo=Python&logoColor=white"><img src="https://img.shields.io/badge/PyTorch-EE4C2C?style=for-the-badge&logo=PyTorch&logoColor=white">
        """)
        with gr.Tab("데모"):
            with gr.Row():
                rd = gr.Radio(MODEL_STRS, value='T5', type='index', label='Model Selection', show_label=True, interactive=True)
                with gr.Column(scale=5): # 챗봇 부분 
                    chatbot = gr.Chatbot(label="T5", bubble_full_width=False)
                    with gr.Row():
                        txt = gr.Textbox(show_label=False, placeholder='연애 관련 질문을 입력하세요!', container=False)

        txt.submit(chat, [txt, chatbot], [txt, chatbot])
        rd.select(change_model_index, [rd, chatbot], [chatbot])
    demo.launch(debug=True, share=True)