#!/usr/bin/python
#
# Copyright 2022 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.

"""Coupling flows for particle models."""

import functools
from typing import Any, Callable, Mapping, Sequence, Union

import chex
import distrax
from flows_for_atomic_solids.models import bijectors
from flows_for_atomic_solids.models import embeddings
from flows_for_atomic_solids.models import utils
import haiku as hk
import jax
import jax.numpy as jnp
from tensorflow_probability.substrates import jax as tfp

Array = chex.Array
PRNGKey = Array


def _reshape_last(x: Array, ndims: int, new_shape: Sequence[int]) -> Array:
  """Reshapes the last `ndims` dimensions of `x` to shape `new_shape`."""
  if ndims <= 0:
    raise ValueError(
        f'Number of dimensions to reshape must be positive, got {ndims}.')
  return jnp.reshape(x, x.shape[:-ndims] + tuple(new_shape))


def make_equivariant_conditioner(
    shape_transformed: Sequence[int],
    num_bijector_params: int,
    lower: float,
    upper: float,
    embedding_size: int,
    conditioner_constructor: Callable[..., Any],
    conditioner_kwargs: Mapping[str, Any],
    num_frequencies: int,
    w_init_final: hk.initializers.Initializer = jnp.zeros,
) -> Callable[[Array], Array]:
  """Make a permutation-equivariant conditioner for the coupling flow."""
  # This conditioner assumes that the input is of shape [..., N, D1]. It returns
  # an output of shape [..., N, D2, K], where:
  #   D2 = `shape_transformed[-1]`
  #   K = `num_bijector_params`
  output_size = shape_transformed[-1]
  conditioner = conditioner_constructor(**conditioner_kwargs)
  return hk.Sequential([
      functools.partial(
          embeddings.circular, lower=lower, upper=upper,
          num_frequencies=num_frequencies),
      hk.Linear(embedding_size),
      conditioner,
      hk.Linear(output_size * num_bijector_params, w_init=w_init_final),
      functools.partial(
          _reshape_last, ndims=1, new_shape=(output_size, num_bijector_params)),
  ])


def make_split_coupling_flow(
    event_shape: Sequence[int],
    lower: float,
    upper: float,
    num_layers: int,
    num_bins: int,
    conditioner: Mapping[str, Any],
    permute_variables: bool,
    split_axis: int,
    use_circular_shift: bool,
    prng: Union[int, PRNGKey],
    circular_shift_init: hk.initializers.Initializer = jnp.zeros,
    ) -> distrax.Bijector:
  """Create a flow that consists of a sequence of split coupling layers.

  All coupling layers use rational-quadratic splines. Each layer of the flow
  is composed of two split coupling bijectors, where each coupling bijector
  transforms a different part of the input.

  The flow maps to and from the range `[lower, upper]`, obeying periodic
  boundary conditions.

  Args:
    event_shape: the shape of the event generated by the flow. Does not
      include the batch dimensions.
    lower: lower range of the flow.
    upper: upper range of the flow.
    num_layers: the number of layers to use. Each layer consists of two split
      coupling bijectors, where each coupling bijector transforms a different
      part of the input.
    num_bins: number of bins to use in the rational-quadratic splines.
    conditioner: a Mapping containing 'constructor' and 'kwargs' keys that
      configures the conditioner used in the coupling layers.
    permute_variables: whether to permute the dimensions along the splitting
      axis between successive layers.
    split_axis: a negative int that defines which axis to split along.
    use_circular_shift: if True, add a learned circular shift between successive
      flow layers.
    prng: either a PRNG key, or an integer seed to convert to a PRNG key. The
      PRNG key will be used to construct the permutations, when these are
      random.
    circular_shift_init: initializer for the circular shifts.

  Returns:
    The flow, a Distrax bijector.
  """
  if isinstance(prng, int):
    prng = jax.random.PRNGKey(prng)

  if split_axis >= 0:
    raise ValueError(f'Expected split axis to be negative, got {split_axis}.')

  def bijector_fn(params: Array):
    return distrax.RationalQuadraticSpline(
        params,
        range_min=lower,
        range_max=upper,
        boundary_slopes='circular',
        min_bin_size=(upper - lower) * 1e-4)

  split_size = event_shape[split_axis]
  split_index = split_size // 2
  cyclic_permutation = jnp.roll(jnp.arange(split_size), shift=1)
  block_ndims = len(event_shape) + split_axis
  layers = []
  for _ in range(num_layers):
    sublayers = []

    # Variable permutation.
    if permute_variables:
      if split_size <= 3:
        permutation = cyclic_permutation
      else:
        prng, key = jax.random.split(prng)
        permutation = jax.random.permutation(key, jnp.arange(split_size))
      permute_layer = tfp.bijectors.Permute(permutation, axis=split_axis)
      permute_layer = distrax.Block(permute_layer, block_ndims)
      sublayers.append(permute_layer)

    # Circular shift.
    if use_circular_shift:
      shift = utils.Parameter(
          name='circular_shift',
          param_name='shift',
          shape=event_shape[-1:],
          init=circular_shift_init)()
      shift_layer = bijectors.CircularShift(
          (upper - lower) * shift, lower, upper)
      shift_layer = distrax.Block(shift_layer, len(event_shape))
      sublayers.append(shift_layer)

    # Coupling layers.
    for swap in [True, False]:
      shape_transformed = list(event_shape)
      shape_transformed[split_axis] = (
          split_index if swap else split_size - split_index)
      coupling_layer = distrax.SplitCoupling(
          swap=swap,
          split_index=split_index,
          split_axis=split_axis,
          event_ndims=len(event_shape),
          bijector=bijector_fn,
          conditioner=conditioner['constructor'](
              shape_transformed=shape_transformed,
              num_bijector_params=3 * num_bins + 1,
              lower=lower,
              upper=upper,
              **conditioner['kwargs']))
      sublayers.append(coupling_layer)
    layers.append(distrax.Chain(sublayers))

  return distrax.Chain(layers)
