import pandas as pd
import logging
import dill as dpickle
import numpy as np
import argparse
from ktext.preprocess import processor
from keras.callbacks import CSVLogger, ModelCheckpoint
from keras.models import Model
from keras.layers import Input, GRU, Dense, Embedding, BatchNormalization
from keras import optimizers
from seq2seq_utils import load_decoder_inputs, load_encoder_inputs, load_text_processor
from seq2seq_utils import Seq2Seq_Inference

from sklearn.model_selection import train_test_split
pd.set_option('display.max_colwidth', 500)
logger = logging.getLogger()
logger.setLevel(logging.WARNING)

parser = argparse.ArgumentParser()
parser.add_argument("--input_csv", default="/data/github_issues.csv")
parser.add_argument("--sample_size", type=int, default="1200000")
parser.add_argument("--output_dir", default=".")
args=parser.parse_args()
print(args)

#data_file = '/data/github_issues.csv'
data_file = args.input_csv

use_sample_data=True
if args.sample_size == -1:
    use_sample_data=False


if use_sample_data:
    training_data_size=args.sample_size
    traindf, testdf = train_test_split(pd.read_csv(data_file).sample(n=training_data_size),
                                   test_size=.10)
else:
    traindf, testdf = train_test_split(pd.read_csv(data_file),test_size=.10)


train_body_raw = traindf.body.tolist()
train_title_raw = traindf.issue_title.tolist()


body_pp = processor(keep_n=8000, padding_maxlen=70)
train_body_vecs = body_pp.fit_transform(train_body_raw)

# Instantiate a text processor for the titles, with some different parameters
#  append_indicators = True appends the tokens '_start_' and '_end_' to each
#                      document
#  padding = 'post' means that zero padding is appended to the end of the
#             of the document (as opposed to the default which is 'pre')
title_pp = processor(append_indicators=True, keep_n=4500,
                     padding_maxlen=12, padding ='post')

# process the title data
train_title_vecs = title_pp.fit_transform(train_title_raw)


body_pkl_file = args.output_dir + '/body_pp.dpkl'
train_body_vecs_file = args.output_dir + '/train_body_vecs.npy'
title_pkl_file = args.output_dir + '/title_pp.dpkl'
train_title_vecs_file = args.output_dir + '/train_title_vecs.npy'

# Save the preprocessor
with open(body_pkl_file, 'wb') as f:
    dpickle.dump(body_pp, f)

with open(title_pkl_file, 'wb') as f:
    dpickle.dump(title_pp, f)

# Save the processed data
np.save(train_title_vecs_file, train_title_vecs)
np.save(train_body_vecs_file, train_body_vecs)


encoder_input_data, doc_length = load_encoder_inputs(train_body_vecs_file)
decoder_input_data, decoder_target_data = load_decoder_inputs(train_title_vecs_file)

num_encoder_tokens, body_pp = load_text_processor(body_pkl_file)
num_decoder_tokens, title_pp = load_text_processor(title_pkl_file)



#arbitrarly set latent dimension for embedding and hidden units
latent_dim = 300

##### Define Model Architecture ######

########################
#### Encoder Model ####
encoder_inputs = Input(shape=(doc_length,), name='Encoder-Input')

# Word embeding for encoder (ex: Issue Body)
x = Embedding(num_encoder_tokens, latent_dim, name='Body-Word-Embedding', mask_zero=False)(encoder_inputs)
x = BatchNormalization(name='Encoder-Batchnorm-1')(x)

# We do not need the `encoder_output` just the hidden state.
_, state_h = GRU(latent_dim, return_state=True, name='Encoder-Last-GRU')(x)

# Encapsulate the encoder as a separate entity so we can just
#  encode without decoding if we want to.
encoder_model = Model(inputs=encoder_inputs, outputs=state_h, name='Encoder-Model')

seq2seq_encoder_out = encoder_model(encoder_inputs)

########################
#### Decoder Model ####
decoder_inputs = Input(shape=(None,), name='Decoder-Input')  # for teacher forcing

# Word Embedding For Decoder (ex: Issue Titles)
dec_emb = Embedding(num_decoder_tokens, latent_dim, name='Decoder-Word-Embedding', mask_zero=False)(decoder_inputs)
dec_bn = BatchNormalization(name='Decoder-Batchnorm-1')(dec_emb)

# Set up the decoder, using `decoder_state_input` as initial state.
decoder_gru = GRU(latent_dim, return_state=True, return_sequences=True, name='Decoder-GRU')
decoder_gru_output, _ = decoder_gru(dec_bn, initial_state=seq2seq_encoder_out)
x = BatchNormalization(name='Decoder-Batchnorm-2')(decoder_gru_output)

# Dense layer for prediction
decoder_dense = Dense(num_decoder_tokens, activation='softmax', name='Final-Output-Dense')
decoder_outputs = decoder_dense(x)

########################
#### Seq2Seq Model ####

#seq2seq_decoder_out = decoder_model([decoder_inputs, seq2seq_encoder_out])
seq2seq_Model = Model([encoder_inputs, decoder_inputs], decoder_outputs)


seq2seq_Model.compile(optimizer=optimizers.Nadam(lr=0.001), loss='sparse_categorical_crossentropy')


script_name_base = 'lstm_seq2seq'
csv_logger = CSVLogger('{:}.log'.format(script_name_base))
model_checkpoint = ModelCheckpoint('{:}.epoch{{epoch:02d}}-val{{val_loss:.5f}}.hdf5'.format(script_name_base),
                                   save_best_only=True)

batch_size = 1200
epochs = 7
history = seq2seq_Model.fit([encoder_input_data, decoder_input_data], np.expand_dims(decoder_target_data, -1),
          batch_size=batch_size,
          epochs=epochs,
          validation_split=0.12, callbacks=[csv_logger, model_checkpoint])

seq2seq_Model.save( args.output_dir + '/seq2seq_model.h5')


seq2seq_inf = Seq2Seq_Inference(encoder_preprocessor=body_pp,
                                 decoder_preprocessor=title_pp,
                                 seq2seq_model=seq2seq_Model)

seq2seq_inf.demo_model_predictions(n=50, issue_df=testdf)

#convenience function that generates predictions on holdout set and calculates BLEU Score

bleu_score = seq2seq_inf.evaluate_model(holdout_bodies=testdf.body.tolist(),
                                        holdout_titles=testdf.issue_title.tolist())

print(f"BLEU Score (avg of BLUE 1-4) on Holdout Set: {bleu_score * 100}")



