import streamlit as st
import requests
import time
from ast import literal_eval

def to_md(text):
    return text.replace("\n", "<br />")

@st.cache
def infer(
    prompt, 
    model_name, 
    max_new_tokens=10, 
    temperature=0.0, 
    top_p=1.0,
    top_k=40,
    num_completions=1,
    seed=42,
    stop="\n"
):
    model_name_map = {
        "GPT-JT-6B-v1": "Together-gpt-JT-6B-v1",
    }
    max_new_tokens = int(max_new_tokens)
    num_completions = int(num_completions)
    temperature = float(temperature)
    top_p = float(top_p)
    stop = stop.split(";")
    seed = seed
    
    assert 0 <= max_new_tokens <= 256
    assert 1 <= num_completions <= 5
    assert 0.0 <= temperature <= 10.0
    assert 0.0 <= top_p <= 1.0
    
    if temperature == 0.0:
        temperature = 0.01

    my_post_dict = {
        "type": "general",
        "payload": {
            "max_tokens": int(max_new_tokens),
            "n": int(num_completions),
            "temperature": float(temperature),
            "top_p": float(top_p),
            "model": model_name_map[model_name],
            "prompt": [prompt],
            "request_type": "language-model-inference",
            "stop": stop,
            "best_of": 1,
            "echo": False,
            "seed": int(seed),
            "prompt_embedding": False,
        },
        "returned_payload": {},
        "status": "submitted",
        "source": "dalle",
    }
    
    job_id = requests.post("https://planetd.shift.ml/jobs", json=my_post_dict).json()['id']
    
    for i in range(100):
    
        time.sleep(0.5)
        ret = requests.get(f"https://planetd.shift.ml/job/{job_id}", json={'id': job_id}).json()
        if ret['status'] == 'finished':
            break

    generated_text = ret['returned_payload']['result']['inference_result'][0]['choices'][0]['text']
    
    for stop_word in stop:
        if stop_word in generated_text:
            generated_text = generated_text[:generated_text.find(stop_word)]

    st.session_state.updated = True
    
    return generated_text


def set_preset():
    if st.session_state.preset == "Classification":
        
        if not st.session_state.updated:
            st.session_state.prompt = '''Please classify the given sentence.
Possible labels:
1. <label_0>
2. <label_1>

Input: <sentence_0>
Label: <label_0>

Input: <sentence_1>
Label:'''
        st.session_state.temperature = "0.0"
        st.session_state.top_p = "1.0"
        st.session_state.max_new_tokens = "10"
        
    elif st.session_state.preset == "Generation":
        
        if not st.session_state.updated:
            st.session_state.prompt = '''Please write a story given keywords.

Input: bear, honey
Story: Once upon a time,'''
        st.session_state.temperature = "0.0"
        st.session_state.top_p = "0.9"
        st.session_state.max_new_tokens = "100"
    
    else:
        pass
    
    
def main():

    if 'preset' not in st.session_state:
        st.session_state.preset = "Classification"

    if 'prompt' not in st.session_state:
        st.session_state.prompt = "Please answer the following question:\n\nQuestion: In which country is Zurich located?\nAnswer:"

    if 'temperature' not in st.session_state:
        st.session_state.temperature = "0.0"

    if 'top_p' not in st.session_state:
        st.session_state.top_p = "1.0"

    if 'top_k' not in st.session_state:
        st.session_state.top_k = "40"
        
    if 'max_new_tokens' not in st.session_state:
        st.session_state.max_new_tokens = "10"
        
    if 'updated' not in st.session_state:
        st.session_state.updated = False
        

    st.title("GPT-JT")

    col1, col2 = st.columns([1, 3])

    with col1:
        model_name = st.selectbox("Model", ["GPT-JT-6B-v1"])
        max_new_tokens = st.text_input('Max new tokens', st.session_state.max_new_tokens)
        temperature = st.text_input('temperature', st.session_state.temperature)
        top_k = st.text_input('top_k', st.session_state.top_k)
        top_p = st.text_input('top_p', st.session_state.top_p)
        # num_completions = st.text_input('num_completions (only the best one will be returend)', "1")
        num_completions = "1"
        stop = st.text_input('stop, split by;', r'\n')
        # seed = st.text_input('seed', "42")
        seed = "42"

    with col2:

        # preset = st.radio(
        #     "Recommended Templates", 
        #     ('Classification', 'Generation'), 
        #     on_change=set_preset,
        #     key="preset",
        #     horizontal=True
        # )

        prompt_area = st.empty()
        prompt = prompt_area.text_area(
            "Prompt",
            value=st.session_state.prompt,
            max_chars=4096,
            height=300,
        )

        generated_area = st.empty()
        generated_area.markdown("(Generate here)")

        button_submit = st.button("Submit")

        if button_submit:
            generated_area.markdown("<b>" + to_md(prompt) + "</b>", unsafe_allow_html=True)
            report_text = infer(
                prompt, model_name=model_name, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k,
                num_completions=num_completions, seed=seed, stop=literal_eval("'''"+stop+"'''"),
            )
            generated_area.markdown("<b>" + to_md(prompt) + "</b>" + to_md(report_text), unsafe_allow_html=True)
        
if __name__ == '__main__':
    main()