"""Head modules

"""

import numpy as np

from bpnet.utils import dict_prefix_key

from bpnet.metrics import ClassificationMetrics, RegressionMetrics

import keras.backend as K

import tensorflow as tf

import keras.layers as kl

import gin

import os

import abc

class BaseHead:

    # loss

    # weight -> loss weight (1 by default)

    # kwargs -> kwargs for the model

    # name -> name of the module

    # _model -> gets setup in the init

    @abc.abstractmethod

    def get_target(self, task):

        pass

    @abc.abstractmethod

    def __call__(self, inp, task):

        """Useful for writing together the model

        Returns the output tensor

        """

        raise NotImplementedError

    @abc.abstractmethod

    def get_preact_tensor(self, graph=None):

        """Return the single pre-activation tensors

        """

        pass

    @abc.abstractmethod

    def intp_tensors(self, preact_only=False, graph=None):

        """Dictionary of all available interpretation tensors

        for `get_interpretation_node`

        """

        raise NotImplementedError

    # @abc.abstractmethod

    # def get_intp_tensor(self, which=None):

    #     """Returns a target tensor which is a scalar

    #     w.r.t. to which to compute the outputs

    #     Args:

    #       which [string]: If None, use the default

    #       **kwargs: optional kwargs for the interpretation method

    #     Returns:

    #       scalar tensor

    #     """

    #     raise NotImplementedError

    def copy(self):

        from copy import deepcopy

        return deepcopy(self)

class BaseHeadWBias(BaseHead):

    @abc.abstractmethod

    def get_bias_input(self, task):

        pass

    @abc.abstractmethod

    def neutral_bias_input(self, task, length, seqlen):

        pass

def id_fn(x):

    return x

def named_tensor(x, name):

    return kl.Lambda(id_fn, name=name)(x)

# --------------------------------------------

# Head implementations

@gin.configurable

class ScalarHead(BaseHeadWBias):

    def __init__(self, target_name,  # "{task}/scalar"

                 net,  # function that takes a keras tensor and returns a keras tensor

                 activation=None,

                 loss='mse',

                 loss_weight=1,

                 metric=RegressionMetrics(),

                 postproc_fn=None,  # post-processing to apply so that we are in the right scale

                 # bias input

                 use_bias=False,

                 bias_net=None,

                 bias_input='bias/{task}/scalar',

                 bias_shape=(1,),

                 ):

        self.net = net

        self.loss = loss

        self.loss_weight = loss_weight

        self.metric = metric

        self.postproc_fn = postproc_fn

        self.target_name = target_name

        self.activation = activation

        self.bias_input = bias_input

        self.bias_net = bias_net

        self.use_bias = use_bias

        self.bias_shape = bias_shape

    def get_target(self, task):

        return self.target_name.format(task=task)

    def __call__(self, inp, task):

        o = self.net(inp)

        # remember the tensors useful for interpretation (referred by name)

        self.pre_act = o.name

        # add the target bias

        if self.use_bias:

            binp = kl.Input(self.bias_shape, name=self.get_bias_input(task))

            bias_inputs = [binp]

            # add the bias term

            if self.bias_net is not None:

                bias_x = self.bias_net(binp)

                # This allows to normalize the bias data first

                # (e.g. when we have profile counts to aggregate it first)

            else:

                # Don't use the nn 'bias' so that when the measurement bias = 0,

                # this term vanishes

                bias_x = kl.Dense(1, use_bias=False)(binp)

            o = kl.add([o, bias_x])

        else:

            bias_inputs = []

        if self.activation is not None:

            if isinstance(self.activation, str):

                o = kl.Activation(self.activation)(o)

            else:

                o = self.activation(o)

        self.post_act = o.name

        # label the target op so that we can use a dictionary of targets

        # to train the model

        return named_tensor(o, name=self.get_target(task)), bias_inputs

    def get_preact_tensor(self, graph=None):

        if graph is None:

            graph = tf.get_default_graph()

        return graph.get_tensor_by_name(self.pre_act)

    def intp_tensors(self, preact_only=False, graph=None):

        """Return the required interpretation tensors

        """

        if graph is None:

            graph = tf.get_default_graph()

        if self.activation is None:

            # the post-activation doesn't

            # have any specific meaning when

            # we don't use any activation function

            return {"pre-act": graph.get_tensor_by_name(self.pre_act)}

        if preact_only:

            return {"pre-act": graph.get_tensor_by_name(self.pre_act)}

        else:

            return {"pre-act": graph.get_tensor_by_name(self.pre_act),

                    "output": graph.get_tensor_by_name(self.post_act)}

    # def get_intp_tensor(self, which='pre-act'):

    #     return self.intp_tensors()[which]

    def get_bias_input(self, task):

        return self.bias_input.format(task=task)

    def neutral_bias_input(self, task, length, seqlen):

        """Create dummy bias input

        Return: (k, v) tuple

        """

        shape = tuple([x if x is not None else seqlen

                       for x in self.bias_shape])

        return (self.get_bias_input(task), np.zeros((length, ) + shape))

@gin.configurable

class BinaryClassificationHead(ScalarHead):

    def __init__(self, target_name,  # "{task}/scalar"

                 net,  # function that takes a keras tensor and returns a keras tensor

                 activation='sigmoid',

                 loss='binary_crossentropy',

                 loss_weight=1,

                 metric=ClassificationMetrics(),

                 postproc_fn=None,

                 # bias input

                 use_bias=False,

                 bias_net=None,

                 bias_input='bias/{task}/scalar',

                 bias_shape=(1,),

                 ):

        # override the default

        super().__init__(target_name,

                         net,

                         activation=activation,

                         loss=loss,

                         metric=metric,

                         postproc_fn=postproc_fn,

                         use_bias=use_bias,

                         bias_net=bias_net,

                         bias_input=bias_input,

                         bias_shape=bias_shape)

        # TODO - mabye override the way we call outputs?

@gin.configurable

class ProfileHead(BaseHeadWBias):

    """Deals with the case where the output are multiple tracks of

    total shape (L, C) (L = sequence length, C = number of channels)

    Note: Since the contribution score will be a single scalar, the

    interpretation method will have to aggregate both across channels

    as well as positions

    """

    def __init__(self, target_name,  # "{task}/profile"

                 net,  # function that takes a keras tensor and returns a keras tensor

                 activation=None,

                 loss='mse',

                 loss_weight=1,

                 metric=RegressionMetrics(),

                 postproc_fn=None,

                 # bias input

                 use_bias=False,

                 bias_net=None,

                 bias_input='bias/{task}/profile',

                 bias_shape=(None, 1),

                 ):

        self.net = net

        self.loss = loss

        self.loss_weight = loss_weight

        self.metric = metric

        self.postproc_fn = postproc_fn

        self.target_name = target_name

        self.activation = activation

        self.bias_input = bias_input

        self.bias_net = bias_net

        self.use_bias = use_bias

        self.bias_shape = bias_shape

    def get_target(self, task):

        return self.target_name.format(task=task)

    def __call__(self, inp, task):

        o = self.net(inp)

        # remember the tensors useful for interpretation (referred by name)

        self.pre_act = o.name

        # add the target bias

        if self.use_bias:

            binp = kl.Input(self.bias_shape, name=self.get_bias_input(task))

            bias_inputs = [binp]

            # add the bias term

            if self.bias_net is not None:

                bias_x = self.bias_net(binp)

                # This allows to normalize the bias data first

                # (e.g. when we have profile counts to aggregate it first)

            else:

                # Don't use the nn 'bias' so that when the measurement bias = 0,

                # this term vanishes

                bias_x = kl.Conv1D(1, kernel_size=1, use_bias=False)(binp)

            o = kl.add([o, bias_x])

        else:

            bias_inputs = []

        if self.activation is not None:

            if isinstance(self.activation, str):

                o = kl.Activation(self.activation)(o)

            else:

                o = self.activation(o)

        self.post_act = o.name

        # label the target op so that we can use a dictionary of targets

        # to train the model

        return named_tensor(o, name=self.get_target(task)), bias_inputs

    def get_preact_tensor(self, graph=None):

        if graph is None:

            graph = tf.get_default_graph()

        return graph.get_tensor_by_name(self.pre_act)

    @staticmethod

    def profile_contrib(p):

        """Summarizing the profile for the contribution scores

        wn: Normalized contribution (weighted sum of the contribution scores)

          where the weighted sum uses softmax(p) to weight it

        w2: Simple sum (p**2)

        w1: sum(p)

        winf: max(p)

        """

        # Note: unfortunately we have to use the kl.Lambda boiler-plate

        # to be able to do Model(inp, outputs) in deep-explain code

        # Normalized contribution  - # TODO - update with tensorflow

        wn = kl.Lambda(lambda p:

                       K.mean(K.sum(K.stop_gradient(tf.nn.softmax(p, dim=-2)) * p, axis=-2), axis=-1)

                       )(p)

        # Squared weight

        w2 = kl.Lambda(lambda p:

                       K.mean(K.sum(p * p, axis=-2), axis=-1)

                       )(p)

        # W1 weight

        w1 = kl.Lambda(lambda preact_m:

                       K.mean(K.sum(preact_m, axis=-2), axis=-1)

                       )(p)

        # Winf

        # 1. max across the positional axis, average the strands

        winf = kl.Lambda(lambda p:

                         K.mean(K.max(p, axis=-2), axis=-1)

                         )(p)

        return {"wn": wn,

                "w1": w1,

                "w2": w2,

                "winf": winf

                }

    def intp_tensors(self, preact_only=False, graph=None):

        """Return the required interpretation tensors (scalars)

        Note: Since we are predicting a track,

            we should return a single scalar here

        """

        if graph is None:

            graph = tf.get_default_graph()

        preact = graph.get_tensor_by_name(self.pre_act)

        postact = graph.get_tensor_by_name(self.post_act)

        # Contruct the profile summary ops

        preact_tensors = self.profile_contrib(preact)

        postact_tensors = dict_prefix_key(self.profile_contrib(postact), 'output_')

        if self.activation is None:

            # the post-activation doesn't

            # have any specific meaning when

            # we don't use any activation function

            return preact_tensors

        if preact_only:

            return preact_tensors

        else:

            return {**preact_tensors, **postact_tensors}

    # def get_intp_tensor(self, which='wn'):

    #     return self.intp_tensors()[which]

    def get_bias_input(self, task):

        return self.bias_input.format(task=task)

    def neutral_bias_input(self, task, length, seqlen):

        """Create dummy bias input

        Return: (k, v) tuple

        """

        shape = tuple([x if x is not None else seqlen

                       for x in self.bias_shape])

        return (self.get_bias_input(task), np.zeros((length, ) + shape))