# 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 haiku as hk
import haiku.experimental as hke
import jax
import jax.numpy as jnp
  """Get variables for gaussian enn linear layer."""
  # TODO(author2): How to import hk._src.base types correctly?
  if _is_linear_bias(context):
    standard_bias = next_getter(value[0])
    gaussian_bias = next_getter(value[1])
    noise = jax.random.normal(hk.next_rng_key(), standard_bias.shape)
    return standard_bias + gaussian_bias * noise
  else:
    return next_getter(value)


class GaussianNoiseEnn(base.EpistemicNetwork):
  """GaussianNoiseEnn from callable module."""

  def __init__(self,
               module_ctor: Callable[[], hk.Module],
               init_scale: float = 1.):
    """GaussianNoiseEnn from callable module."""
    enn_creator = make_enn_creator(init_scale=init_scale)
    def net_fn(inputs: base.Array) -> base.Array:
      with hke.custom_getter(enn_getter), hke.custom_creator(enn_creator):
        output = module_ctor()(inputs)  # pytype: disable=not-callable
        return output

    # TODO(author2): Note that the GaussianENN requires a rng_key in place of an
    # index. Therefore we do *not* hk.without_apply_rng.
    transformed = hk.transform(net_fn)
    super().__init__(
        apply=lambda params, x, z: transformed.apply(params, z, x),
        init=lambda rng, x, z: transformed.init(rng, 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/gaussian_enn.py
from typing import Callable, Sequence
from enn import base
from enn.networks import indexers
import haiku as hk
import haiku.experimental as hke
import jax
import jax.numpy as jnp

  """Get variables for gaussian enn linear layer."""
  # TODO(author2): How to import hk._src.base types correctly?
  if _is_linear_bias(context):
    standard_bias = next_getter(value[0])
    gaussian_bias = next_getter(value[1])
    noise = jax.random.normal(hk.next_rng_key(), standard_bias.shape)
    return standard_bias + gaussian_bias * noise
  else:
    return next_getter(value)


class GaussianNoiseEnn(base.EpistemicNetwork):
  """GaussianNoiseEnn from callable module."""

  def __init__(self,
               module_ctor: Callable[[], hk.Module],
               init_scale: float = 1.):
    """GaussianNoiseEnn from callable module."""
    enn_creator = make_enn_creator(init_scale=init_scale)
    def net_fn(inputs: base.Array) -> base.Array:
      with hke.custom_getter(enn_getter), hke.custom_creator(enn_creator):
        output = module_ctor()(inputs)  # pytype: disable=not-callable
        return output

    # TODO(author2): Note that the GaussianENN requires a rng_key in place of an
    # index. Therefore we do *not* hk.without_apply_rng.
    transformed = hk.transform(net_fn)
    super().__init__(
        apply=lambda params, x, z: transformed.apply(params, z, x),
        init=lambda rng, x, z: transformed.init(rng, x),

        indexer=indexers.PrngIndexer(),