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


from typing import Callable
from enn import base
from enn import utils
from tensorflow_probability.substrates import jax as tfp
import chex
import haiku as hk
import jax
import jax.numpy as jnp
# 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.
# ============================================================================
"""Tools for computing VI loss."""


tfd = tfp.distributions


def get_awgn_loglike_fn(sigma_w: float
                        ) -> Callable[[base.Output, base.Batch], float]:
  """Returns a function that computes the simple unnormalized log likelihood.

  It assumes response variable is perturbed with additive iid Gaussian noise.

  Args:
    sigma_w: standard deviation of the additive Gaussian noise.

  Returns:
    A function that computes the log likelihood given data and output.
  """
  def log_likelihood_fn(out: base.Output, batch: base.Batch):
    chex.assert_shape(batch.y, (None, 1))

# 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/losses/vi_losses.py
from typing import Callable
from enn import base
from enn import utils
from tensorflow_probability.substrates import jax as tfp
import chex
import haiku as hk
import jax
import jax.numpy as jnp

# 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.
# ============================================================================
"""Tools for computing VI loss."""


tfd = tfp.distributions


def get_awgn_loglike_fn(sigma_w: float
                        ) -> Callable[[base.Output, base.Batch], float]:
  """Returns a function that computes the simple unnormalized log likelihood.

  It assumes response variable is perturbed with additive iid Gaussian noise.

  Args:
    sigma_w: standard deviation of the additive Gaussian noise.

  Returns:
    A function that computes the log likelihood given data and output.
  """
  def log_likelihood_fn(out: base.Output, batch: base.Batch):
    chex.assert_shape(batch.y, (None, 1))

    err_sq = jnp.mean(jnp.square(utils.parse_net_output(out) - batch.y))