import gradio as gr
from todset import todset
import numpy as np
from keras.models import Model
from keras.saving import load_model
from keras.layers import *
from keras.preprocessing.text import Tokenizer
import os

os.mkdir("cache")

emb_size = 128
inp_len = 16
maxshift = 4

def hash_str(data: str):
    return hashlib.md5(data.encode('utf-8')).hexdigest()

def train(data: str, message: str):
    if "→" not in data or "\n" not in data:
        return "Dataset example:\nquestion→answer\nquestion→answer\netc."
    dset, responses = todset(data)
    resps_len = len(responses)
    tokenizer = Tokenizer()
    tokenizer.fit_on_texts(list(dset.keys()))
    
    vocab_size = len(tokenizer.word_index) + 1
    if hash_str(data)+".keras" in os.listdir("cache"):
        model = load_model(hash_str(data)+".keras")
    else:
        input_layer = Input(shape=(inp_len,))
        emb_layer = Embedding(input_dim=vocab_size, output_dim=emb_size, input_length=inp_len)(input_layer)
        attn_layer = MultiHeadAttention(num_heads=4, key_dim=128)(emb_layer, emb_layer, emb_layer)
        noise_layer = GaussianNoise(0.1)(attn_layer)
        conv1_layer = Conv1D(64, 8, padding='same', activation='relu', strides=1, input_shape=(64, 128))(noise_layer)
        conv2_layer = Conv1D(16, 4, padding='valid', activation='relu', strides=1)(conv1_layer)
        conv3_layer = Conv1D(8, 2, padding='valid', activation='relu', strides=1)(conv2_layer)
        flatten_layer = Flatten()(conv3_layer)
        attn_flatten_layer = Flatten()(attn_layer)
        conv1_flatten_layer = Flatten()(conv1_layer)
        conv3_flatten_layer = Flatten()(conv3_layer)
        concat1_layer = Concatenate()([flatten_layer, attn_flatten_layer, conv1_flatten_layer, conv2_layer, conv3_flatten_layer])
        dense1_layer = Dense(512, activation="linear")(concat1_layer)
        prelu1_layer = PReLU()(dense1_layer)
        dropout_layer = Dropout(0.3)(prelu1_layer)
        dense2_layer = Dense(256, activation="tanh")(dropout_layer)
        dense3_layer = Dense(256, activation="relu")(dense2_layer)
        dense4_layer = Dense(100, activation="tanh")(dense3_layer)
        concat2_layer = Concatenate()([dense4_layer, prelu1_layer, attn_flatten_layer, conv1_flatten_layer])
        dense4_layer = Dense(resps_len, activation="softmax")(concat2_layer)
        
        X = []
        y = []
        
        for key in dset:
            for p in range(maxshift):
                tokens = tokenizer.texts_to_sequences([key,])[0]
                X.append(np.array(([0,]*p+list(tokens)+[0,]*inp_len)[:inp_len]))
                output_array = np.zeros(resps_len)
                output_array[dset[key]] = 1
                y.append(output_array)
        
        X = np.array(X)
        y = np.array(y)
        
        model.compile(loss="categorical_crossentropy", metrics=["accuracy",])
        
        model.fit(X, y, epochs=10, batch_size=8, workers=4, use_multiprocessing=True)
        model.save("{data_hash}.keras")
    tokens = tokenizer.texts_to_sequences([message,])[0]
    prediction = model.predict(np.array([(list(tokens)+[0,]*inp_len)[:inp_len],]))[0]
    max_o = 0
    max_v = 0
    for ind, i in enumerate(prediction):
        if max_v < i:
            max_v = i
            max_o = ind
    return responses[ind]

iface = gr.Interface(fn=train, inputs=["text", "text"], outputs="text")
iface.launch()