import os
import tensorflow as tf
from tensorflow.keras.backend import set_session
from tf_crf_layer.loss import ConditionalRandomFieldLoss
from tf_crf_layer.metrics import SequenceCorrectness
# tf.enable_eager_execution()
from mtnlpmodel.utils.input_process_util import _read_configure
from mtnlpmodel.utils.ctrldir_util import (create_dir_if_needed,
                                           create_file_dir_if_needed,
                                           create_or_rm_dir_if_needed)
from mtnlpmodel.utils.deliverablemodel_util import (ConverterForVector,
                                                    ConverterForMTRequest,
                                                    ConverterForMTResponse_VirtualPad,
                                                    mtinput_export_as_deliverable_model, 
                                                    mtinput_export_as_vector_deliverable_model)


def main():
    # get configure
    config = _read_configure("./configure.yaml")

    # get Parameters (controller)
    EPOCHS = config.get("epochs", 10)
    PRETRAIN_EPOCHS =config.get("pretrain_epochs", 5)
    BATCHSIZE = config.get("batch_size", 32)
    PRETRAIN_BATCHSIZE = config.get("pretrain_batchsize", 32)
    LEARNINGRATE = config.get("learning_rate", 0.001)
    MAX_SENTENCE_LEN = config.get("max_sentence_len", 25)
    LRDECAY = config.get('lr_decay', False)
    EARLYSTOP = config.get('early_stop', False)

    # get Parameters (model select)
    MODEL_CHOICE = config.get("model_choice", "VIRTUAL_EMBEDDING")

    # get Parameters (model structure)
    CLS2NER_KEYWORD_LEN = config.get("cls2ner_keyword_len", 5)
    EMBED_DIM = config.get("embedding_dim", 128)
    ARCLOSS = config.get("Arcloss", True)
    USE_ATTENTION_LAYER = config.get("use_attention_layer", False)
    BiLSTM_STACK_CONFIG = config.get("bilstm_stack_config", [])
    CRF_PARAMS = config.get("crf_params", {})
    DECODER_STYLE = config.get("decoder_style", "BILUO")


    # set gpu occupancy rate
    gpu_cfg = tf.ConfigProto()
    gpu_cfg.gpu_options.allocator_type = 'BFC'  # A "Best-fit with coalescing" algorithm, simplified from a version of dlmalloc.
    gpu_cfg.gpu_options.per_process_gpu_memory_fraction = 0.5
    gpu_cfg.gpu_options.allow_growth = True
    set_session(tf.Session(config=gpu_cfg))

    # get preprocessed input data dict
    from mtnlpmodel.utils.input_process_util import input_data_process
    # to build a fixed training environment, input data should be fixed.
    # input_data should be shuffled and remove duplication outside the trainer before running the program.
    # input data should be corpus(no duplication, shuffle well)

    if MODEL_CHOICE=='VIRTUAL_EMBEDDING' or MODEL_CHOICE=='CLS2NER_INPUT':  # different model structures have different input process way
        data_dict = input_data_process(config, **{'MAX_SENTENCE_LEN': MAX_SENTENCE_LEN,      # preprocess the input_data
                                       'CLS2NER_KEYWORD_LEN': CLS2NER_KEYWORD_LEN,})
    else:
        data_dict = input_data_process(config, **{'MAX_SENTENCE_LEN': MAX_SENTENCE_LEN,  # preprocess the input_data
                                                  'CLS2NER_KEYWORD_LEN': 0, })
        PRETRAIN_EPOCHS = 0

    # get lookupers
    ner_tag_lookuper = data_dict['ner_tag_lookuper']
    cls_label_lookuper = data_dict['cls_label_lookuper']
    vocabulary_lookuper = data_dict['vocabulary_lookuper']

    # get train/test data for training model
    ner_train_x, ner_train_y = data_dict['ner_train_x'], data_dict['ner_train_y']
    ner_test_x, ner_test_y = data_dict['ner_test_x'], data_dict['ner_test_y']

    cls_train_x, cls_train_y = data_dict['cls_train_x'], data_dict['cls_train_y']
    cls_test_x, cls_test_y = data_dict['cls_test_x'], data_dict['cls_test_y']


    # build model or finetuning
    from mtnlpmodel.core import build_model_multi_input
    params = {'EMBED_DIM': EMBED_DIM,
              'PRETRAIN_EPOCHS': PRETRAIN_EPOCHS,
              'BiLSTM_STACK_CONFIG': BiLSTM_STACK_CONFIG,
              'MAX_SENTENCE_LEN': MAX_SENTENCE_LEN,
              'CLS2NER_KEYWORD_LEN': CLS2NER_KEYWORD_LEN,
              'USE_ATTENTION_LAYER': USE_ATTENTION_LAYER,
              'Arcloss': ARCLOSS,
              'ner_tag_lookuper': ner_tag_lookuper,
              'cls_label_lookuper': cls_label_lookuper,
              'vocabulary_lookuper': vocabulary_lookuper,
              'CRF_PARAMS': CRF_PARAMS
             }
    model_choice = MODEL_CHOICE   # VIRTUAL_EMBEDDING, CLS2NER_INPUT, OTHER
    print("Model structure choosing {}".format(model_choice))

    # train the model by random initializer(make a fresh start to train a model)
    model, semantic_vector = build_model_multi_input(model_choice, **params)     # to build the model and get cls_vector
    
    model.summary()


    # build callbacks list
    callbacks_list = []

    tensorboard_callback = tf.keras.callbacks.TensorBoard(
        #log_dir=create_dir_if_needed(config["summary_log_dir"]),
        log_dir='.\\results\\summary_log_dir',  # for windows
        batch_size=BATCHSIZE,
    )
    callbacks_list.append(tensorboard_callback)

    checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
        os.path.join(create_dir_if_needed(config["model_dir"]), "cp-{epoch:04d}.ckpt"),
        load_weights_on_restart=True,
        verbose=1,
    )
    callbacks_list.append(checkpoint_callback)
    
    # early stop util
    if EARLYSTOP:
        early_stop = tf.keras.callbacks.EarlyStopping(monitor='val_loss',  # early stop index
                                                      patience=3,          # early stop delay epoch
                                                      verbose=2,           # display mode
                                                      mode='auto')
        callbacks_list.append(early_stop)

    #learning rate decay util
    if LRDECAY:
        reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.75, patience=3, verbose=1,
                                                         mode='auto', epsilon=0.0001, cooldown=0, min_lr=0.00001)
        callbacks_list.append(reduce_lr)
    
    # ner_loss_func
    ner_loss_func = ConditionalRandomFieldLoss()
    
    # set optimizer
    adam_optimizer = tf.keras.optimizers.Adam(learning_rate=LEARNINGRATE, beta_1=0.9, beta_2=0.999, amsgrad=True)

    NER_out_name = 'crf'
    CLS_out_name = 'cls'

    # pretrain model -> train cls branch
    model.compile(optimizer=adam_optimizer,
                  loss={NER_out_name: ner_loss_func, CLS_out_name: 'categorical_crossentropy'},
                  loss_weights={NER_out_name: 0., CLS_out_name: 10.},  # set weight of loss
                  metrics={NER_out_name: SequenceCorrectness(), CLS_out_name: 'categorical_accuracy'})

    model.fit(
              {'ner_input': ner_train_x, 'cls_input': cls_train_x},
              {NER_out_name: ner_train_y, CLS_out_name: cls_train_y},
              epochs=PRETRAIN_EPOCHS,
              batch_size=PRETRAIN_BATCHSIZE,
              class_weight={NER_out_name: None, CLS_out_name: 'auto'},   # cls loss multiply the class weights
              validation_data=[{'ner_input': ner_test_x, 'cls_input': cls_test_x},
                               {NER_out_name: ner_test_y, CLS_out_name: cls_test_y}],
              callbacks=callbacks_list,)

    # train model
    model.compile(optimizer=adam_optimizer,
                  loss={NER_out_name: ner_loss_func, CLS_out_name: 'categorical_crossentropy'},
                  loss_weights={NER_out_name: 15., CLS_out_name: 10.},  # set weight of loss
                  metrics={NER_out_name: SequenceCorrectness(), CLS_out_name: 'categorical_accuracy'})
    
    model.fit(
            {'ner_input': ner_train_x, 'cls_input': cls_train_x},
            {NER_out_name: ner_train_y, CLS_out_name: cls_train_y},
            epochs=EPOCHS,
            batch_size=BATCHSIZE,
            class_weight={NER_out_name: None, CLS_out_name: 'auto'},   # cls loss multiply the class weights
            validation_data=[{'ner_input': ner_test_x, 'cls_input': cls_test_x},
                             {NER_out_name: ner_test_y, CLS_out_name: cls_test_y}],
            callbacks=callbacks_list,)


    # save model
    model.save(create_file_dir_if_needed(config["h5_model_file"]))

    model.save_weights(create_file_dir_if_needed(config["h5_weights_file"]))

    semantic_vector.save(create_file_dir_if_needed(config["vector_model_file"]))
    
    semantic_vector.save_weights(create_file_dir_if_needed(config["vector_weights_file"]))

    tf.keras.experimental.export_saved_model(
            model, create_or_rm_dir_if_needed(config["saved_model_dir"])
        )

    tf.keras.experimental.export_saved_model(
            semantic_vector, create_or_rm_dir_if_needed(config["vector_model_dir"])
        )

    # save cls&NER deliverable model
    mtinput_export_as_deliverable_model(
        create_dir_if_needed(config["deliverable_model_dir"]),
        keras_saved_model=config["saved_model_dir"],
        converter_for_request=ConverterForMTRequest(),
        converter_for_response=ConverterForMTResponse_VirtualPad(prepad=CLS2NER_KEYWORD_LEN),
        lookup_tables={'vocab_lookup':vocabulary_lookuper,
                       'tag_lookup':ner_tag_lookuper,
                       'label_lookup':cls_label_lookuper},
        padding_parameter={"maxlen": MAX_SENTENCE_LEN, "value": 0, "padding": "post"},
        addition_model_dependency=["tf-crf-layer"],
        custom_object_dependency=["tf_crf_layer"],
        decoder_style=DECODER_STYLE,
    )

    # save semantic vector deliverable model
    mtinput_export_as_vector_deliverable_model(
        create_dir_if_needed(config["vector_deliverable_model_dir"]),
        keras_saved_model=config["vector_model_dir"],
        converter_for_request=ConverterForMTRequest(),
        converter_for_response=ConverterForVector(), 
        lookup_tables={'vocab_lookup':vocabulary_lookuper},
        padding_parameter={"maxlen": MAX_SENTENCE_LEN, "value": 0, "padding": "post"},
        addition_model_dependency=["tf-crf-layer"],
        custom_object_dependency=["tf_crf_layer"],
        decoder_style=DECODER_STYLE,
    )



if __name__ == "__main__":
    main()