import os
import json
from greek_normalisation.normalise import Normaliser
import numpy as np
from gensim.models import KeyedVectors
from tensorflow.keras.models import load_model
import tarfile
import gdown


class Morphs:
    """Hold data for one aspect of morphology."""
    def __init__(self, title, tags, lstm1, dnn, lstm2):
        self.title = title
        self.tags = tags
        self.lstm1 = lstm1
        self.lstm2 = lstm2
        self.dnn = dnn
        self.output1 = []
        self.output2 = []
        self.output3 = []
        self.predicted_tags1 = []
        self.predicted_tags2 = []
        self.predicted_tags3 = []
        self.confidence1 = []
        self.confidence2 = []
        self.confidence3 = []


def create_morph_classes():
    """Create a class instance for each part of speech aspect."""
    print('Part-of-speech models loading...')
    pos_lstm1 = load_model(os.path.join(model_folder, 'pos-lstm1-3x128-0.927val0.939-AGDTfirst26last7.h5'))
    pos_dnn = load_model(os.path.join(model_folder, 'pos-dnn-2x20-0.939val0.942-AGDTfirst26last7.h5'))
    pos_lstm2 = load_model(os.path.join(model_folder, 'pos-lstm2-3x128-0.958val0.955-AGDTfirst26last7.h5'))

    print('Person models loading...')
    person_lstm1 = load_model(os.path.join(model_folder, 'person-lstm1-3x128-0.983val0.990-AGDTfirst26last7.h5'))
    person_dnn = load_model(os.path.join(model_folder, 'person-dnn-2x20-0.994val0.992-AGDTfirst26last7.h5'))
    person_lstm2 = load_model(os.path.join(model_folder, 'person-lstm2-3x128-0.994val0.994-AGDTfirst26last7.h5'))

    print('Number models loading...')
    number_lstm1 = load_model(os.path.join(model_folder, 'number-lstm1-3x128-0.955val0.980-AGDTfirst26last7.h5'))
    number_dnn = load_model(os.path.join(model_folder, 'number-dnn-2x20-0.977val0.981-AGDTfirst26last7.h5'))
    number_lstm2 = load_model(os.path.join(model_folder, 'number-lstm2-3x128-0.985val0.987-AGDTfirst26last7.h5'))

    print('Tense models loading...')
    tense_lstm1 = load_model(os.path.join(model_folder, 'tense-lstm1-3x128-0.976val0.990-AGDTfirst26last7.h5'))
    tense_dnn = load_model(os.path.join(model_folder, 'tense-dnn-2x20-0.990val0.992-AGDTfirst26last7.h5'))
    tense_lstm2 = load_model(os.path.join(model_folder, 'tense-lstm2-3x128-0.986val0.992-AGDTfirst26last7.h5'))

    print('Mood models loading...')
    mood_lstm1 = load_model(os.path.join(model_folder, 'mood-lstm1-3x128-0.981val0.992-AGDTfirst26last7.h5'))
    mood_dnn = load_model(os.path.join(model_folder, 'mood-dnn-2x20-0.994val0.992-AGDTfirst26last7.h5'))
    mood_lstm2 = load_model(os.path.join(model_folder, 'mood-lstm2-3x128-0.994val0.995-AGDTfirst26last7.h5'))

    print('Voice models loading...')
    voice_lstm1 = load_model(os.path.join(model_folder, 'voice-lstm1-3x128-0.978val0.991-AGDTfirst26last7.h5'))
    voice_dnn = load_model(os.path.join(model_folder, 'voice-dnn-2x20-0.992val0.993-AGDTfirst26last7.h5'))
    voice_lstm2 = load_model(os.path.join(model_folder, 'voice-lstm2-3x128-0.989val0.993-AGDTfirst26last7.h5'))

    print('Gender models loading...')
    gender_lstm1 = load_model(os.path.join(model_folder, 'gender-lstm1-3x128-0.923val0.934-AGDTfirst26last7.h5'))
    gender_dnn = load_model(os.path.join(model_folder, 'gender-dnn-2x20-0.952val0.937-AGDTfirst26last7.h5'))
    gender_lstm2 = load_model(os.path.join(model_folder, 'gender-lstm2-4x128-0.960val0.958-AGDTfirst26last7.h5'))

    print('Case models loading...')
    case_lstm1 = load_model(os.path.join(model_folder, 'case-lstm1-3x128-0.934val0.962-AGDTfirst26last7.h5'))
    case_dnn = load_model(os.path.join(model_folder, 'case-dnn-2x20-0.957val0.963-AGDTfirst26last7.h5'))
    case_lstm2 = load_model(os.path.join(model_folder, 'case-lstm2-2x128-0.975val0.977-AGDTfirst26last7.h5'))

    print('Degree models loading...')
    degree_lstm1 = load_model(os.path.join(model_folder, 'degree-lstm1-3x128-0.998val0.999-AGDTfirst26last7.h5'))
    degree_dnn = load_model(os.path.join(model_folder, 'degree-dnn-2x20-0.999val0.999-AGDTfirst26last7.h5'))
    degree_lstm2 = load_model(os.path.join(model_folder, 'degree-lstm2-2x128-0.998val0.999-AGDTfirst26last7.h5'))

    # The possible tags for each aspect of morphology
    pos_tags = ('l', 'n', 'a', 'r', 'c', 'i', 'p', 'v', 'd', 'm', 'g', 'u')
    person_tags = ('1', '2', '3')
    number_tags = ('s', 'p', 'd')
    tense_tags = ('p', 'i', 'r', 'l', 't', 'f', 'a')
    mood_tags = ('i', 's', 'n', 'm', 'p', 'o')
    voice_tags = ('a', 'p', 'm', 'e')
    gender_tags = ('m', 'f', 'n')
    case_tags = ('n', 'g', 'd', 'a', 'v')
    degree_tags = ('p', 'c', 's')

    # Create a class instance for each aspect of morphology
    fpos = Morphs('pos', pos_tags, pos_lstm1, pos_dnn, pos_lstm2)
    fperson = Morphs('person', person_tags, person_lstm1, person_dnn, person_lstm2)
    fnumber = Morphs('number', number_tags, number_lstm1, number_dnn, number_lstm2)
    ftense = Morphs('tense', tense_tags, tense_lstm1, tense_dnn, tense_lstm2)
    fmood = Morphs('mood', mood_tags, mood_lstm1, mood_dnn, mood_lstm2)
    fvoice = Morphs('voice', voice_tags, voice_lstm1, voice_dnn, voice_lstm2)
    fgender = Morphs('gender', gender_tags, gender_lstm1, gender_dnn, gender_lstm2)
    fcase = Morphs('case', case_tags, case_lstm1, case_dnn, case_lstm2)
    fdegree = Morphs('degree', degree_tags, degree_lstm1, degree_dnn, degree_lstm2)

    return fpos, fperson, fnumber, ftense, fmood, fvoice, fgender, fcase, fdegree


def elision_normalize(s):
    """Turn unicode characters which look similar to 2019 into 2019."""
    return s.replace("\u02BC", "\u2019").replace("\u1FBF", "\u2019").replace("\u0027", "\u2019").\
        replace("\u1FBD", "\u2019")


def isolate_greek_punctuation(fsentence):
    return fsentence.replace(',', ' , ').replace('·', ' · ').replace(';', ' ; ').replace('.', ' . ').\
        replace('?', ' ? ').replace('»', ' » ').replace('«', ' « ').replace('“', ' “ ').replace('„', ' „ ')


def vector_lookup(gword):
    """Return a vector for a given Greek word."""
    try:
        return wv[gword]
    except KeyError:
        return np.array([0]*100)


def tag(greek_text, annotator='Vanessa Gorman'):
    """Take in a string of Greek text and return that text morphologically tagged."""
    print('Loading models...')
    pos, person, number, tense, mood, voice, gender, case, degree = create_morph_classes()
    morphs = (pos, person, number, tense, mood, voice, gender, case, degree)

    # Load character list, annotator list, and vector dictionary
    print('Loading character list, annotator list, and vector dictionary...')
    with open(os.path.join('data', 'jsons', 'all_norm_characters.json'), encoding='utf-8') as json_file:
        all_norm_characters = json.load(json_file)
    with open(os.path.join('data', 'jsons', 'annotators.json'), encoding='utf-8') as json_file:
        all_annotators = json.load(json_file)

    # Create the normalizer
    normalise = Normaliser().normalise

    # Create annotator tensor
    annotator_tensor = [0] * 37
    try:
        annotator_tensor[all_annotators.index(annotator)] = 1

    # Make Vanessa Gorman the default annotator
    except IndexError:
        annotator_tensor[0] = 1

    print('Pre-processing text...')
    blank_character_tensor = np.array([0]*174, dtype=np.float32)
    punc_separated_text = isolate_greek_punctuation(greek_text)
    split_text = punc_separated_text.split()
    print(f'Text and punctuation split into {len(split_text)} individual tokens.')
    one_hotted_tokens = []
    dnn_input = []
    blank_lstm2_token = np.array([0]*192)
    lstm2_padding = np.tile(blank_lstm2_token, (7, 1))
    lstm2_input = []
    return_list = []

    # Create character tensors and word tensors composed of those character tensors
    for word in split_text:

        # The whole token tensor starts out blank because it's challenging to fill out the empty characters.
        token_tensor = np.array([blank_character_tensor]*21, dtype=np.float32)

        # Normalize each token before tensorizing its characters.
        normalized_form = normalise(elision_normalize(word))[0]
        token_length = len(normalized_form)

        # Create token tensors for tokens longer than 21 characters
        if token_length > 21:
            token_tensor = []
            for character in normalized_form[:10]:
                character_tensor = [0]*137
                try:
                    character_tensor[all_norm_characters.index(character)] = 1
                except ValueError:
                    character_tensor[136] = 1

                # Append the annotator tensor at the end of every character tensor
                character_tensor = character_tensor + annotator_tensor
                character_tensor = np.array(character_tensor, dtype=np.float32)
                token_tensor.append(character_tensor)
            character_tensor = [0]*137
            character_tensor[135] = 1

            # Append the annotator tensor at the end of every character tensor
            character_tensor = character_tensor + annotator_tensor
            character_tensor = np.array(character_tensor, dtype=np.float32)
            token_tensor.append(character_tensor)
            for character in normalized_form[-10:]:
                character_tensor = [0]*137
                try:
                    character_tensor[all_norm_characters.index(character)] = 1
                except ValueError:
                    character_tensor[136] = 1

                # Append the annotator tensor at the end of every character tensor
                character_tensor = character_tensor + annotator_tensor
                character_tensor = np.array(character_tensor, dtype=np.float32)
                token_tensor.append(character_tensor)
            token_tensor = np.array(token_tensor, dtype=np.float32)

        # Create token tensors for tokens shorter than 22 characters
        else:
            for i, character in enumerate(normalized_form):
                character_tensor = [0]*137
                try:
                    character_tensor[all_norm_characters.index(character)] = 1
                except ValueError:
                    character_tensor[136] = 1

                # Append the annotator tensor at the end of every character tensor
                character_tensor = character_tensor + annotator_tensor
                character_tensor = np.array(character_tensor, dtype=np.float32)
                token_tensor[21-token_length+i] = character_tensor

        # Add each tensor token to the samples
        one_hotted_tokens.append(token_tensor)
    one_hots_np = np.array(one_hotted_tokens, dtype=np.float32)

    # Process through the first LSTM...
    print('Making LSTM 1 predictions...')
    for aspect in morphs:
        aspect.output1 = aspect.lstm1.predict(one_hots_np)

    for aspect in morphs:
        for tensor in aspect.output1:
            try:
                aspect.predicted_tags1.append(aspect.tags[int(np.argmax(tensor))])
            except IndexError:
                aspect.predicted_tags1.append('-')
            aspect.confidence1.append(np.amax(tensor))

    for i, token in enumerate(punc_separated_text.split()):
        dnn_input.append(np.concatenate((pos.output1[i], person.output1[i], number.output1[i], tense.output1[i],
                                        mood.output1[i], voice.output1[i], gender.output1[i], case.output1[i],
                                        degree.output1[i], annotator_tensor), axis=0))
    np_dnn_input = np.array(dnn_input)

    # Run outputs through DNN
    print('Making DNN predictions...')
    for aspect in morphs:
        aspect.output2 = aspect.dnn.predict(np_dnn_input)

    for aspect in morphs:
        for tensor in aspect.output2:
            try:
                aspect.predicted_tags2.append(aspect.tags[int(np.argmax(tensor))])
            except IndexError:
                aspect.predicted_tags2.append('-')
            aspect.confidence2.append(np.amax(tensor))

    # Prepare inputs for LSTM2
    for i, token in enumerate(punc_separated_text.split()):
        lstm2_input.append(np.concatenate((pos.output2[i], person.output2[i], number.output2[i], tense.output2[i],
                                           mood.output2[i], voice.output2[i], gender.output2[i], case.output2[i],
                                           degree.output2[i], annotator_tensor,
                                           vector_lookup(normalise(elision_normalize(token))[0])), axis=0))

    padded_lstm2_input = np.concatenate((lstm2_padding, lstm2_input, lstm2_padding))

    time_series = []
    for i in range(0, len(padded_lstm2_input)-14):
        time_series.append(padded_lstm2_input[i:i+15])

    lstm2_ts = np.array(time_series)

    # In the future, convert to a tf.data format:
    # dataset = tf.data.Dataset.from_tensor_slices(padded_lstm2_input).window(15, 1, 1)

    # Run outputs through LSTM2
    print('Making LSTM 2 predictions...')
    for aspect in morphs:
        aspect.output3 = aspect.lstm2.predict(lstm2_ts)

    for aspect in morphs:
        for tensor in aspect.output3:
            try:
                aspect.predicted_tags3.append(aspect.tags[int(np.argmax(tensor))])
            except IndexError:
                aspect.predicted_tags3.append('-')
            aspect.confidence3.append(np.amax(tensor))

    for i, token in enumerate(punc_separated_text.split()):
        return_list.append((token, pos.predicted_tags3[i] + person.predicted_tags3[i] + number.predicted_tags3[i] +
                            tense.predicted_tags3[i] + mood.predicted_tags3[i] + voice.predicted_tags3[i] +
                            gender.predicted_tags3[i] + case.predicted_tags3[i] + degree.predicted_tags3[i]))
    return tuple(return_list)


# This is a check to see which 'home' environmental variable is the correct one.
if 'HOME' in os.environ:
    model_folder = os.path.join(os.environ['HOME'], 'angel_models')
elif 'HOMEPATH' in os.environ:
    model_folder = os.path.join(os.environ['HOMEPATH'], 'angel_models')
else:
    model_folder = 'angel_models'

# See if models have been downloaded. If not, download them.
if os.path.isdir(model_folder) and 'fasttext.wordvectors' in os.listdir(model_folder):
    wv = KeyedVectors.load(os.path.join(model_folder, 'fasttext.wordvectors'))
else:
    # Download the models if they don't exist
    print('Ancient Greek language models need to be downloaded. Once downloaded, the models will be saved locally. '
          'This should only be required once. It may take a minute. These are big files.')
    print('Downloading models...')

    # URL needs to look like https://drive.google.com/uc?id=1-w6Ld2r1Z9kKI2oZ3RWTsmNdJSb12hxP
    url = 'https://drive.google.com/uc?id=1MPTNoRNnTEY818BOdCjRhs7nSG3PgRWW'
    output = 'models.tar.xz'
    gdown.download(url, output, quiet=False)
    if 'HOME' in os.environ:
        extract_folder = os.path.join(os.environ['HOME'], 'angel_models')
    elif 'HOMEPATH' in os.environ:
        extract_folder = os.path.join(os.environ['HOMEPATH'], 'angel_models')
    else:
        # This is a really bad option. You'll have to re-download for every project. I need to look up the most
        # widespread environmental variables
        extract_folder = 'angel_models'
    print(f'Unpacking models to {extract_folder}...')
    tar = tarfile.open(output, 'r:xz')
    tar.extractall(path=extract_folder)
    tar.close()

    # Delete the downloaded file
    os.remove(output)
    with open(os.path.join('data', 'jsons', 'model_folder.json'), 'w') as outfile:
        json.dump(extract_folder, outfile)
    wv = KeyedVectors.load(os.path.join(extract_folder, 'fasttext.wordvectors'))
