# -*- coding: utf-8 -*-
"""
Created on Tue Jul 02 17:50:32 2019

@author: uzaymacar

Script containing an example and guideline for training machine translation models
using the implemented @Attention() layer. This script is adapted from
https://www.tensorflow.org/beta/tutorials/text/nmt_with_attention

Usage:
  # Train a Encoder-Decoder model:
  python machine_translation.py
  # Train a Encoder-Decoder model with Global Attention:
  python machine_translation.py --config=1
  # Check CONFIG OPTIONS comment block for more options
"""

import argparse
import time
import os
import unicodedata
import re
import numpy as np
from sklearn.model_selection import train_test_split
import tensorflow as tf
from tensorflow.keras.utils import get_file, to_categorical
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras import Model
from tensorflow.keras.layers import Input, Embedding, Bidirectional, Dense, RepeatVector, \
    TimeDistributed, Flatten, Lambda, Concatenate, Permute, Reshape
from tensorflow.compat.v1.keras.layers import CuDNNLSTM   # CuDNNLSTM not yet released for TF 2.0
from tensorflow.keras.backend import permute_dimensions

import sys
sys.path.append('..')  # add parent directory to Python path for layers.py access
from layers import Attention

# Argument specification
parser = argparse.ArgumentParser()
parser.add_argument("--config",
                    default=0,
                    help="Integer value representing a model configuration")
# CONFIG OPTIONS:
# 0: Encoder-Decoder Model
# 1: Encoder-Decoder Model w/ Global Attention
# 2: Encoder-Decoder Model w/ Local-m Attention
# 3: Encoder-Decoder Model w/ Local-p Attention
args = parser.parse_args()

# Set seeds for reproducibility
np.random.seed(500)
tf.random.set_seed(500)

# Set global constants
embedding_dim = 128     # number of dimensions to represent each character in vector space
batch_size = 100        # feed in the neural network in 100-example training batches
num_epochs = 30         # number of times the neural network goes over EACH training example
config = int(args.config)  # model-configuration

# Load Spanish-to-English dataset (.zip)
zipped = get_file(
    fname='spa-eng.zip',
    origin='http://storage.googleapis.com/download.tensorflow.org/data/spa-eng.zip',
    extract=True
)
file = os.path.join(os.path.dirname(zipped), 'spa-eng/spa.txt')


def unicode_to_ascii(string):
    """Function to convert the string from unicode file to ascii format"""
    return ''.join(char for char in unicodedata.normalize('NFD', string)
                   if unicodedata.category(char) != 'Mn')


def preprocess_sentence(sentence):
    """
    Function to preprocess sentences according to machine translation conventions. Includes
    conversion to ascii characters, general cleaning operations, and removal of accents.
    """
    sentence = unicode_to_ascii(sentence.lower().strip())
    # Creates a space between a word and the punctuation following it, ex: "hi dad." => "hi dad ."
    sentence = re.sub(r'([?.!,¿])', r' \1 ', sentence)
    sentence = re.sub(r'[" "]+', ' ', sentence)

    # Replace everything with space except (a-z, A-Z, '.', '?', '!', ',')
    sentence = re.sub(r'[^a-zA-Z?.!,¿]+', ' ', sentence)

    # Remove spaces
    sentence = sentence.rstrip().strip()

    # Add a start and an end token to the sentence for the model to recognize
    sentence = '<start> ' + sentence + ' <end>'
    return sentence


def create_dataset(path, num_examples):
    """Returns sentence pairs in [ENGLISH, SPANISH] format"""
    lines = open(path, encoding='utf8').read().strip().split('\n')
    word_pairs = [[preprocess_sentence(sentence)
                   for sentence in line.split('\t')]
                  for line in lines[:num_examples]]
    return zip(*word_pairs)


# Load and process pairwise English and Spanish sentences
english_sentences, spanish_sentences = create_dataset(path=file, num_examples=None)


def max_length(tensor):
    """Function that returns the maximum length of any element in a given tensor"""
    return max(len(tensor_unit) for tensor_unit in tensor)


def tokenize(language):
    """Function to tokenize language by mapping words to integer indices"""
    # Perform tokenization
    language_tokenizer = Tokenizer(filters='')
    language_tokenizer.fit_on_texts(language)
    tensor = language_tokenizer.texts_to_sequences(language)
    # Pad sequences to maximum found sequence length by appending 0s to end
    tensor = pad_sequences(sequences=tensor, padding='post')

    return tensor, language_tokenizer


def load_dataset(path, num_examples=None):
    """Function to load dataset"""
    # Create cleaned input-output pairs
    target_language, input_language = create_dataset(path, num_examples)
    # Create language tokenizers and extract tensors
    input_tensor, input_language_tokenizer = tokenize(input_language)
    target_tensor, target_language_tokenizer = tokenize(target_language)

    return input_tensor, target_tensor, input_language_tokenizer, target_language_tokenizer


# Get example (input) tensors, label (target) tensors, and distinct tokenizers for both languages
input_tensor, target_tensor, input_language_tokenizer, target_language_tokenizer = load_dataset(
    path=file, num_examples=None
)

# Setup more global constants
input_vocabulary_size = len(input_language_tokenizer.word_index) + 1
target_vocabulary_size = len(target_language_tokenizer.word_index) + 1
# Calculate maximum sequence lengths of the input and target tensors
input_sequence_length, target_sequence_length = max_length(input_tensor), max_length(target_tensor)

# Split data to training and validation sets
X_train, X_test, Y_train, Y_test = train_test_split(input_tensor,
                                                    target_tensor,
                                                    test_size=0.2,
                                                    random_state=500)

# Compute batch size and cutoff training & validation examples to fit
training_cutoff, test_cutoff = len(X_train) % batch_size, len(X_test) % batch_size
X_train, Y_train = X_train[:-training_cutoff], Y_train[:-training_cutoff]
X_test, Y_test = X_test[:-test_cutoff], Y_test[:-test_cutoff]

# Feed in current labels (Y) as decoder inputs, and pad current labels by 1 word
# Check https://www.tensorflow.org/images/seq2seq/attention_mechanism.jpg for better understanding
X_train_target, X_test_target = Y_train, Y_test
Y_train = np.array(pad_sequences(sequences=np.array([sequence[1:] for sequence in Y_train]),
                                 maxlen=target_sequence_length,
                                 padding='post'))
Y_test = np.array(pad_sequences(sequences=np.array([sequence[1:] for sequence in Y_test]),
                                maxlen=target_sequence_length,
                                padding='post'))

# Create word-level multi-class classification (machine translation), sequence-to-sequence model
# Input Layers
# i)  Initialize input & target sequences
X_input = Input(shape=(input_sequence_length,), batch_size=batch_size, name='input_sequences')
X_target = Input(shape=(target_sequence_length,), batch_size=batch_size, name='target_sequences')
# ii) Initialize hidden & cell states
initial_hidden_state = Input(shape=(128,), batch_size=batch_size, name='hidden_state')
initial_cell_state = Input(shape=(128,), batch_size=batch_size, name='cell_state')
hidden_state, cell_state = initial_hidden_state, initial_cell_state
# NOTE: Here hidden state refers to the recurrently propagated input to the cell, whereas cell
# state refer to the cell state directly from the previous cell.

# Word-Embedding Layers
# i)  Embed input sequences from the input language
embedded_input = Embedding(input_dim=input_vocabulary_size, output_dim=embedding_dim)(X_input)
# ii) Embed target sequences from the target language
embedded_target = Embedding(input_dim=target_vocabulary_size, output_dim=embedding_dim)(X_target)
# NOTE: The embedded target sequences (deriving from X_target) allow us to enforce Teacher Forcing:
# using the actual output (correct translation) from the training dataset at the current time step
# as input in the next time step, rather than the output generated by the network.

# Recurrent Layers
# i)  Encoder
encoder_output = CuDNNLSTM(units=128, return_sequences=True)(embedded_input)
# ii) Decoder
decoder_recurrent_layer = CuDNNLSTM(units=128, return_state=True)
# NOTE: The encoder is always fully vectorized and returns the hidden representations of the whole
# sequence at once, whereas the decoder does this step by step.

# Optional Attention Mechanism
if config == 1:
    attention_layer = Attention(context='many-to-many', alignment_type='global')
elif config == 2:
    attention_layer = Attention(context='many-to-many', alignment_type='local-m')
elif config == 3:
    attention_layer = Attention(context='many-to-many', alignment_type='local-p')

# Prediction Layer
decoder_dense_layer = Dense(units=target_vocabulary_size, activation='softmax')

# Training Loop
outputs = []
for timestep in range(target_sequence_length):
    # Get current input in from embedded target sequences
    current_word = Lambda(lambda x: x[:, timestep: timestep+1, :])(embedded_target)
    # Apply optional attention mechanism
    if config != 0:
        context_vector, attention_weights = attention_layer([encoder_output,
                                                             hidden_state,
                                                             timestep])
    # Combine information
    decoder_input = Concatenate(axis=1)([context_vector if config != 0
                                         else encoder_output, current_word])
    # Decode target word hidden representation at t = timestep
    output, hidden_state, cell_state = decoder_recurrent_layer(
        decoder_input, initial_state=[hidden_state, cell_state]
    )
    # Predict next word & append to outputs
    decoder_outputs = decoder_dense_layer(output)
    outputs.append(decoder_outputs)

# Reshape outputs to (B, S', V)
outputs = Lambda(lambda x: permute_dimensions(tf.stack(x), pattern=(1, 0, 2)))(outputs)

# Compile model
model = Model(inputs=[X_input, X_target, initial_hidden_state, initial_cell_state],
              outputs=outputs)
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
# print(model.summary())
# NOTE: Summary is omitted due to length, deriving from number of simulated recurrent connections

# Create placeholder variables of 0s
placeholder = tf.zeros(shape=(len(X_train), 128))

# Train multi-step, multi-class classification model
model.fit(x={'input_sequences': X_train, 'target_sequences': X_train_target,
             'hidden_state': placeholder, 'cell_state': placeholder},
          y=Y_train,
          validation_data=([X_test, X_test_target,
                            placeholder, placeholder],
                           Y_test),
          epochs=num_epochs, batch_size=batch_size)
