from __future__ import absolute_import, division, print_function, unicode_literals

import os
import pandas as pd
from sklearn.model_selection import train_test_split
import numpy as np

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

from docproduct.bert import build_model_from_config

from keras_bert.loader import load_model_weights_from_checkpoint


class FFN(tf.keras.layers.Layer):
    def __init__(
            self,
            hidden_size=768,
            dropout=0.2,
            residual=True,
            name='FFN',
            **kwargs):
        """Simple Dense wrapped with various layers
        """

        super(FFN, self).__init__(name=name, **kwargs)
        self.hidden_size = hidden_size
        self.dropout = dropout
        self.residual = residual
        self.ffn_layer = tf.keras.layers.Dense(
            units=hidden_size,
            use_bias=True
        )

    def call(self, inputs):
        ffn_embedding = self.ffn_layer(inputs)
        ffn_embedding = tf.keras.layers.ReLU()(ffn_embedding)
        if self.dropout > 0:
            ffn_embedding = tf.keras.layers.Dropout(
                self.dropout)(ffn_embedding)

        if self.residual:
            ffn_embedding += inputs
        return ffn_embedding


class MedicalQAModel(tf.keras.Model):
    def __init__(self, name=''):
        super(MedicalQAModel, self).__init__(name=name)
        self.q_ffn = FFN(name='q_ffn', input_shape=(768,))
        self.a_ffn = FFN(name='a_ffn', input_shape=(768,))

    def call(self, inputs):
        q_bert_embedding, a_bert_embedding = tf.unstack(inputs, axis=1)
        q_embedding, a_embedding = self.q_ffn(
            q_bert_embedding), self.a_ffn(a_bert_embedding)
        return tf.stack([q_embedding, a_embedding], axis=1)


class MedicalQAModelwithBert(tf.keras.Model):
    def __init__(
            self,
            hidden_size=768,
            dropout=0.2,
            residual=True,
            config_file=None,
            checkpoint_file=None,
            bert_trainable=True,
            layer_ind=-1,
            name=''):
        super(MedicalQAModelwithBert, self).__init__(name=name)
        build = checkpoint_file != None
        self.biobert, config = build_model_from_config(
            config_file=config_file,
            training=False,
            trainable=bert_trainable,
            build=build)
        if checkpoint_file is not None:
            load_model_weights_from_checkpoint(
                model=self.biobert, config=config, checkpoint_file=checkpoint_file, training=False)
        self.q_ffn_layer = FFN(
            hidden_size=hidden_size,
            dropout=dropout,
            residual=residual,
            name='q_ffn')
        self.a_ffn_layer = FFN(
            hidden_size=hidden_size,
            dropout=dropout,
            residual=residual,
            name='a_ffn')
        self.layer_ind = layer_ind

    def call(self, inputs):

        if 'q_input_ids' in inputs:
            with_question = True
        else:
            with_question = False

        if 'a_input_ids' in inputs:
            with_answer = True
        else:
            with_answer = False
        # according to USE, the DAN network average embedding across tokens
        if with_question:
            q_bert_embedding = self.biobert(
                (inputs['q_input_ids'], inputs['q_segment_ids'], inputs['q_input_masks']))[self.layer_ind]
            q_bert_embedding = tf.reduce_mean(q_bert_embedding, axis=1)
        if with_answer:
            a_bert_embedding = self.biobert(
                (inputs['a_input_ids'], inputs['a_segment_ids'], inputs['a_input_masks']))[self.layer_ind]
            a_bert_embedding = tf.reduce_mean(a_bert_embedding, axis=1)

        if with_question:
            q_embedding = self.q_ffn_layer(q_bert_embedding)
            output = q_embedding
        if with_answer:
            a_embedding = self.a_ffn_layer(a_bert_embedding)
            output = a_embedding

        if with_question and with_answer:
            output = tf.stack([q_embedding, a_embedding], axis=1)

        return output