"""Graph representation aggregation layer."""
from abc import abstractmethod
from typing import List, NamedTuple

import tensorflow as tf
from dpu_utils.tf2utils import MLP, get_activation_function_by_name, unsorted_segment_softmax


class NodesToGraphRepresentationInput(NamedTuple):
    """A named tuple to hold input to layers computing graph representations from nodes
    representations."""

    node_embeddings: tf.Tensor
    node_to_graph_map: tf.Tensor
    num_graphs: tf.Tensor


class NodesToGraphRepresentation(tf.keras.layers.Layer):
    """Abstract class to compute graph representations from node representations.

    Throughout we use the following abbreviations in shape descriptions:
        * V: number of nodes (across all graphs)
        * VD: node representation dimension
        * G: number of graphs
        * GD: graph representation dimension
    """

    def __init__(self, graph_representation_size: int, **kwargs):
        super().__init__(**kwargs)
        self._graph_representation_size = graph_representation_size

    @abstractmethod
    def call(self, inputs: NodesToGraphRepresentationInput, training: bool = False):
        """Call the layer.

        Args:
            inputs: A tuple containing two items:
                node_embeddings: float32 tensor of shape [V, VD], the representation of each
                    node in all graphs.
                node_to_graph_map: int32 tensor of shape [V] with values in range [0, G-1],
                    mapping each node to a graph ID.
                num_graphs: int32 scalar, specifying the number G of graphs.
            training: A bool that denotes whether we are in training mode.

        Returns:
            float32 tensor of shape [G, GD]
        """
        pass


class WeightedSumGraphRepresentation(NodesToGraphRepresentation):
    """Layer computing graph representations as weighted sum of node representations.
    The weights are either computed from the original node representations ("self-attentional")
    or by a softmax across the nodes of a graph.
    Supports splitting operation into parallely computed independent "heads" which can focus
    on different aspects.

    Throughout we use the following abbreviations in shape descriptions:
        * V: number of nodes (across all graphs)
        * VD: node representation dimension
        * G: number of graphs
        * GD: graph representation dimension
        * H: number of heads
    """

    def __init__(
        self,
        graph_representation_size: int,
        num_heads: int,
        weighting_fun: str = "softmax",  # One of {"softmax", "sigmoid"}
        scoring_mlp_layers: List[int] = [128],
        scoring_mlp_activation_fun: str = "ReLU",
        scoring_mlp_dropout_rate: float = 0.2,
        transformation_mlp_layers: List[int] = [128],
        transformation_mlp_activation_fun: str = "ReLU",
        transformation_mlp_dropout_rate: float = 0.2,
        **kwargs,
    ):
        """
        Args:
            graph_representation_size: Size of the computed graph representation.
            num_heads: Number of independent heads to use to compute weights.
            weighting_fun: "sigmoid" ([0, 1] weights for each node computed from its
                representation), "softmax" ([0, 1] weights for each node computed
                from all nodes in same graph), "average" (weight is fixed to 1/num_nodes),
                or "none" (weight is fixed to 1).
            scoring_mlp_layers: MLP layer structure for computing raw scores turned into
                weights.
            scoring_mlp_activation_fun: MLP activcation function for computing raw scores
                turned into weights.
            scoring_mlp_dropout_rate: MLP inter-layer dropout rate for computing raw scores
                turned into weights.
            transformation_mlp_layers: MLP layer structure for computing graph representations.
            transformation_mlp_activation_fun: MLP activcation function for computing graph
                representations.
            transformation_mlp_dropout_rate: MLP inter-layer dropout rate for computing graph
                representations.
        """
        super().__init__(graph_representation_size, **kwargs)
        assert (
            graph_representation_size % num_heads == 0
        ), f"Number of heads {num_heads} needs to divide final representation size {graph_representation_size}!"
        assert weighting_fun.lower() in {
            "none",
            "average",
            "softmax",
            "sigmoid",
        }, f"Weighting function {weighting_fun} unknown, {{'softmax', 'sigmoid', 'none', 'average'}} supported."

        self._num_heads = num_heads
        self._weighting_fun = weighting_fun.lower()
        self._scoring_mlp_layers = scoring_mlp_layers
        self._scoring_mlp_activation_fun = get_activation_function_by_name(
            scoring_mlp_activation_fun
        )
        self._scoring_mlp_dropout_rate = scoring_mlp_dropout_rate
        self._transformation_mlp_layers = transformation_mlp_layers
        self._transformation_mlp_activation_fun = get_activation_function_by_name(
            transformation_mlp_activation_fun
        )
        self._transformation_mlp_dropout_rate = transformation_mlp_dropout_rate

    def build(self, input_shapes: NodesToGraphRepresentationInput):
        with tf.name_scope("WeightedSumGraphRepresentation"):
            if self._weighting_fun not in ("none", "average"):
                self._scoring_mlp = MLP(
                    out_size=self._num_heads,
                    hidden_layers=self._scoring_mlp_layers,
                    activation_fun=self._scoring_mlp_activation_fun,
                    dropout_rate=self._scoring_mlp_dropout_rate,
                    name="ScoringMLP",
                )
                self._scoring_mlp.build(tf.TensorShape((None, input_shapes.node_embeddings[-1])))

            self._transformation_mlp = MLP(
                out_size=self._graph_representation_size,
                hidden_layers=self._transformation_mlp_layers,
                activation_fun=self._transformation_mlp_activation_fun,
                dropout_rate=self._transformation_mlp_dropout_rate,
                name="TransformationMLP",
            )
            self._transformation_mlp.build(tf.TensorShape((None, input_shapes.node_embeddings[-1])))

            super().build(input_shapes)

    @tf.function(
        input_signature=(
            NodesToGraphRepresentationInput(
                node_embeddings=tf.TensorSpec(shape=tf.TensorShape((None, None)), dtype=tf.float32),
                node_to_graph_map=tf.TensorSpec(shape=tf.TensorShape((None,)), dtype=tf.int32),
                num_graphs=tf.TensorSpec(shape=(), dtype=tf.int32),
            ),
            tf.TensorSpec(shape=(), dtype=tf.bool),
        )
    )
    def call(self, inputs: NodesToGraphRepresentationInput, training: bool = False):
        # (1) compute weights for each node/head pair:
        if self._weighting_fun not in ("none", "average"):
            scores = self._scoring_mlp(inputs.node_embeddings, training=training)  # Shape [V, H]
            if self._weighting_fun == "sigmoid":
                weights = tf.nn.sigmoid(scores)  # Shape [V, H]
            elif self._weighting_fun == "softmax":
                weights_per_head = []
                for head_idx in range(self._num_heads):
                    head_scores = scores[:, head_idx]  # Shape [V]
                    head_weights = unsorted_segment_softmax(
                        logits=head_scores,
                        segment_ids=inputs.node_to_graph_map,
                        num_segments=inputs.num_graphs,
                    )  # Shape [V]
                    weights_per_head.append(tf.expand_dims(head_weights, -1))
                weights = tf.concat(weights_per_head, axis=1)  # Shape [V, H]
            else:
                raise ValueError()

        # (2) compute representations for each node/head pair:
        node_reprs = self._transformation_mlp(
            inputs.node_embeddings, training=training
        )  # Shape [V, GD]
        node_reprs = tf.reshape(
            node_reprs,
            shape=(-1, self._num_heads, self._graph_representation_size // self._num_heads),
        )  # Shape [V, H, GD//H]

        # (3) if necessary, weight representations and aggregate by graph:
        if self._weighting_fun == "none":
            node_reprs = tf.reshape(
                node_reprs, shape=(-1, self._graph_representation_size)
            )  # Shape [V, GD]
            graph_reprs = tf.math.segment_sum(
                data=node_reprs, segment_ids=inputs.node_to_graph_map
            )  # Shape [G, GD]
        elif self._weighting_fun == "average":
            node_reprs = tf.reshape(
                node_reprs, shape=(-1, self._graph_representation_size)
            )  # Shape [V, GD]
            graph_reprs = tf.math.segment_mean(
                data=node_reprs, segment_ids=inputs.node_to_graph_map
            )  # Shape [G, GD]
        else:
            weights = tf.expand_dims(weights, -1)  # Shape [V, H, 1]
            weighted_node_reprs = weights * node_reprs  # Shape [V, H, GD//H]

            weighted_node_reprs = tf.reshape(
                weighted_node_reprs, shape=(-1, self._graph_representation_size)
            )  # Shape [V, GD]
            graph_reprs = tf.math.segment_sum(
                data=weighted_node_reprs, segment_ids=inputs.node_to_graph_map
            )  # Shape [G, GD]

        return graph_reprs
