# -*- coding: utf-8 -*-
"""
@author: mwahdan
"""

from readers.goo_format_reader import Reader
from vectorizers.trans_vectorizer import TransVectorizer
from vectorizers.tags_vectorizer import TagsVectorizer
from models.trans_auto_model import create_joint_trans_model, load_joint_trans_model
from utils import str2bool

import argparse
from sklearn.preprocessing import LabelEncoder
import numpy as np
import os
import pickle
import tensorflow as tf


# read command-line parameters
parser = argparse.ArgumentParser('Training the Joint Transformer NLU model')
parser.add_argument('--train', '-t', help = 'Path to training data in Goo et al format', type = str, required = True)
parser.add_argument('--val', '-v', help = 'Path to validation data in Goo et al format', type = str, required = True)
parser.add_argument('--save', '-s', help = 'Folder path to save the trained model', type = str, required = True)
parser.add_argument('--epochs', '-e', help = 'Number of epochs', type = int, default = 5, required = False)
parser.add_argument('--batch', '-bs', help = 'Batch size', type = int, default = 64, required = False)
parser.add_argument('--model', '-m', help = 'Path to joint trans NLU model for incremental training', type = str, required = False)
parser.add_argument('--trans', '-tr', help = 'Pretrained transformer model name or path. Is optional. Either --model OR --trans should be provided'
                    , type = str, required = False)
parser.add_argument('--from_pt', '-pt', help = 'Whether the --trans (if provided) is from pytorch or not', 
                    type=str2bool, nargs='?', const=True, default=False, required=False)
parser.add_argument('--cache_dir', '-c', help = 'The cache_dir for transformers library. Is optional', type = str, required = False, default = None)


VALID_TYPES = ['bert', 'albert']

args = parser.parse_args()
train_data_folder_path = args.train
val_data_folder_path = args.val
save_folder_path = args.save
epochs = args.epochs
batch_size = args.batch
start_model_folder_path = args.model
pretrained_model_name_or_path = args.trans#"bert-base-uncased"
from_pt = args.from_pt
cache_dir = args.cache_dir#"D:\\transformers"
if start_model_folder_path is None and pretrained_model_name_or_path is None:
    raise argparse.ArgumentTypeError("Either --model OR --trans should be provided")


max_length = 128

print('read data ...')
train_text_arr, train_tags_arr, train_intents = Reader.read(train_data_folder_path)
val_text_arr, val_tags_arr, val_intents = Reader.read(val_data_folder_path)

print('vectorize data ...')
trans_vectorizer = TransVectorizer(pretrained_model_name_or_path, max_length, cache_dir)
train_input_ids, train_input_mask, train_segment_ids, train_valid_positions, train_sequence_lengths = trans_vectorizer.transform(train_text_arr)
val_input_ids, val_input_mask, val_segment_ids, val_valid_positions, val_sequence_lengths = trans_vectorizer.transform(val_text_arr)


print('vectorize tags ...')
tags_vectorizer = TagsVectorizer()
tags_vectorizer.fit(train_tags_arr)
train_tags = tags_vectorizer.transform(train_tags_arr, train_valid_positions)
val_tags = tags_vectorizer.transform(val_tags_arr, val_valid_positions)
slots_num = len(tags_vectorizer.label_encoder.classes_)

id2label = {i:v for i, v in enumerate(tags_vectorizer.label_encoder.classes_)}


print('encode labels ...')
intents_label_encoder = LabelEncoder()
train_intents = intents_label_encoder.fit_transform(train_intents).astype(np.int32)
val_intents = intents_label_encoder.transform(val_intents).astype(np.int32)
intents_num = len(intents_label_encoder.classes_)


if start_model_folder_path is None or start_model_folder_path == '':
    config = {
        "slots_num": slots_num,
        "intents_num": intents_num,
        "pretrained_model_name_or_path": pretrained_model_name_or_path,
        "cache_dir": cache_dir,
        "from_pt": from_pt,
        "num_bert_fine_tune_layers": 10,
        "intent_loss_weight": 1.0,#0.2,
        "slots_loss_weight": 3.0,#2.0,
        
#        "layer_pruning": {
#            "strategy": "top",
#            "k": 2#6#2
#        }
#        "layer_pruning": {
#            "strategy": "custom",
#            "layers_indexes": [2, 3]
#        }
    }
    model = create_joint_trans_model(config)
else:
    model = load_joint_trans_model(start_model_folder_path)     

print('training model ...')
model.fit([train_input_ids, train_input_mask, train_segment_ids, train_valid_positions], [train_tags, train_intents],
          validation_data=([val_input_ids, val_input_mask, val_segment_ids, val_valid_positions], [val_tags, val_intents]),
          epochs=epochs, batch_size=batch_size, id2label=id2label)

### saving
print('Saving ..')
if not os.path.exists(save_folder_path):
    os.makedirs(save_folder_path)
    print('Folder `%s` created' % save_folder_path)
model.save(save_folder_path)
with open(os.path.join(save_folder_path, 'tags_vectorizer.pkl'), 'wb') as handle:
    pickle.dump(tags_vectorizer, handle, protocol=pickle.HIGHEST_PROTOCOL)
with open(os.path.join(save_folder_path, 'intents_label_encoder.pkl'), 'wb') as handle:
    pickle.dump(intents_label_encoder, handle, protocol=pickle.HIGHEST_PROTOCOL)


tf.compat.v1.reset_default_graph()