# Path: enn/experiments/neurips_2021/base.py
# class Data(NamedTuple):
# class PriorKnowledge:
# class ENNQuality:
# class EpistemicSampler(typing_extensions.Protocol):
# class TestbedAgent(typing_extensions.Protocol):
# class TestbedProblem(abc.ABC):
#   def __call__(self, x: chex.Array, seed: int = 0) -> chex.Array:
#   def __call__(self,
#                data: Data,
#                prior: Optional[PriorKnowledge] = None) -> EpistemicSampler:
#   def train_data(self) -> Data:
#   def evaluate_quality(self, enn_sampler: EpistemicSampler) -> ENNQuality:
#   def prior_knowledge(self) -> PriorKnowledge:


import dataclasses
import chex
import haiku as hk
import jax
import jax.numpy as jnp
import neural_tangents as nt
from enn.experiments.neurips_2021 import base as testbed_base
from neural_tangents.utils import typing as nt_types
               kernel_fn: nt_types.KernelFn,
               x_train: chex.Array,
               x_test: chex.Array,
               key: chex.PRNGKey,
               tau: int = 1,
               noise_std: float = 1,
               kernel_ridge: float = 1e-6):

    # Checking the dimensionality of our data coming in.
    num_train, input_dim = x_train.shape
    num_test_x_cache, input_dim_test = x_test.shape
    assert input_dim == input_dim_test

    rng = hk.PRNGSequence(key)
    self._tau = tau
    self._input_dim = input_dim
    self._x_train = jnp.array(x_train)
    self._x_test = jnp.array(x_test)
    self._num_train = num_train
    self._num_test_x_cache = num_test_x_cache
    self._noise_std = noise_std
    self._kernel_ridge = kernel_ridge

    # Form the training data
    mean = jnp.zeros(num_train)
    k_train_train = kernel_fn(self._x_train, x2=None, get='nngp')
    k_train_train += kernel_ridge * jnp.eye(num_train)
    y_function = jax.random.multivariate_normal(next(rng), mean, k_train_train)
    y_noise = jax.random.normal(next(rng), [num_train, 1]) * noise_std
    y_train = y_function[:, None] + y_noise

# Path: enn/experiments/neurips_2021/base.py
# class Data(NamedTuple):
# class PriorKnowledge:
# class ENNQuality:
# class EpistemicSampler(typing_extensions.Protocol):
# class TestbedAgent(typing_extensions.Protocol):
# class TestbedProblem(abc.ABC):
#   def __call__(self, x: chex.Array, seed: int = 0) -> chex.Array:
#   def __call__(self,
#                data: Data,
#                prior: Optional[PriorKnowledge] = None) -> EpistemicSampler:
#   def train_data(self) -> Data:
#   def evaluate_quality(self, enn_sampler: EpistemicSampler) -> ENNQuality:
#   def prior_knowledge(self) -> PriorKnowledge:


# Path: enn/experiments/neurips_2021/testbed.py
import dataclasses
import chex
import haiku as hk
import jax
import jax.numpy as jnp
import neural_tangents as nt
from enn.experiments.neurips_2021 import base as testbed_base
from neural_tangents.utils import typing as nt_types

               kernel_fn: nt_types.KernelFn,
               x_train: chex.Array,
               x_test: chex.Array,
               key: chex.PRNGKey,
               tau: int = 1,
               noise_std: float = 1,
               kernel_ridge: float = 1e-6):

    # Checking the dimensionality of our data coming in.
    num_train, input_dim = x_train.shape
    num_test_x_cache, input_dim_test = x_test.shape
    assert input_dim == input_dim_test

    rng = hk.PRNGSequence(key)
    self._tau = tau
    self._input_dim = input_dim
    self._x_train = jnp.array(x_train)
    self._x_test = jnp.array(x_test)
    self._num_train = num_train
    self._num_test_x_cache = num_test_x_cache
    self._noise_std = noise_std
    self._kernel_ridge = kernel_ridge

    # Form the training data
    mean = jnp.zeros(num_train)
    k_train_train = kernel_fn(self._x_train, x2=None, get='nngp')
    k_train_train += kernel_ridge * jnp.eye(num_train)
    y_function = jax.random.multivariate_normal(next(rng), mean, k_train_train)
    y_noise = jax.random.normal(next(rng), [num_train, 1]) * noise_std
    y_train = y_function[:, None] + y_noise

    self._train_data = testbed_base.Data(x_train, y_train)