# 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/utils.py
# FLAGS = flags.FLAGS
# def epistemic_network_from_module(
#     enn_ctor: Callable[[], base.EpistemicModule],
#     indexer: base.EpistemicIndexer,
# ) -> base.EpistemicNetwork:
#   def enn_fn(inputs: base.Array, index: base.Index) -> base.Output:
# def wrap_transformed_as_enn(
#     transformed: hk.Transformed) -> base.EpistemicNetwork:
# def wrap_enn_as_enn_with_state(
#     enn: base.EpistemicNetwork) -> base.EpistemicNetworkWithState:
#   def init(key, x, index):
#   def apply(params, unused_state, x, index):
# def make_centered_enn(enn: base.EpistemicNetwork,
#                       x_train: base.Array) -> base.EpistemicNetwork:
#   def centered_apply(params: hk.Params,
#                      x: base.Array,
#                      z: base.Index) -> base.Output:
# def parse_net_output(net_out: base.Output) -> base.Array:
# def make_batch_indexer(indexer: base.EpistemicIndexer,
#                        batch_size: int) -> base.EpistemicIndexer:
#   def batch_indexer(key: base.RngKey) -> base.Index:
# def _clean_batch_data(data: base.Batch) -> base.Batch:
# def make_batch_iterator(data: base.Batch,
#                         batch_size: Optional[int] = None,
#                         seed: int = 0) -> base.BatchIterator:
# def make_test_data(n_samples: int = 20) -> base.BatchIterator:
# 
# 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, Type
from enn import base
from enn import utils
from enn.networks import priors
import chex
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np




# TODO(author2): Current implementation will produce a hypermodel that can
# *only* work for a single index at a time. However, note that jax.vmap means
# this can easily be converted into a form that works with batched index.


class MLPHypermodel(base.EpistemicNetwork):
  """MLP hypermodel for transformed_base as EpistemicNetwork."""

  def __init__(self,
               transformed_base: hk.Transformed,
               dummy_input: base.Array,
               indexer: base.EpistemicIndexer,
               hidden_sizes: Optional[Sequence[int]] = None,
               return_generated_params: bool = False,
               scale: bool = True,
               w_init: Optional[hk.initializers.Initializer] = None,
               b_init: Optional[hk.initializers.Initializer] = None,
               ):
    """MLP hypermodel for transformed_base as EpistemicNetwork."""

    if hidden_sizes is None:
      hyper_torso = lambda x: x
    else:
      def hyper_torso(index):
        return hk.nets.MLP(hidden_sizes, w_init=w_init, b_init=b_init)(index)


# 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/utils.py
# FLAGS = flags.FLAGS
# def epistemic_network_from_module(
#     enn_ctor: Callable[[], base.EpistemicModule],
#     indexer: base.EpistemicIndexer,
# ) -> base.EpistemicNetwork:
#   def enn_fn(inputs: base.Array, index: base.Index) -> base.Output:
# def wrap_transformed_as_enn(
#     transformed: hk.Transformed) -> base.EpistemicNetwork:
# def wrap_enn_as_enn_with_state(
#     enn: base.EpistemicNetwork) -> base.EpistemicNetworkWithState:
#   def init(key, x, index):
#   def apply(params, unused_state, x, index):
# def make_centered_enn(enn: base.EpistemicNetwork,
#                       x_train: base.Array) -> base.EpistemicNetwork:
#   def centered_apply(params: hk.Params,
#                      x: base.Array,
#                      z: base.Index) -> base.Output:
# def parse_net_output(net_out: base.Output) -> base.Array:
# def make_batch_indexer(indexer: base.EpistemicIndexer,
#                        batch_size: int) -> base.EpistemicIndexer:
#   def batch_indexer(key: base.RngKey) -> base.Index:
# def _clean_batch_data(data: base.Batch) -> base.Batch:
# def make_batch_iterator(data: base.Batch,
#                         batch_size: Optional[int] = None,
#                         seed: int = 0) -> base.BatchIterator:
# def make_test_data(n_samples: int = 20) -> base.BatchIterator:
# 
# 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/hypermodels.py
from typing import Callable, Optional, Sequence, Type
from enn import base
from enn import utils
from enn.networks import priors
import chex
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np





# TODO(author2): Current implementation will produce a hypermodel that can
# *only* work for a single index at a time. However, note that jax.vmap means
# this can easily be converted into a form that works with batched index.


class MLPHypermodel(base.EpistemicNetwork):
  """MLP hypermodel for transformed_base as EpistemicNetwork."""

  def __init__(self,
               transformed_base: hk.Transformed,
               dummy_input: base.Array,
               indexer: base.EpistemicIndexer,
               hidden_sizes: Optional[Sequence[int]] = None,
               return_generated_params: bool = False,
               scale: bool = True,
               w_init: Optional[hk.initializers.Initializer] = None,
               b_init: Optional[hk.initializers.Initializer] = None,
               ):
    """MLP hypermodel for transformed_base as EpistemicNetwork."""

    if hidden_sizes is None:
      hyper_torso = lambda x: x
    else:
      def hyper_torso(index):
        return hk.nets.MLP(hidden_sizes, w_init=w_init, b_init=b_init)(index)


    enn = utils.epistemic_network_from_module(