# 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/losses/single_index.py
# class SingleIndexLossFn(typing_extensions.Protocol):
# class L2Loss(SingleIndexLossFn):
# class XentLoss(SingleIndexLossFn):
# class AccuracyErrorLoss(SingleIndexLossFn):
# class ElboLoss(SingleIndexLossFn):
# class VaeLoss(SingleIndexLossFn):
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
# def average_single_index_loss(single_loss: SingleIndexLossFn,
#                               num_index_samples: int = 1) -> base.LossFn:
#   def loss_fn(enn: base.EpistemicNetwork,
#               params: hk.Params,
#               batch: base.Batch,
#               key: base.RngKey) -> Tuple[base.Array, base.LossMetrics]:
# def add_data_noise(single_loss: SingleIndexLossFn,
#                    noise_fn: data_noise.DataNoise) -> SingleIndexLossFn:
#   def noisy_loss(apply: base.ApplyFn,
#                  params: hk.Params,
#                  batch: base.Batch,
#                  index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __post_init__(self):
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:


import dataclasses
import haiku as hk
import jax
import jax.numpy as jnp
from typing import Any, Dict, Callable, List, Optional, Sequence, Tuple, Union
from enn import base
from enn.losses import single_index

def l2_weights_with_predicate(
    params: hk.Params,
    predicate: Optional[PredicateFn] = None) -> jnp.DeviceArray:
  """Sum of squares of parameter weights that passes predicate_fn."""
  if predicate is not None:
    params = hk.data_structures.filter(predicate, params)
  return sum(jnp.sum(jnp.square(p)) for p in jax.tree_leaves(params))


def add_l2_weight_decay(loss_fn: base.LossFn,
                        scale: Union[float, Callable[[hk.Params], hk.Params]],
                        predicate: Optional[PredicateFn] = None) -> base.LossFn:
  """Adds scale * l2 weight decay to an existing loss function."""
  try:  # Scale is numeric.
    scale = jnp.sqrt(scale)
    scale_fn = lambda ps: jax.tree_map(lambda p: scale * p, ps)
  except TypeError:
    scale_fn = scale  # Assuming scale is a Callable.
  def new_loss(a, params: hk.Params, c, d) -> base.Array:
    loss, metrics = loss_fn(a, params, c, d)
    decay = l2_weights_with_predicate(scale_fn(params), predicate)
    total_loss = loss +  decay
    metrics['decay'] = decay
    metrics['raw_loss'] = loss
    return total_loss, metrics
  return new_loss


def combine_single_index_losses_as_metric(

# 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/losses/single_index.py
# class SingleIndexLossFn(typing_extensions.Protocol):
# class L2Loss(SingleIndexLossFn):
# class XentLoss(SingleIndexLossFn):
# class AccuracyErrorLoss(SingleIndexLossFn):
# class ElboLoss(SingleIndexLossFn):
# class VaeLoss(SingleIndexLossFn):
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
# def average_single_index_loss(single_loss: SingleIndexLossFn,
#                               num_index_samples: int = 1) -> base.LossFn:
#   def loss_fn(enn: base.EpistemicNetwork,
#               params: hk.Params,
#               batch: base.Batch,
#               key: base.RngKey) -> Tuple[base.Array, base.LossMetrics]:
# def add_data_noise(single_loss: SingleIndexLossFn,
#                    noise_fn: data_noise.DataNoise) -> SingleIndexLossFn:
#   def noisy_loss(apply: base.ApplyFn,
#                  params: hk.Params,
#                  batch: base.Batch,
#                  index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __post_init__(self):
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:
#   def __call__(self,
#                apply: base.ApplyFn,
#                params: hk.Params,
#                batch: base.Batch,
#                index: base.Index) -> Tuple[base.Array, base.LossMetrics]:


# Path: enn/losses/utils.py
import dataclasses
import haiku as hk
import jax
import jax.numpy as jnp
from typing import Any, Dict, Callable, List, Optional, Sequence, Tuple, Union
from enn import base
from enn.losses import single_index


def l2_weights_with_predicate(
    params: hk.Params,
    predicate: Optional[PredicateFn] = None) -> jnp.DeviceArray:
  """Sum of squares of parameter weights that passes predicate_fn."""
  if predicate is not None:
    params = hk.data_structures.filter(predicate, params)
  return sum(jnp.sum(jnp.square(p)) for p in jax.tree_leaves(params))


def add_l2_weight_decay(loss_fn: base.LossFn,
                        scale: Union[float, Callable[[hk.Params], hk.Params]],
                        predicate: Optional[PredicateFn] = None) -> base.LossFn:
  """Adds scale * l2 weight decay to an existing loss function."""
  try:  # Scale is numeric.
    scale = jnp.sqrt(scale)
    scale_fn = lambda ps: jax.tree_map(lambda p: scale * p, ps)
  except TypeError:
    scale_fn = scale  # Assuming scale is a Callable.
  def new_loss(a, params: hk.Params, c, d) -> base.Array:
    loss, metrics = loss_fn(a, params, c, d)
    decay = l2_weights_with_predicate(scale_fn(params), predicate)
    total_loss = loss +  decay
    metrics['decay'] = decay
    metrics['raw_loss'] = loss
    return total_loss, metrics
  return new_loss


def combine_single_index_losses_as_metric(

    train_loss: single_index.SingleIndexLossFn,