import json
import numpy as np
import pickle

from keras.models import load_model
from keras.preprocessing.sequence import pad_sequences


model = load_model("multi_task_lstm_ner_srl_model.keras")

with open("word2idx.pkl", "rb") as f:
    word2idx = pickle.load(f)

with open("tag2idx_ner.pkl", "rb") as f:
    tag2idx_ner = pickle.load(f)

with open("tag2idx_srl.pkl", "rb") as f:
    tag2idx_srl = pickle.load(f)


idx2tag_ner = {i: t for t, i in tag2idx_ner.items()}
idx2tag_srl = {i: t for t, i in tag2idx_srl.items()}


max = 50


def predict_sentence(sentence):
    tokens = sentence.strip().lower().split()
    print(tokens)

    x = [word2idx.get(w.lower(), word2idx["UNK"]) for w in tokens]
    x = pad_sequences([x], maxlen=50, padding="post", value=word2idx["PAD"])

    preds = model.predict(x)
    pred_labels_ner = np.argmax(preds[0], axis=-1)[0]
    pred_labels_srl = np.argmax(preds[1], axis=-1)[0]

    result = {
        "tokens": tokens,
        "labels_ner": [
            idx2tag_ner[int(label)] for label in pred_labels_ner[: len(tokens)]
        ],
        "labels_srl": [
            idx2tag_srl[int(label)] for label in pred_labels_srl[: len(tokens)]
        ],
    }

    return result


if __name__ == "__main__":
    try:
        sentence = "sore ini aku pergi ke indonesia"
        print(predict_sentence(sentence))
    except KeyboardInterrupt:
        print("\n\nSelesai.")