import json
import requests
from mtranslate import translate
from prompts import PROMPT_LIST
import streamlit as st
import random
import fasttext

headers = {}
LOGO = "huggingwayang.png"
MODELS = {
    "GPT-2 Small": {
        "url": "https://api-inference.huggingface.co/models/flax-community/gpt2-small-indonesian"
    },
    "GPT-2 Medium": {
        "url": "https://api-inference.huggingface.co/models/flax-community/gpt2-medium-indonesian"
    },
}


def get_image(text: str):
    url = "https://wikisearch.uncool.ai/get_image/"
    try:
        payload = {
            "text": text,
            "image_width": 400
        }
        data = json.dumps(payload)
        response = requests.request("POST", url, headers=headers, data=data)
        print(response.content)
        image = json.loads(response.content.decode("utf-8"))["url"]
    except:
        image = ""
    return image

def query(payload, model_name):
    data = json.dumps(payload)
    # print("model url:", MODELS[model_name]["url"])
    response = requests.request("POST", MODELS[model_name]["url"], headers=headers, data=data)
    return json.loads(response.content.decode("utf-8"))


def process(text: str,
            model_name: str,
            max_len: int,
            temp: float,
            top_k: int,
            top_p: float):

    payload = {
        "inputs": text,
        "parameters": {
            "max_new_tokens": max_len,
            "top_k": top_k,
            "top_p": top_p,
            "temperature": temp,
            "repetition_penalty": 2.0,
        },
        "options": {
            "use_cache": True,
        }
    }
    return query(payload, model_name)

st.set_page_config(page_title="Indonesian GPT-2 Demo")
st.title("Indonesian GPT-2")

try:
    token = st.secrets["flax_community_token"]
    headers = {"Authorization": f"Bearer {token}"}
except FileNotFoundError:
    print(f"Token is not found")

ft_model = fasttext.load_model('lid.176.ftz')
# Sidebar
st.sidebar.image(LOGO)
st.sidebar.subheader("Configurable parameters")

max_len = st.sidebar.number_input(
    "Maximum length",
    value=100,
    help="The maximum length of the sequence to be generated."
)

temp = st.sidebar.slider(
    "Temperature",
    value=1.0,
    min_value=0.1,
    max_value=100.0,
    help="The value used to module the next token probabilities."
)

top_k = st.sidebar.number_input(
    "Top k",
    value=10,
    help="The number of highest probability vocabulary tokens to keep for top-k-filtering."
)

top_p = st.sidebar.number_input(
    "Top p",
    value=0.95,
    help=" If set to float < 1, only the most probable tokens with probabilities that add up to top_p or higher are kept for generation."
)

# do_sample = st.sidebar.selectbox('Sampling?', (True, False), help="Whether or not to use sampling; use greedy decoding otherwise.")

st.markdown(
    """
    This demo uses the [small](https://huggingface.co/flax-community/gpt2-small-indonesian) and 
    [medium](https://huggingface.co/flax-community/gpt2-medium-indonesian) Indonesian GPT2 model 
    trained on the Indonesian [Oscar](https://huggingface.co/datasets/oscar), [MC4](https://huggingface.co/datasets/mc4) 
    and [Wikipedia](https://huggingface.co/datasets/wikipedia) dataset. We created it as part of the 
    [Huggingface JAX/Flax event](https://discuss.huggingface.co/t/open-to-the-community-community-week-using-jax-flax-for-nlp-cv/).
    
    The demo supports "multi language" ;-), feel free to try a prompt on your language. We are also experimenting with 
    the sentence based image search using Wikipedia passages encoded with distillbert, and search it using Faiss.  
    """
)

model_name = st.selectbox('Model',(['GPT-2 Small', 'GPT-2 Medium']))

ALL_PROMPTS = list(PROMPT_LIST.keys())+["Custom"]
prompt = st.selectbox('Please choose a predefined prompt or create your custom text.', ALL_PROMPTS, index=len(ALL_PROMPTS)-1)

if prompt == "Custom":
    prompt_box = "Feel free to write text in any language"
else:
    prompt_box = random.choice(PROMPT_LIST[prompt])

text = st.text_area("Enter text", prompt_box)

if st.button("Run"):
    with st.spinner(text="Getting results..."):
        lang_predictions, lang_probability = ft_model.predict(text.replace("\n", " "), k=3)
        # print(f"lang: {lang_predictions}, {lang_probability}")
        if "__label__id" in lang_predictions:
            lang = "id"
        else:
            lang = lang_predictions[0].replace("__label__", "")
            text = translate(text, "id", lang)
        # print(f"{lang}: {text}")
        st.subheader("Result")
        # print(f"maxlen:{max_len}, temp:{temp}, top_k:{top_k}, top_p:{top_p}")
        result = process(text=text,
                         model_name=model_name,
                         max_len=int(max_len),
                         temp=temp,
                         top_k=int(top_k),
                         top_p=float(top_p))

        # print("result:", result)
        if "error" in result:
            if type(result["error"]) is str:
                st.write(f'{result["error"]}.', end=" ")
                if "estimated_time" in result:
                    st.write(f'Please try it again in about {result["estimated_time"]:.0f} seconds')
            else:
                if type(result["error"]) is list:
                    for error in result["error"]:
                        st.write(f'{error}')
        else:
            result = result[0]["generated_text"]
            st.write(result.replace("\n", "  \n"))
            st.text("Translation")
            translation = translate(result, "en", "id")
            if lang == "id":
                st.write(translation.replace("\n", "  \n"))
            else:
                st.write(translate(result, lang, "id").replace("\n", "  \n"))

            image_cat = "https://media.giphy.com/media/vFKqnCdLPNOKc/giphy.gif"
            image = get_image(translation.replace("\"", "'"))
            if image is not "":
                st.image(image, width=400)
            else:
                st.image(image_cat, width=400)