# 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:


from typing import Callable, Sequence
from enn import base
from enn.networks import indexers
import chex
import haiku as hk
import jax
import jax.numpy as jnp

  Args:
    output_sizes: Sequence of integer sizes for the MLPs.
    num_ensemble: Integer number of elements in the ensemble.
    nonzero_bias: Whether to make the initial layer bias nonzero.
    activation: Jax callable defining activation per layer.
  Returns:
    EpistemicNetwork as an ensemble of MLP.
  """
  def ensemble_forward(x: base.Array) -> base.OutputWithPrior:
    """Forwards the entire ensemble at given input x."""
    model = EnsembleMLP(output_sizes, num_ensemble, nonzero_bias, activation)
    return model(x)

  transformed = hk.without_apply_rng(hk.transform(ensemble_forward))

  # Apply function selects the appropriate index of the ensemble output.
  def apply(params: hk.Params,
            x: base.Array,
            z: base.Index) -> base.OutputWithPrior:
    net_out = transformed.apply(params, x)
    one_hot_index = jax.nn.one_hot(z, num_ensemble)
    return jnp.dot(net_out, one_hot_index)

  def init(key: base.RngKey,
           x: base.Array,
           z: base.Index) -> hk.Params:
    del z
    return transformed.init(key, x)


# 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/einsum_mlp.py
from typing import Callable, Sequence
from enn import base
from enn.networks import indexers
import chex
import haiku as hk
import jax
import jax.numpy as jnp


  Args:
    output_sizes: Sequence of integer sizes for the MLPs.
    num_ensemble: Integer number of elements in the ensemble.
    nonzero_bias: Whether to make the initial layer bias nonzero.
    activation: Jax callable defining activation per layer.
  Returns:
    EpistemicNetwork as an ensemble of MLP.
  """
  def ensemble_forward(x: base.Array) -> base.OutputWithPrior:
    """Forwards the entire ensemble at given input x."""
    model = EnsembleMLP(output_sizes, num_ensemble, nonzero_bias, activation)
    return model(x)

  transformed = hk.without_apply_rng(hk.transform(ensemble_forward))

  # Apply function selects the appropriate index of the ensemble output.
  def apply(params: hk.Params,
            x: base.Array,
            z: base.Index) -> base.OutputWithPrior:
    net_out = transformed.apply(params, x)
    one_hot_index = jax.nn.one_hot(z, num_ensemble)
    return jnp.dot(net_out, one_hot_index)

  def init(key: base.RngKey,
           x: base.Array,
           z: base.Index) -> hk.Params:
    del z
    return transformed.init(key, x)


  indexer = indexers.EnsembleIndexer(num_ensemble)