import keras
import re
import string
import tensorflow
import numpy as np
from tensorflow.keras.preprocessing.sequence import pad_sequences

from keras.layers import *
from keras.models import *
from keras import backend as K
import os
from keras_contrib.layers.crf import CRF, crf_loss, crf_viterbi_accuracy
from keras_contrib import losses
import h5py
import pickle


ceb_stopwords = [
    "ako",
    "amua",
    "ato",
    "busa",
    "ikaw",
    "ila",
    "ilang",
    "imo",
    "imong",
    "iya",
    "iyang",
    "kaayo",
    "kana",
    "kaniya",
    "kaugalingon",
    "kay",
    "kini",
    "kinsa",
    "kita",
    "lamang",
    "mahimong",
    "mga",
    "mismo",
    "nahimo",
    "nga",
    "pareho",
    "pud",
    "sila",
    "siya",
    "unsa"
]


class attention(Layer):
    def __init__(self, name=None, return_sequences=True, **kwargs):
        super(attention, self).__init__(name=name, **kwargs)
        self.return_sequences = return_sequences

        super(attention, self).__init__()

    def build(self, input_shape):
        self.W = self.add_weight(name="att_weight", shape=(input_shape[-1], 1),
                                 initializer="normal")
        self.b = self.add_weight(name="att_bias", shape=(input_shape[1], 1),
                                 initializer="zeros")

        super(attention, self).build(input_shape)

    def call(self, x):
        e = K.tanh(K.dot(x, self.W)+self.b)
        a = K.softmax(e, axis=1)
        output = x*a
        if self.return_sequences:

            return output
        return K.sum(output, axis=1)


vocabs_pkl_filepath = os.path.join(
    os.path.dirname(__file__), 'models', 'vocab_model3.pkl')

pos_h5_filepath = os.path.join(os.path.dirname(
    __file__), 'models', 'POS_TAG_WORD_validated.h5')

with open(vocabs_pkl_filepath, "rb") as f:
    vocab_model = pickle.load(f)

modelx = h5py.File(pos_h5_filepath, 'r')


def POS_TAGGER_MODEL(test_samples):

    model2 = keras.models.load_model(modelx, custom_objects={
                                     'attention': attention, "CRF": CRF, 'crf_loss': crf_loss, 'crf_viterbi_accuracy': crf_viterbi_accuracy}, compile=True)
    model2.compile(optimizer=tf.optimizers.Adam(lr=0.008), loss=losses.crf_loss, metrics=[
                   tf.keras.metrics.Recall(), tf.keras.metrics.AUC(), tf.keras.metrics.Precision(), 'accuracy'])
    # model2.summary()

    test_samples_X = []
    for s in test_samples:
        s_int = []
        for w in s:
            try:
                s_int.append(vocab_model[w.lower()])
            except KeyError:
                s_int.append(vocab_model['<_UNK_>'])
        test_samples_X.append(s_int)

    test_samples_X = pad_sequences(test_samples_X, maxlen=70, padding='post')

    predictions = model2.predict(test_samples_X)

    return predictions


def logits_to_tokens(sequences, index):
    token_sequences = []
    for categorical_sequence in sequences:
        token_sequence = []
        for categorical in categorical_sequence:
            token_sequence.append(index[np.argmax(categorical)])

        token_sequences.append(token_sequence)

    return token_sequences


def remove_stopwords(text):
    return " ".join([word for word in text.split() if word.lower() not in ceb_stopwords])


postagsindex_pkl_filepath = os.path.join(
    os.path.dirname(__file__), 'models', 'tagindex_model2.pkl')

with open(postagsindex_pkl_filepath, "rb") as f:
    tagindex_model = pickle.load(f)


def predict_POS_model():
    test_samples = []

    test_sample = input('INPUT YOUR TEST SENTENCE: ')

    decide = input('DO YOU WANT TO PREPROCESS TEXT? Y/N: ').lower()

    if decide == 'Y'.lower():

        test_sample = test_sample.lower()

        test_sample = remove_stopwords(test_sample)

        test_samples.append(test_sample.split())

    else:

        test_samples.append(test_sample.split())

    max_len = len(test_samples[0])

    predictions = POS_TAGGER_MODEL(test_samples)

    predix = logits_to_tokens(
        predictions, {i: t for t, i in tagindex_model.items()})

    # print('\n')
    # print(test_samples[0])
    # print(predix[0][:max_len])
    sentence = test_samples[0]
    tagged = predix[0][:max_len]
    return sentence, tagged
