# 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
# class GPRegression:
# class TestbedGPRegression(testbed_base.TestbedProblem):
#   def __init__(self,
#                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):
#   def x_test(self) -> chex.Array:
#   def test_mean(self) -> chex.Array:
#   def test_cov(self) -> chex.Array:
#   def train_data(self) -> testbed_base.Data:
#   def train_data(self) -> testbed_base.Data:
#   def prior_knowledge(self) -> testbed_base.PriorKnowledge:
#   def evaluate_quality(
#       self,
#       enn_sampler: testbed_base.EpistemicSampler) -> testbed_base.ENNQuality:
# def _kl_gaussian(
#     mean_1: float, std_1: float, mean_2: float, std_2: float) -> float:


import dataclasses
import chex
import haiku as hk
import jax
import jax.config
import numpy as np
import typing_extensions
from typing import Tuple
from enn.experiments.neurips_2021 import base as testbed_base
from enn.experiments.neurips_2021 import testbed
from neural_tangents import stax
from neural_tangents.utils import typing as nt_types
  return kernel_ctor(input_dim)


def gaussian_data(key: chex.PRNGKey,
                  num_train: int,
                  input_dim: int,
                  num_test: int) -> Tuple[chex.Array, chex.Array]:
  """Generate Gaussian training and test data."""
  train_key, test_key = jax.random.split(key)
  x_train = jax.random.normal(train_key, [num_train, input_dim])
  x_test = jax.random.normal(test_key, [num_test, input_dim])
  return x_train, x_test


@dataclasses.dataclass
class RegressionTestbedConfig:
  """Configuration options for regression testbed instance."""
  num_train: int
  input_dim: int
  seed: int
  noise_std: float
  tau: int = 1  # TODO(author2): Consider design of this parameter.
  num_test_cache: int = 1000
  target_test_seeds: int = 1000  # num_test_seeds = target_test_seeds / tau
  num_enn_samples: int = 100
  kernel_ctor: KernelCtor = make_benchmark_kernel
  num_layers: int = 1  # Output to prior knowledge


def regression_load_from_config(

# 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
# class GPRegression:
# class TestbedGPRegression(testbed_base.TestbedProblem):
#   def __init__(self,
#                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):
#   def x_test(self) -> chex.Array:
#   def test_mean(self) -> chex.Array:
#   def test_cov(self) -> chex.Array:
#   def train_data(self) -> testbed_base.Data:
#   def train_data(self) -> testbed_base.Data:
#   def prior_knowledge(self) -> testbed_base.PriorKnowledge:
#   def evaluate_quality(
#       self,
#       enn_sampler: testbed_base.EpistemicSampler) -> testbed_base.ENNQuality:
# def _kl_gaussian(
#     mean_1: float, std_1: float, mean_2: float, std_2: float) -> float:


# Path: enn/experiments/neurips_2021/load.py
import dataclasses
import chex
import haiku as hk
import jax
import jax.config
import numpy as np
import typing_extensions
from typing import Tuple
from enn.experiments.neurips_2021 import base as testbed_base
from enn.experiments.neurips_2021 import testbed
from neural_tangents import stax
from neural_tangents.utils import typing as nt_types

  return kernel_ctor(input_dim)


def gaussian_data(key: chex.PRNGKey,
                  num_train: int,
                  input_dim: int,
                  num_test: int) -> Tuple[chex.Array, chex.Array]:
  """Generate Gaussian training and test data."""
  train_key, test_key = jax.random.split(key)
  x_train = jax.random.normal(train_key, [num_train, input_dim])
  x_test = jax.random.normal(test_key, [num_test, input_dim])
  return x_train, x_test


@dataclasses.dataclass
class RegressionTestbedConfig:
  """Configuration options for regression testbed instance."""
  num_train: int
  input_dim: int
  seed: int
  noise_std: float
  tau: int = 1  # TODO(author2): Consider design of this parameter.
  num_test_cache: int = 1000
  target_test_seeds: int = 1000  # num_test_seeds = target_test_seeds / tau
  num_enn_samples: int = 100
  kernel_ctor: KernelCtor = make_benchmark_kernel
  num_layers: int = 1  # Output to prior knowledge


def regression_load_from_config(

    config: RegressionTestbedConfig) -> testbed_base.TestbedProblem: