import argparse

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

from docproduct.dataset import create_dataset_for_ffn
from docproduct.models import MedicalQAModel
from docproduct.loss import qa_pair_loss, qa_pair_cross_entropy_loss
from docproduct.metrics import qa_pair_batch_accuracy

DEVICE = ["/gpu:0", "/gpu:1"]


def multi_gpu_train(batch_size, num_gpu, data_path, num_epochs, model_path, loss=qa_pair_loss):
    mirrored_strategy = tf.distribute.MirroredStrategy(
        devices=DEVICE[:num_gpu])
    global_batch_size = batch_size*num_gpu
    learning_rate = learning_rate*1.5**num_gpu
    with mirrored_strategy.scope():
        d = create_dataset_for_ffn(
            data_path, batch_size=global_batch_size, shuffle_buffer=100000)

        d_iter = mirrored_strategy.make_dataset_iterator(d)

        medical_qa_model = tf.keras.Sequential()
        medical_qa_model.add(tf.keras.layers.Input((2, 768)))
        medical_qa_model.add(MedicalQAModel())
        optimizer = tf.keras.optimizers.Adam(lr=learning_rate)
        medical_qa_model.compile(
            optimizer=optimizer, loss=loss)

    epochs = num_epochs
    loss_metric = tf.keras.metrics.Mean()

    medical_qa_model.fit(d_iter, epochs=epochs, metrics=[
                         qa_pair_batch_accuracy])
    medical_qa_model.save_weights(model_path)
    return medical_qa_model


def single_gpu_train(batch_size, num_gpu, data_path, num_epochs, model_path, loss=qa_pair_loss):
    global_batch_size = batch_size*num_gpu
    learning_rate = learning_rate
    d = create_dataset_for_ffn(
        data_path, batch_size=global_batch_size, shuffle_buffer=500000)
    eval_d = create_dataset_for_ffn(
        data_path, batch_size=batch_size, mode='eval')

    medical_qa_model = MedicalQAModel()
    optimizer = tf.keras.optimizers.Adam(lr=learning_rate)
    medical_qa_model.compile(
        optimizer=optimizer, loss=loss, metrics=[
            qa_pair_batch_accuracy])

    epochs = num_epochs

    medical_qa_model.fit(d, epochs=epochs, validation_data=eval_d)
    medical_qa_model.save_weights(model_path)
    return medical_qa_model


def train_ffn(model_path='models/ffn_crossentropy/ffn',
              data_path='data/mqa_csv',
              num_epochs=300,
              num_gpu=1,
              batch_size=64,
              learning_rate=0.0001,
              validation_split=0.2,
              loss='categorical_crossentropy'):

    if loss == 'categorical_crossentropy':
        loss_fn = qa_pair_cross_entropy_loss
    else:
        loss_fn = qa_pair_loss
    eval_d = create_dataset_for_ffn(
        data_path, batch_size=batch_size, mode='eval')

    if num_gpu > 1:
        medical_qa_model = multi_gpu_train(
            batch_size, num_gpu, data_path, num_epochs, model_path, loss_fn)
    else:
        medical_qa_model = single_gpu_train(
            batch_size, num_gpu, data_path, num_epochs, model_path, loss_fn)

    medical_qa_model.summary()
    medical_qa_model.save_weights(model_path, overwrite=True)
    # K.set_learning_phase(0)
    # q_embedding, a_embedding = tf.unstack(
    #     medical_qa_model(next(iter(eval_d))[0]), axis=1)

    # q_embedding = q_embedding / tf.norm(q_embedding, axis=-1, keepdims=True)
    # a_embedding = a_embedding / tf.norm(a_embedding, axis=-1, keepdims=True)

    # batch_score = tf.reduce_sum(q_embedding*a_embedding, axis=-1)
    # baseline_score = tf.reduce_mean(
    #     tf.matmul(q_embedding, tf.transpose(a_embedding)))

    # print('Eval Batch Cos similarity')
    # print(tf.reduce_mean(batch_score))
    # print('Baseline: {0}'.format(baseline_score))

    # medical_qa_model.save_weights(model_path, overwrite=True)


if __name__ == "__main__":

    train_ffn()