# -*- coding: utf-8 -*-
import numpy as np
import streamlit as st
from transformers import AutoModelWithLMHead, PreTrainedTokenizerFast

model_dir = "snoop2head/kogpt-conditional-2"
tokenizer = PreTrainedTokenizerFast.from_pretrained(
    model_dir,
    bos_token="<s>",
    eos_token="</s>",
    unk_token="<unk>",
    pad_token="<pad>",
    mask_token="<mask>",
)


@st.cache
def load_model(model_name):
    model = AutoModelWithLMHead.from_pretrained(model_name)
    return model


model = load_model(model_dir)
print("loaded model completed")


def find_nth(haystack, needle, n):
    start = haystack.find(needle)
    while start >= 0 and n > 1:
        start = haystack.find(needle, start + len(needle))
        n -= 1
    return start


def infer(input_ids, max_length, temperature, top_k, top_p):
    output_sequences = model.generate(
        input_ids=input_ids,
        max_length=max_length,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        do_sample=True,
        num_return_sequences=1,
    )
    return output_sequences


# prompts
st.title("주어진 감정에 맞게 문장을 만드는 KoGPT입니다 🦄")
st.write("좌측에 감정상태의 변화를 주고, CTRL+Enter(CMD+Enter)를 누르세요 🤗")

# text and sidebars
default_value = "수상한 밤들이 계속되던 날 언젠가부터 나는"
sent = st.text_area("Text", default_value, max_chars=30, height=50)
max_length = st.sidebar.slider("생성 문장 길이를 선택해주세요!", min_value=42, max_value=64)
temperature = st.sidebar.slider(
    "Temperature", value=0.9, min_value=0.0, max_value=1.0, step=0.05
)
top_k = st.sidebar.slider("Top-k", min_value=0, max_value=5, value=0)
top_p = st.sidebar.slider("Top-p", min_value=0.0, max_value=1.0, step=0.05, value=1.0)

print("slider sidebars rendering completed")

# make input sentence
emotion_list = ["행복", "놀람", "분노", "혐오", "슬픔", "공포", "중립"]
main_emotion = st.sidebar.radio("주요 감정을 선택하세요", emotion_list)
emotion_list.reverse()
sub_emotion = st.sidebar.radio("두 번째 감정을 선택하세요", emotion_list)

print("radio sidebars rendering completed")

# create condition sentence
random_main_logit = np.random.normal(loc=3.368, scale=1.015, size=1)[0].round(1)
random_sub_logit = np.random.normal(loc=1.333, scale=0.790, size=1)[0].round(1)
condition_sentence = f"{random_main_logit}만큼 {main_emotion}감정인 문장이다. {random_sub_logit}만큼 {sub_emotion}감정인 문장이다. "
condition_plus_input = condition_sentence + sent
print(condition_plus_input)


def infer_sentence(
    condition_plus_input=condition_plus_input, tokenizer=tokenizer, top_k=2
):
    encoded_prompt = tokenizer.encode(
        condition_plus_input, add_special_tokens=False, return_tensors="pt"
    )
    if encoded_prompt.size()[-1] == 0:
        input_ids = None
    else:
        input_ids = encoded_prompt
    output_sequences = infer(input_ids, max_length, temperature, top_k, top_p)
    print(output_sequences)

    generated_sequence = output_sequences[0]
    print(generated_sequence)

    # Decode text
    text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
    print(text)

    # Remove all text after the pad token
    stop_token = tokenizer.pad_token
    print(stop_token)
    text = text[: text.find(stop_token) if stop_token else None]
    print(text)
    
    # Remove condition sentence
    condition_index = find_nth(text, "문장이다", 2)
    text = text[condition_index + 5 :]
    text = text.strip()
    return text


return_text = infer_sentence(
    condition_plus_input=condition_plus_input, tokenizer=tokenizer
)

print(return_text)

st.write(return_text)