import gradio as gr
import numpy as np
import json
import tensorflow as tf
from tensorflow.keras.models import load_model

# ====================
# Load models
# ====================
encoder_model = load_model("encoder_model.h5")
decoder_model = load_model("decoder_model.h5")

# ====================
# Load assets
# ====================
with open("assets/char2idx.json", "r", encoding="utf-8") as f:
    char2idx = json.load(f)
with open("assets/idx2char.json", "r", encoding="utf-8") as f:
    idx2char = json.load(f)
with open("assets/config.json", "r", encoding="utf-8") as f:
    config = json.load(f)

num_chars = config["num_chars"]
max_len_input = config["max_len_input"]
max_len_target = config["max_len_target"]
latent_dim = config["latent_dim"]

# ====================
# Encode input text
# ====================
def encode_text(text, max_len):
    x = np.zeros((1, max_len, num_chars), dtype="float32")
    words = text.lower().split(" ")
    for t, w in enumerate(words[:max_len]):
        if w in char2idx:
            x[0, t, char2idx[w]] = 1.0
    return x

# ====================
# Decode sequence
# ====================
def decode_sequence(input_seq):
    states_value = encoder_model.predict(input_seq)

    target_seq = np.zeros((1, 1, num_chars))
    target_seq[0, 0, char2idx["\t"]] = 1.0

    decoded_sentence = ""
    stop_condition = False
    while not stop_condition:
        output_tokens, h, c = decoder_model.predict([target_seq] + states_value)
        sampled_token_index = np.argmax(output_tokens[0, -1, :])
        sampled_char = idx2char[str(sampled_token_index)]  # key là str khi lưu json

        if sampled_char == "\n" or len(decoded_sentence.split(" ")) >= max_len_target - 1:
            stop_condition = True
            decoded_sentence += sampled_char
        else:
            decoded_sentence += sampled_char + " "

        target_seq = np.zeros((1, 1, num_chars))
        target_seq[0, 0, sampled_token_index] = 1.0
        states_value = [h, c]

    return decoded_sentence.strip()

# ====================
# Gradio interface
# ====================
def chatbot(user_input, history=[]):
    seq = encode_text(user_input, max_len_input)
    response = decode_sequence(seq)
    history.append((user_input, response))
    return history, history

with gr.Blocks() as demo:
    gr.Markdown("# Basic Chatbot")
    chatbot_ui = gr.Chatbot()
    msg = gr.Textbox(placeholder="Type a message...")
    clear = gr.Button("Clear")

    def respond(message, chat_history):
        response = chatbot(message, chat_history)[1][-1][1]
        return "", chat_history

    msg.submit(respond, [msg, chatbot_ui], [msg, chatbot_ui])
    clear.click(lambda: None, None, chatbot_ui, queue=False)

demo.launch()