import argparse
import os

import tensorflow as tf
import tensorflow.keras.backend as K

from docproduct.dataset import create_dataset_for_bert
from docproduct.models import MedicalQAModelwithBert
from docproduct.loss import qa_pair_loss, qa_pair_cross_entropy_loss
from docproduct.tokenization import FullTokenizer
from docproduct.metrics import qa_pair_batch_accuracy


def train_bertffn(model_path='models/bertffn_crossentropy/bertffn',
                  data_path='data/mqa_csv',
                  num_epochs=20,
                  num_gpu=1,
                  batch_size=64,
                  learning_rate=2e-5,
                  validation_split=0.2,
                  loss='categorical_crossentropy',
                  pretrained_path='models/pubmed_pmc_470k/',
                  max_seq_len=256):
    """A function to train BertFFNN similarity embedding model.

    Input file format:
        question,answer
        my eyes hurts, go see a doctor

    For more information about training details:
    https://github.com/Santosh-Gupta/DocProduct/blob/master/README.md

    Keyword Arguments:
        model_path {str} -- Path to save embedding model weights, ends with prefix of model files (default: {'models/bertffn_crossentropy/bertffn'})
        data_path {str} -- CSV data path (default: {'data/mqa_csv'})
        num_epochs {int} -- Number of Epochs to train (default: {20})
        num_gpu {int} -- Number of GPU to use(Currently only support single GPU) (default: {1})
        batch_size {int} -- Batch size (default: {64})
        learning_rate {float} -- learning rate (default: {2e-5})
        validation_split {float} -- validation split (default: {0.2})
        loss {str} -- Loss type, either MSE or crossentropy (default: {'categorical_crossentropy'})
        pretrained_path {str} -- Pretrained bioBert model path (default: {'models/pubmed_pmc_470k/'})
        max_seq_len {int} -- Max sequence length of model(No effects if dynamic padding is enabled) (default: {256})
    """
    tf.compat.v1.disable_eager_execution()
    if loss == 'categorical_crossentropy':
        loss_fn = qa_pair_cross_entropy_loss
    else:
        loss_fn = qa_pair_loss
    K.set_floatx('float32')
    tokenizer = FullTokenizer(os.path.join(pretrained_path, 'vocab.txt'))
    d = create_dataset_for_bert(
        data_path, tokenizer=tokenizer, batch_size=batch_size,
        shuffle_buffer=500000, dynamic_padding=True, max_seq_length=max_seq_len)
    eval_d = create_dataset_for_bert(
        data_path, tokenizer=tokenizer, batch_size=batch_size,
        mode='eval', dynamic_padding=True, max_seq_length=max_seq_len,
        bucket_batch_sizes=[64, 64, 64])

    medical_qa_model = MedicalQAModelwithBert(
        config_file=os.path.join(
            pretrained_path, 'bert_config.json'),
        checkpoint_file=os.path.join(pretrained_path, 'biobert_model.ckpt'))
    optimizer = tf.keras.optimizers.Adam(lr=learning_rate)
    medical_qa_model.compile(
        optimizer=optimizer, loss=loss_fn, metrics=[qa_pair_batch_accuracy])

    epochs = num_epochs

    callback = tf.keras.callbacks.ModelCheckpoint(
        model_path, verbose=1, save_weights_only=True, save_best_only=False, period=1)

    medical_qa_model.fit(d, epochs=epochs, callbacks=[callback])
    medical_qa_model.summary()
    medical_qa_model.save_weights(model_path)
    medical_qa_model.evaluate(eval_d)


if __name__ == "__main__":

    train_bertffn()