# Path: enn/base.py
# class OutputWithPrior(NamedTuple):
# class EpistemicModule(abc.ABC, hk.Module):
# class ApplyFn(typing_extensions.Protocol):
# class InitFn(typing_extensions.Protocol):
# class EpistemicIndexer(typing_extensions.Protocol):
# class EpistemicNetwork:
# class Batch(NamedTuple):
# class LossFn(typing_extensions.Protocol):
# class EpistemicNetworkWithState:
# class LossFnWithState(typing_extensions.Protocol):
#   def preds(self) -> Array:
#   def __call__(self, inputs: Array, index: Index) -> Output:
#   def __call__(self, params: hk.Params, inputs: Array, index: Index) -> Output:
#   def __call__(self, rng_key: RngKey, inputs: Array, index: Index) -> hk.Params:
#   def __call__(self, key: RngKey) -> Index:
#   def __call__(self,
#                enn: EpistemicNetwork,
#                params: hk.Params,
#                batch: Batch,
#                key: RngKey) -> Tuple[Array, LossMetrics]:
#   def __call__(self,
#                enn: EpistemicNetworkWithState,
#                params: hk.Params,
#                state: hk.State,
#                batch: Batch,
#                key: RngKey) -> Tuple[Array, Tuple[hk.State, LossMetrics]]:
# 
# Path: enn/networks/indexers.py
# class PrngIndexer(base.EpistemicIndexer):
# class EnsembleIndexer(base.EpistemicIndexer):
# class GaussianIndexer(base.EpistemicIndexer):
# class ScaledGaussianIndexer(base.EpistemicIndexer):
# class GaussianWithUnitIndexer(base.EpistemicIndexer):
# class DirichletIndexer(base.EpistemicIndexer):
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def mean_index(self) -> base.Array:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
# 
# Path: enn/networks/priors.py
# class EnnWithAdditivePrior(base.EpistemicNetwork):
# class NetworkWithAdditivePrior(hk.Module):
#   def __init__(self,
#                enn: base.EpistemicNetwork,
#                prior_fn: PriorFn,
#                prior_scale: float = 1.):
#     def apply_fn(params: hk.Params,
#                  inputs: base.Array,
#                  index: base.Index) -> base.OutputWithPrior:
# def convert_enn_to_prior_fn(enn: base.EpistemicNetwork,
#                             dummy_input: base.Array,
#                             key: base.RngKey) -> PriorFn:
#   def prior_fn(x: base.Array, z: base.Index) -> base.Output:
# def make_null_prior(output_dim: int) -> Callable[[base.Array], base.Array]:
#   def null_prior(inputs: base.Array) -> base.Array:
# def _parse_gamma(gamma: GpGamma,
#                  num_feat: int,
#                  key: base.RngKey) -> Union[float, base.Array]:
# def make_random_feat_gp(
#     input_dim: int,
#     output_dim: int,
#     num_feat: int,
#     key: base.RngKey,
#     gamma: GpGamma = 1.,
#     scale: float = 1.,
# ) -> Callable[[base.Array], base.Array]:
#   def gp_instance(inputs: base.Array) -> base.Array:
# def get_random_mlp_with_index(
#     x_sample: base.Array,
#     z_sample: base.Array,
#     rng: chex.PRNGKey,
#     prior_output_sizes: Optional[Iterable[int]] = None,
#     prior_weight_std: float = 3,
#     prior_bias_std: float = 0.1
# ) -> PriorFn:
#   def net_fn(x: base.Array, z: base.Array):
#   def __call__(self, *args, **kwargs) -> base.Array:
# WARN = ('WARNING: prior parameters will be included as hk.Params for module.'
#         'If possible, you should use EnnWithAdditivePrior instead.')


from typing import Callable, Optional, Sequence, Tuple
from enn import base
from enn.networks import indexers
from enn.networks import priors
import chex
import haiku as hk
import jax
import jax.numpy as jnp
# ============================================================================

"""Implementing some types of ENN ensembles in JAX."""



class Ensemble(base.EpistemicNetwork):
  """Ensemble ENN that uses a dot product in param space.

  Repeats parameters by an additional *ensemble* dimension in axis=0.
  Applying the parameters selects one single component of parameters per index.
  """

  def __init__(self,
               model: hk.Transformed,
               num_ensemble: int):
    self.model = model
    self.num_ensemble = num_ensemble

    def init(key: chex.PRNGKey, inputs: chex.Array, index: int) -> hk.Params:
      del index  # Unused
      batched_init = jax.vmap(model.init, in_axes=[0, None], out_axes=0)
      return batched_init(jax.random.split(key, num_ensemble), inputs)

    def apply(params: hk.Params, inputs: chex.Array, index: int) -> base.Output:
      one_hot_index = jax.nn.one_hot(index, num_ensemble)
      param_selector = lambda p: jnp.einsum('i...,i->...', p, one_hot_index)
      sub_params = jax.tree_map(param_selector, params)
      return model.apply(sub_params, inputs)


# Path: enn/base.py
# class OutputWithPrior(NamedTuple):
# class EpistemicModule(abc.ABC, hk.Module):
# class ApplyFn(typing_extensions.Protocol):
# class InitFn(typing_extensions.Protocol):
# class EpistemicIndexer(typing_extensions.Protocol):
# class EpistemicNetwork:
# class Batch(NamedTuple):
# class LossFn(typing_extensions.Protocol):
# class EpistemicNetworkWithState:
# class LossFnWithState(typing_extensions.Protocol):
#   def preds(self) -> Array:
#   def __call__(self, inputs: Array, index: Index) -> Output:
#   def __call__(self, params: hk.Params, inputs: Array, index: Index) -> Output:
#   def __call__(self, rng_key: RngKey, inputs: Array, index: Index) -> hk.Params:
#   def __call__(self, key: RngKey) -> Index:
#   def __call__(self,
#                enn: EpistemicNetwork,
#                params: hk.Params,
#                batch: Batch,
#                key: RngKey) -> Tuple[Array, LossMetrics]:
#   def __call__(self,
#                enn: EpistemicNetworkWithState,
#                params: hk.Params,
#                state: hk.State,
#                batch: Batch,
#                key: RngKey) -> Tuple[Array, Tuple[hk.State, LossMetrics]]:
# 
# Path: enn/networks/indexers.py
# class PrngIndexer(base.EpistemicIndexer):
# class EnsembleIndexer(base.EpistemicIndexer):
# class GaussianIndexer(base.EpistemicIndexer):
# class ScaledGaussianIndexer(base.EpistemicIndexer):
# class GaussianWithUnitIndexer(base.EpistemicIndexer):
# class DirichletIndexer(base.EpistemicIndexer):
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def mean_index(self) -> base.Array:
#   def __call__(self, key: base.RngKey) -> base.Index:
#   def __call__(self, key: base.RngKey) -> base.Index:
# 
# Path: enn/networks/priors.py
# class EnnWithAdditivePrior(base.EpistemicNetwork):
# class NetworkWithAdditivePrior(hk.Module):
#   def __init__(self,
#                enn: base.EpistemicNetwork,
#                prior_fn: PriorFn,
#                prior_scale: float = 1.):
#     def apply_fn(params: hk.Params,
#                  inputs: base.Array,
#                  index: base.Index) -> base.OutputWithPrior:
# def convert_enn_to_prior_fn(enn: base.EpistemicNetwork,
#                             dummy_input: base.Array,
#                             key: base.RngKey) -> PriorFn:
#   def prior_fn(x: base.Array, z: base.Index) -> base.Output:
# def make_null_prior(output_dim: int) -> Callable[[base.Array], base.Array]:
#   def null_prior(inputs: base.Array) -> base.Array:
# def _parse_gamma(gamma: GpGamma,
#                  num_feat: int,
#                  key: base.RngKey) -> Union[float, base.Array]:
# def make_random_feat_gp(
#     input_dim: int,
#     output_dim: int,
#     num_feat: int,
#     key: base.RngKey,
#     gamma: GpGamma = 1.,
#     scale: float = 1.,
# ) -> Callable[[base.Array], base.Array]:
#   def gp_instance(inputs: base.Array) -> base.Array:
# def get_random_mlp_with_index(
#     x_sample: base.Array,
#     z_sample: base.Array,
#     rng: chex.PRNGKey,
#     prior_output_sizes: Optional[Iterable[int]] = None,
#     prior_weight_std: float = 3,
#     prior_bias_std: float = 0.1
# ) -> PriorFn:
#   def net_fn(x: base.Array, z: base.Array):
#   def __call__(self, *args, **kwargs) -> base.Array:
# WARN = ('WARNING: prior parameters will be included as hk.Params for module.'
#         'If possible, you should use EnnWithAdditivePrior instead.')


# Path: enn/networks/ensembles.py
from typing import Callable, Optional, Sequence, Tuple
from enn import base
from enn.networks import indexers
from enn.networks import priors
import chex
import haiku as hk
import jax
import jax.numpy as jnp

# ============================================================================

"""Implementing some types of ENN ensembles in JAX."""



class Ensemble(base.EpistemicNetwork):
  """Ensemble ENN that uses a dot product in param space.

  Repeats parameters by an additional *ensemble* dimension in axis=0.
  Applying the parameters selects one single component of parameters per index.
  """

  def __init__(self,
               model: hk.Transformed,
               num_ensemble: int):
    self.model = model
    self.num_ensemble = num_ensemble

    def init(key: chex.PRNGKey, inputs: chex.Array, index: int) -> hk.Params:
      del index  # Unused
      batched_init = jax.vmap(model.init, in_axes=[0, None], out_axes=0)
      return batched_init(jax.random.split(key, num_ensemble), inputs)

    def apply(params: hk.Params, inputs: chex.Array, index: int) -> base.Output:
      one_hot_index = jax.nn.one_hot(index, num_ensemble)
      param_selector = lambda p: jnp.einsum('i...,i->...', p, one_hot_index)
      sub_params = jax.tree_map(param_selector, params)
      return model.apply(sub_params, inputs)


    indexer = indexers.EnsembleIndexer(num_ensemble)