import math

import tensorflow as tf
from transformers import TFAutoModel

from capreolus import ConfigOption, Dependency, get_logger
from capreolus.reranker import Reranker
from capreolus.reranker.common import NewRbfKernelBankTF, new_similarity_matrix_tf

logger = get_logger(__name__)


class TFCEDRKNRM_Class(tf.keras.layers.Layer):
    def __init__(self, extractor, config, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.extractor = extractor
        self.config = config

        if config["pretrained"] == "electra-base-msmarco":
            self.bert = TFAutoModel.from_pretrained(
                "Capreolus/electra-base-msmarco", hidden_dropout_prob=config["hidden_dropout_prob"], output_hidden_states=True
            )
        elif config["pretrained"] == "electra-base":
            self.bert = TFAutoModel.from_pretrained(
                "google/electra-base-discriminator", hidden_dropout_prob=config["hidden_dropout_prob"], output_hidden_states=True
            )
        elif config["pretrained"] == "bert-base-msmarco":
            self.bert = TFAutoModel.from_pretrained(
                "Capreolus/bert-base-msmarco", hidden_dropout_prob=config["hidden_dropout_prob"], output_hidden_states=True
            )
        elif config["pretrained"] == "bert-base-uncased":
            self.bert = TFAutoModel.from_pretrained(
                "bert-base-uncased", hidden_dropout_prob=config["hidden_dropout_prob"], output_hidden_states=True
            )

        self.hidden_size = self.bert.config.hidden_size
        mus = list(self.config["mus"]) + [1.0]
        sigmas = [self.config["sigma"] for _ in self.config["mus"]] + [0.01]
        logger.debug("mus: %s", mus)
        self.kernels = NewRbfKernelBankTF(mus, sigmas, dim=1, requires_grad=self.config["gradkernels"])

        if -1 in self.config["simmat_layers"]:
            assert len(self.config["simmat_layers"]) == 1
            assert self.config["cls"] is not None
            self._compute_simmat = False
            combine_size = 0
        else:
            self._compute_simmat = True
            combine_size = self.kernels.count() * len(self.config["simmat_layers"])

        assert self.config["cls"] in ("avg", "max", None)
        if self.config["cls"]:
            combine_size += self.hidden_size

        # use weight init from PyTorch 0.4
        if config["combine_hidden"] == 0:
            stdv = 1.0 / math.sqrt(combine_size)
            weight_init = tf.keras.initializers.RandomUniform(minval=-stdv, maxval=stdv)
            combine_steps = [
                tf.keras.layers.Dense(1, input_shape=(combine_size,), kernel_initializer=weight_init, dtype=tf.float32)
            ]
        else:
            stdv1 = 1.0 / math.sqrt(combine_size)
            weight_init1 = tf.keras.initializers.RandomUniform(minval=-stdv1, maxval=stdv1)
            stdv2 = 1.0 / math.sqrt(config["combine_hidden"])
            weight_init2 = tf.keras.initializers.RandomUniform(minval=-stdv2, maxval=stdv2)
            combine_steps = [
                tf.keras.layers.Dense(
                    config["combine_hidden"], input_shape=(combine_size,), kernel_initializer=weight_init1, dtype=tf.float32
                ),
                tf.keras.layers.Dense(
                    1, input_shape=(config["combine_hidden"],), kernel_initializer=weight_init2, dtype=tf.float32
                ),
            ]

        self.combine = tf.keras.Sequential(combine_steps)

        self.num_passages = extractor.config["numpassages"]
        self.maxseqlen = extractor.config["maxseqlen"]
        # TODO we include SEP in maxqlen due to the way the simmat is constructed... (and another SEP in document)
        # (maxqlen is the actual query length and does not count CLS or SEP)
        self.maxqlen = extractor.config["maxqlen"] + 1
        # decreased by 1 because we remove CLS before generating embeddings
        self.maxdoclen = self.maxseqlen - 1

    def masked_simmats(self, embeddings, bert_mask, bert_segments):
        bert_mask = tf.cast(bert_mask, embeddings.dtype)
        # segment 0 contains '[CLS] query [SEP]' and segment 1 contains 'document [SEP]'
        query_mask = bert_mask * tf.cast(bert_segments == 0, bert_mask.dtype)
        query_mask = tf.expand_dims(query_mask, axis=-1)
        padded_query = (query_mask * embeddings)[:, : self.maxqlen]
        query_mask = query_mask[:, : self.maxqlen]

        doc_mask = bert_mask * tf.cast(bert_segments == 1, bert_mask.dtype)
        doc_mask = tf.expand_dims(doc_mask, axis=-1)
        # padded_doc length is maxsdoclen; zero padding on  both left and right of doc
        padded_doc = doc_mask * embeddings

        # (maxqlen, maxdoclen)
        simmat = new_similarity_matrix_tf(padded_query, padded_doc, query_mask, doc_mask, 0)
        return simmat, doc_mask, query_mask

    def knrm(self, bert_output, bert_mask, bert_segments, batch_size):
        # create similarity matrix for each passage (skipping CLS)
        passage_simmats, passage_doc_mask, passage_query_mask = self.masked_simmats(
            bert_output[:, 1:], bert_mask[:, 1:], bert_segments[:, 1:]
        )

        passage_simmats = tf.reshape(passage_simmats, [batch_size, self.num_passages, self.maxqlen, self.maxdoclen])
        passage_doc_mask = tf.reshape(passage_doc_mask, [batch_size, self.num_passages, 1, -1])

        # concat similarity matrices along document dimension; query mask is the same across passages
        doc_simmat = tf.concat([passage_simmats[:, PIDX, :, :] for PIDX in range(self.num_passages)], axis=2)
        doc_mask = tf.concat([passage_doc_mask[:, PIDX, :, :] for PIDX in range(self.num_passages)], axis=2)
        query_mask = tf.reshape(passage_query_mask, [batch_size, self.num_passages, -1, 1])[:, 0, :, :]

        # KNRM on similarity matrix
        prepooled_doc = self.kernels(doc_simmat)
        prepooled_doc = (
            prepooled_doc * tf.reshape(doc_mask, [batch_size, 1, 1, -1]) * tf.reshape(query_mask, [batch_size, 1, -1, 1])
        )

        # sum over document
        knrm_features = tf.reduce_sum(prepooled_doc, axis=3)
        knrm_features = tf.math.log(tf.maximum(knrm_features, 1e-6)) * 0.01
        # sum over query
        knrm_features = tf.reduce_sum(knrm_features, axis=2)

        return knrm_features

    def call(self, x, **kwargs):
        doc_input, doc_mask, doc_seg = x[0], x[1], x[2]
        batch_size = tf.shape(doc_input)[0]

        doc_input = tf.reshape(doc_input, [batch_size * self.num_passages, self.maxseqlen])
        doc_mask = tf.reshape(doc_mask, [batch_size * self.num_passages, self.maxseqlen])
        doc_seg = tf.reshape(doc_seg, [batch_size * self.num_passages, self.maxseqlen])

        # get BERT embeddings (including CLS) for each passage
        # TODO switch to hgf's ModelOutput after bumping tranformers version
        outputs = self.bert(doc_input, attention_mask=doc_mask, token_type_ids=doc_seg)
        if self.config["pretrained"].startswith("bert-"):
            outputs = (outputs[0], outputs[2])
        bert_output, all_layer_output = outputs

        #  embeddings to create the CLS feature
        cls = bert_output[:, 0, :]
        if self.config["cls"] == "max":
            cls_features = tf.reshape(cls, [batch_size, self.num_passages, self.hidden_size])
            cls_features = tf.reduce_max(cls_features, axis=1)
        elif self.config["cls"] == "avg":
            cls_features = tf.reshape(cls, [batch_size, self.num_passages, self.hidden_size])
            cls_features = tf.reduce_mean(cls_features, axis=1)

        # create KNRM features for each output layer
        if self._compute_simmat:
            layer_knrm_features = [
                self.knrm(all_layer_output[LIDX], doc_mask, doc_seg, batch_size) for LIDX in self.config["simmat_layers"]
            ]

        # concat CLS+KNRM features and pass to linear layer
        if self.config["cls"] and self._compute_simmat:
            all_features = tf.concat([cls_features] + layer_knrm_features, axis=1)
        elif self._compute_simmat:
            all_features = tf.concat(layer_knrm_features, axis=1)
        elif self.config["cls"]:
            all_features = cls_features
        else:
            raise ValueError("invalid config: %s" % self.config)

        score = self.combine(all_features)
        return score

    def predict_step(self, data):
        posdoc_bert_input, posdoc_mask, posdoc_seg, negdoc_bert_input, negdoc_mask, negdoc_seg = data
        doc_scores = self.call((posdoc_bert_input, posdoc_mask, posdoc_seg), training=False)

        return doc_scores

    def score(self, x, **kwargs):
        posdoc_bert_input, posdoc_mask, posdoc_seg, negdoc_bert_input, negdoc_mask, negdoc_seg = x

        return self.call((posdoc_bert_input, posdoc_mask, posdoc_seg), **kwargs)

    def score_pair(self, x, **kwargs):
        posdoc_bert_input, posdoc_mask, posdoc_seg, negdoc_bert_input, negdoc_mask, negdoc_seg = x

        pos_score = self.call((posdoc_bert_input, posdoc_mask, posdoc_seg), **kwargs)
        neg_score = self.call((negdoc_bert_input, negdoc_mask, negdoc_seg), **kwargs)

        return pos_score, neg_score


@Reranker.register
class TFCEDRKNRM(Reranker):
    """
    TensorFlow implementation of CEDR-KNRM.
    Equivalant to BERT-KNRM when cls=None.

    CEDR: Contextualized Embeddings for Document Ranking
    Sean MacAvaney, Andrew Yates, Arman Cohan, and Nazli Goharian. SIGIR 2019.
    https://arxiv.org/pdf/1904.07094
    """

    module_name = "TFCEDRKNRM"

    dependencies = [
        Dependency(key="extractor", module="extractor", name="pooledbertpassage"),
        Dependency(key="trainer", module="trainer", name="tensorflow"),
    ]
    config_spec = [
        ConfigOption(
            "pretrained",
            "electra-base",
            "Pretrained model: bert-base-uncased, bert-base-msmarco, electra-base, or electra-base-msmarco",
        ),
        ConfigOption("mus", [-0.9, -0.7, -0.5, -0.3, -0.1, 0.1, 0.3, 0.5, 0.7, 0.9], "mus", value_type="floatlist"),
        ConfigOption("sigma", 0.1, "sigma"),
        ConfigOption("gradkernels", True, "tune mus and sigmas"),
        ConfigOption("hidden_dropout_prob", 0.1, "The dropout probability of BERT-like model's hidden layers."),
        ConfigOption("simmat_layers", "0..12,1", "Layer outputs to include in similarity matrix", value_type="intlist"),
        ConfigOption("combine_hidden", 1024, "Hidden size to use with combination FC layer (0 to disable)"),
        ConfigOption("cls", "avg", "Handling of CLS token: avg, max, or None"),
    ]

    def build_model(self):
        if not hasattr(self, "model"):
            with self.trainer.strategy.scope():
                self.model = TFCEDRKNRM_Class(self.extractor, self.config)
        return self.model
