# coding=utf-8
# Copyright 2023 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Bound propagation utilities."""

import abc
from typing import Any, Callable, Generic, Mapping, Tuple, TypeVar

import jax
import jax.numpy as jnp
import jax_verify
from jax_verify.src import bound_propagation
from jax_verify.src import bound_utils
from jax_verify.src import graph_traversal
from jax_verify.src import utils
from jax_verify.src.types import Index, Nest, Primitive, Tensor  # pylint: disable=g-multiple-import


T = TypeVar("T", bound=graph_traversal.TransformedNode)


class BackwardConcretizingTransform(
    graph_traversal.BackwardGraphTransform[T],
    Generic[T],  # Explicitly restate this, to aid PyType resolution.
    metaclass=abc.ABCMeta,
):
    """Abstract class for a Backward Transformation that can concretize bounds."""

    @abc.abstractmethod
    def concretize_args(self, primitive: Primitive) -> bool:
        """Return whether the arguments needs to be concretized.

        Args:
          primitive: Primitive that we are encountering.
        """

    @abc.abstractmethod
    def concrete_bound_chunk(
        self,
        graph: graph_traversal.PropagationGraph,
        inputs: Nest[graph_traversal.GraphInput],
        env: Mapping[jax.extend.core.Var, bound_propagation.LayerInput],
        node: jax.extend.core.Var,
        obj: Tensor,
    ) -> Tensor:
        """Computes concrete bounds for a chunk of neurons in the given layer.

        Args:
          graph: Graph to perform Backward Propagation on.
          inputs: Bounds on the inputs.
          env: Environment containing intermediate bound and shape information.
          node: Graph node to obtain a bound for.
          obj: One-hot tensor of shape (chunk_size, *node_shape) specifying, for
            each index in the chunk, an element of the objective. Non-zero entries
            may be +1 or -1 to request lower or upper bounds respectively.

        Returns:
          Bound of shape (chunk_size,) on the activation at `node`.
        """


class BackwardConcretizer(metaclass=abc.ABCMeta):
    """Abstract producer of concretize bounds by back-propagation."""

    @abc.abstractmethod
    def should_handle_as_subgraph(self, primitive: Primitive) -> bool:
        """Returns whether the primitive should be handled via its sub-graph."""

    @abc.abstractmethod
    def concretize_args(self, primitive: Primitive) -> bool:
        """Return whether the arguments needs to be concretized.

        Args:
          primitive: Primitive that we are encountering.
        """

    @abc.abstractmethod
    def concrete_bound(
        self,
        graph: graph_traversal.PropagationGraph,
        inputs: Nest[graph_traversal.GraphInput],
        env: Mapping[jax.extend.core.Var, bound_propagation.LayerInput],
        node_ref: jax.extend.core.Var,
    ) -> jax_verify.IntervalBound:
        """Perform backward linear bound computation for the node `index`.

        Args:
          graph: Graph to perform Backward Propagation on.
          inputs: Bounds on the inputs.
          env: Environment containing intermediate bound and shape information.
          node_ref: Graph node to obtain a bound for.

        Returns:
          concrete_bound: IntervalBound on the activation at `node`.
        """


class ChunkedBackwardConcretizer(BackwardConcretizer):
    """Concretizer that invokes the given transform in chunks for each layer."""

    def __init__(
        self,
        concretizing_transform: BackwardConcretizingTransform[Any],
        max_chunk_size: int = 0,
    ):
        self._concretizing_transform = concretizing_transform
        self._max_chunk_size = max_chunk_size

    def should_handle_as_subgraph(self, primitive: Primitive) -> bool:
        return self._concretizing_transform.should_handle_as_subgraph(primitive)

    def concretize_args(self, primitive: Primitive) -> bool:
        return self._concretizing_transform.concretize_args(primitive)

    def concrete_bound(
        self,
        graph: graph_traversal.PropagationGraph,
        inputs: Nest[graph_traversal.GraphInput],
        env: Mapping[jax.extend.core.Var, bound_propagation.LayerInput],
        node_ref: jax.extend.core.Var,
    ) -> jax_verify.IntervalBound:
        """Perform backward linear bound computation for the node `index`.

        Args:
          graph: Graph to perform Backward Propagation on.
          inputs: Bounds on the inputs.
          env: Environment containing intermediate bound and shape information.
          node_ref: Reference of the node to obtain a bound for.

        Returns:
          concrete_bound: IntervalBound on the activation at `node_ref`.
        """
        node = graph_traversal.read_env(env, node_ref)

        def bound_fn(obj: Tensor) -> Tuple[Tensor, Tensor]:
            # Handle lower bounds and upper bounds independently in the same chunk.
            obj = jnp.concatenate([obj, -obj], axis=0)

            all_bounds = self._concretizing_transform.concrete_bound_chunk(
                graph, inputs, env, node_ref, obj
            )

            # Separate out the lower and upper bounds.
            lower_bound, neg_upper_bound = jnp.split(all_bounds, 2, axis=0)
            upper_bound = -neg_upper_bound
            return lower_bound, upper_bound

        return jax_verify.IntervalBound(
            *utils.chunked_bounds(node.shape, self._max_chunk_size, bound_fn)
        )


def stop_gradient_postprocess(
    index: Index,
    concrete_bound: bound_propagation.Bound,
) -> bound_propagation.Bound:
    del index  # unused
    return jax_verify.IntervalBound(
        jax.lax.stop_gradient(concrete_bound.lower),
        jax.lax.stop_gradient(concrete_bound.upper),
    )


class BackwardConcretizingAlgorithm(
    bound_propagation.PropagationAlgorithm[bound_propagation.Bound]
):
    """Abstract Backward graph propagation method with forward concretization.

    A trace through the network is first obtained, then backward bound
    computations are performed for each of the node requiring to be concretized,
    such as intermediate nodes that are inputs to non-linearity or output values.

    Note that the resulting environment of intermediate bounds (returned by
    `propagate` alongside the concrete output bound) is sparse, in the sense that
    many nodes will only contain a placeholder `GraphNode` with vacuous bounds.
    The only nodes guaranteed to contain genuine bounds are the args to
    concretised nodes, and the final outputs. In addition, the graph inputs are
    included in concrete form for convenience.
    """

    def __init__(
        self,
        backward_concretizer: BackwardConcretizer,
        bound_postprocess_fn: Callable[
            [Index, bound_propagation.LayerInput], bound_propagation.LayerInput
        ] = lambda _, x: x,
    ):
        self._backward_concretizer = backward_concretizer
        self._bound_postprocess_fn = bound_postprocess_fn

    def propagate(
        self,
        graph: graph_traversal.PropagationGraph,
        inputs: Nest[graph_traversal.GraphInput],
    ) -> Tuple[
        Nest[bound_propagation.LayerInput],
        Mapping[jax.extend.core.Var, bound_propagation.LayerInput],
    ]:
        subgraph_decider = self._backward_concretizer.should_handle_as_subgraph
        graph_inspector = bound_utils.GraphInspector(subgraph_decider)
        inspector_algorithm = bound_propagation.ForwardPropagationAlgorithm(
            graph_inspector
        )
        gn_outvals, env = inspector_algorithm.propagate(graph, inputs)
        env = dict(env)  # take a modifiable copy

        def lazily_concretize(index, *, is_output):
            jaxpr_node = graph.jaxpr_node(index)
            graph_node = env[jaxpr_node]
            if isinstance(graph_node, bound_utils.GraphNode):
                # This node has not yet been concretized. Perform concretization.
                concrete_bound = self._backward_concretizer.concrete_bound(
                    graph, inputs, env, jaxpr_node
                )
                if not is_output:
                    concrete_bound = self._bound_postprocess_fn(index, concrete_bound)
                # Update the environment, replacing the graph node with the concretised
                # bound. Note that the same graph node may occur multiple times, under
                # separate variables as we step into or out of a subgraph.
                for k, v in env.items():
                    if v is graph_node:
                        env[k] = concrete_bound

        # Iterate over the nodes in order so that we get intermediate bounds in
        # the order where we need them.
        for node in graph_inspector.nodes.values():
            if not node.is_input() and self._backward_concretizer.concretize_args(
                node.primitive
            ):
                for node_arg in node.args:
                    if isinstance(node_arg, bound_propagation.Bound):
                        lazily_concretize(node_arg.index, is_output=False)

        # Iterate over the outputs, making sure to concretize all of them.
        outvals = []
        for gn in gn_outvals:
            lazily_concretize(gn.index, is_output=True)
            outvals.append(env[graph.jaxpr_node(gn.index)])

        # Fill in the bounds for the inputs.
        # This is unnecessary for backward methods themselves, but may be useful
        # if the resulting `env` is used as a set of base bounds for a forward
        # method.
        flat_inputs, _ = jax.tree_util.tree_flatten(inputs)
        for in_jaxpr_node, in_bound in zip(graph.inputs, flat_inputs):
            if isinstance(env[in_jaxpr_node], bound_utils.GraphNode):
                env[in_jaxpr_node] = in_bound

        return outvals, env


class BackwardAlgorithmForwardConcretization(
    bound_propagation.PropagationAlgorithm[bound_propagation.Bound]
):
    """Abstract Backward graph propagation with forward concretization."""

    def __init__(
        self,
        forward_transform: bound_propagation.BoundTransform,
        backward_concretizer: BackwardConcretizer,
    ):
        self._forward_algorithm = bound_propagation.ForwardPropagationAlgorithm(
            forward_transform
        )
        self._backward_concretizer = backward_concretizer

    def propagate(
        self,
        graph: graph_traversal.PropagationGraph,
        inputs: Nest[graph_traversal.GraphInput],
    ) -> Tuple[
        Nest[bound_propagation.LayerInput],
        Mapping[jax.extend.core.Var, bound_propagation.LayerInput],
    ]:
        # Perform the forward propagation so that all intermediate bounds are
        # concretized.
        _, env = self._forward_algorithm.propagate(graph, inputs)
        env = dict(env)  # take a modifiable copy

        # Iterate over the outputs, computing each one separately.
        outvals = []
        for out_var in graph.outputs:
            concrete_outvar_bound = self._backward_concretizer.concrete_bound(
                graph, inputs, env, out_var
            )
            outvals.append(concrete_outvar_bound)
            env[out_var] = concrete_outvar_bound
        return outvals, env
