import streamlit as st # type: ignore
from translate_app import tr
import numpy as np
import openai
from dotenv import load_dotenv
import os
import warnings
warnings.filterwarnings('ignore')

title = "Playground ChatGPT"
sidebar_name = "Playground"

def display_proba_next_token(prompt,temperature,max_token):

    # Définir un ou plusieurs motifs d'arrêt
    stop_sequences = [".","\n", "FIN", "Merci"]  # Ajoutez d'autres motifs selon vos besoins
    response = openai.Completion.create(
        engine="davinci-002",
        prompt=prompt,
        max_tokens=max_token,       # Limite à un seul token pour voir uniquement la probabilité du prochain
        logprobs=5,                 # Obtenir les 5 tokens les plus probables
        temperature=temperature,     # Contrôle la créativité, 0 pour des réponses plus prédictibles
        stop=stop_sequences 
    )

    # Mettre logprobs dans un dictionnaire
    if response['choices'][0]['text'] !="":
        logprobs = response['choices'][0]['logprobs']['top_logprobs'][0]
        logprobs_dict = dict(logprobs)

        # Convertir logprobs en probabilités
        probabilities = {key: np.exp(value) for key, value in logprobs_dict.items()}

        # Afficher le texte de la réponse
        st.write("**Token suivant :** ->**:red["+list(probabilities.keys())[0]+"]**<-")  # Afficher le texte de la réponse
        # st.write("Texte de la réponse : ->"+response['choices'][0]['text'][0]+"<-")  # Afficher le texte de la réponse
        st.write("")

        # Afficher les probabilités pour le prochain token
        st.write("**"+tr("Probabilité d'apparition des token:")+"**")
        for i,(token, proba) in enumerate(probabilities.items()):
            st.write(f"&nbsp;&nbsp;&nbsp;&nbsp;{i+1} - Token: ->**:red[{token}]**<- ; proba: {proba:.0%}")

        # Afficher le texte de la réponse complete
        st.write("")
        try:
            text = response['choices'][0]['text']
            st.write("**"+tr("Texte de la réponse complète")+":**")
            st.write("->**:red[" + str(text) + "]**<-")
        except KeyError:
            st.error(tr("La réponse ne contient pas de texte supplémentaire."))
    else:
        st.error(tr("La réponse ne contient pas de texte supplémentaire."))
    # st.write("Texte de la réponse complete: ->"+str(response['choices'][0]['text'])+"<-")  # Afficher le texte de la réponse
    # st.write("")
    # st.write(response)


def run():
    global temperature, max_token
    
    st.write("")
    st.write("")

    st.title(tr(title))
    st.markdown('''
                ---
                ''')

    st.header("**"+tr("Prédiction du token suivant")+"**")
    st.markdown(tr(
        """
        Cet espace a pour objectif d'observer la génération de token par ChatGPT  
        """)
    , unsafe_allow_html=True)
    st.write("")
    if 'OPENAI_API_KEY' in st.session_state:
        try:
            openai.api_key = st.session_state['OPENAI_API_KEY'] 
            col1, col2 = st.columns([3, 1])
            with col2:
                temperature = st.slider(
                    label=tr("Temperature"),
                    min_value=0.0,
                    max_value=1.0,
                    value=0.0  # Valeur initiale
                    )
                max_token = st.slider(
                    label=tr("Max tokens"),
                    min_value=1,
                    max_value=500,
                    value=100 # Valeur initiale
                    ) 
            with col1:
                prompt = st.text_area(label=tr("Prompt:"), value="La souris est mangée par le",height=100)
                display_proba_next_token(prompt,temperature,max_token)
        except Exception as e:
            st.error(f"An unexpected error occurred: {e}")
    else:
        st.write(tr("Saisissez d'abord votre clé API OpenAI !"))