# Path: enn/experiments/neurips_2021/agent_factories.py
# def make_ensemble_ctor(num_ensemble: int,
#                        noise_scale: float,
#                        prior_scale: float,
#                        hidden_size: int = 50,
#                        num_layers: int = 2,
#                        seed: int = 0) -> ConfigCtor:
#   def make_enn(prior: testbed_base.PriorKnowledge) -> enn_base.EpistemicNetwork:
#   def make_agent_config() -> agents.VanillaEnnConfig:
# def make_dropout_ctor(dropout_rate: float,
#                       regularization_scale: float,
#                       hidden_size: int = 50,
#                       num_layers: int = 2,
#                       dropout_input: bool = True,
#                       regularization_tau: float = 1,
#                       seed: int = 0) -> ConfigCtor:
#   def make_enn(prior: testbed_base.PriorKnowledge) -> enn_base.EpistemicNetwork:
#   def make_agent_config() -> agents.VanillaEnnConfig:
# def make_hypermodel_ctor(index_dim: int,
#                          noise_scale: float,
#                          prior_scale: float,
#                          hidden_size: int = 50,
#                          num_layers: int = 2,
#                          seed: int = 0) -> ConfigCtor:
#   def make_enn(prior: testbed_base.PriorKnowledge) -> enn_base.EpistemicNetwork:
#   def make_agent_config() -> agents.VanillaEnnConfig:
# def make_bbb_ctor(sigma_0: float,
#                   learning_rate: float,
#                   hidden_size: int = 50,
#                   num_layers: int = 2,
#                   num_index_samples: int = 64,
#                   seed: int = 0) -> ConfigCtor:
#   def make_enn(prior: testbed_base.PriorKnowledge) -> enn_base.EpistemicNetwork:
#   def make_agent_config() -> agents.VanillaEnnConfig:
# def make_ensemble_sweep() -> List[AgentCtorConfig]:
# def make_dropout_sweep() -> List[AgentCtorConfig]:
# def make_hypermodel_sweep() -> List[AgentCtorConfig]:
# def make_bbb_sweep() -> List[AgentCtorConfig]:
# def make_agent_sweep(agent: str = 'all') -> Sequence[AgentCtorConfig]:
# def load_agent_config(agent_id: int,
#                       agent: str = 'all') -> agents.VanillaEnnConfig:
# def xm_agent_sweep(agent: str = 'all') -> Sequence[int]:
# def join_metadata(df: pd.DataFrame) -> pd.DataFrame:
# class AgentCtorConfig:
# 
# Path: enn/experiments/neurips_2021/agents.py
# class VanillaEnnConfig:
# class VanillaEnnAgent(testbed_base.TestbedAgent):
# def extract_enn_sampler(
#     experiment: supervised.Experiment) -> testbed_base.EpistemicSampler:
#   def enn_sampler(x: enn_base.Array, seed: int = 0) -> enn_base.Array:
#   def __call__(
#       self,
#       data: testbed_base.Data,
#       prior: testbed_base.PriorKnowledge,
#   ) -> testbed_base.EpistemicSampler:
# def _round_to_integer(x: float) -> int:
# def logging_freq(num_batches: int,
#                  num_points: int = 100,
#                  log_freq: Optional[int] = None) -> int:
# 
# Path: enn/experiments/neurips_2021/load.py
# class KernelCtor(typing_extensions.Protocol):
# class MLPKernelCtor(KernelCtor):
# class RegressionTestbedConfig:
#   def __call__(self, input_dim: int) -> nt_types.KernelFn:
#   def __post_init__(self):
#   def __call__(self, input_dim: int = 1) -> nt_types.KernelFn:
# def make_benchmark_kernel(input_dim: int = 1) -> nt_types.KernelFn:
# def gaussian_data(key: chex.PRNGKey,
#                   num_train: int,
#                   input_dim: int,
#                   num_test: int) -> Tuple[chex.Array, chex.Array]:
# def regression_load_from_config(
#     config: RegressionTestbedConfig) -> testbed_base.TestbedProblem:
# def regression_load(input_dim: int,
#                     data_ratio: float,
#                     seed: int,
#                     noise_std: float) -> testbed_base.TestbedProblem:


from absl import app
from absl import flags
from enn.experiments.neurips_2021 import agent_factories
from enn.experiments.neurips_2021 import agents
from enn.experiments.neurips_2021 import load
from jax.config import config

# Double-precision in JAX helps with numerical stability
config.update('jax_enable_x64', True)

# GP configuration
flags.DEFINE_integer('input_dim', 1, 'Input dimension.')
flags.DEFINE_float('data_ratio', 1., 'Ratio of num_train to input_dim.')
flags.DEFINE_float('noise_std', 0.1, 'Additive noise standard deviation.')
flags.DEFINE_integer('seed', 1, 'Seed for testbed problem.')


# ENN agent
flags.DEFINE_integer('agent_id', 0, 'Which agent id')
flags.DEFINE_enum('agent', 'all',
                  ['all', 'ensemble', 'dropout', 'hypermodel', 'bbb'],
                  'Which agent family.')

FLAGS = flags.FLAGS


def main(_):
  # Load the appropriate testbed problem
  problem = load.regression_load(
      input_dim=FLAGS.input_dim,
      data_ratio=FLAGS.data_ratio,
      seed=FLAGS.seed,
      noise_std=FLAGS.noise_std,
  )

  # Form the appropriate agent for training

# Path: enn/experiments/neurips_2021/agent_factories.py
# def make_ensemble_ctor(num_ensemble: int,
#                        noise_scale: float,
#                        prior_scale: float,
#                        hidden_size: int = 50,
#                        num_layers: int = 2,
#                        seed: int = 0) -> ConfigCtor:
#   def make_enn(prior: testbed_base.PriorKnowledge) -> enn_base.EpistemicNetwork:
#   def make_agent_config() -> agents.VanillaEnnConfig:
# def make_dropout_ctor(dropout_rate: float,
#                       regularization_scale: float,
#                       hidden_size: int = 50,
#                       num_layers: int = 2,
#                       dropout_input: bool = True,
#                       regularization_tau: float = 1,
#                       seed: int = 0) -> ConfigCtor:
#   def make_enn(prior: testbed_base.PriorKnowledge) -> enn_base.EpistemicNetwork:
#   def make_agent_config() -> agents.VanillaEnnConfig:
# def make_hypermodel_ctor(index_dim: int,
#                          noise_scale: float,
#                          prior_scale: float,
#                          hidden_size: int = 50,
#                          num_layers: int = 2,
#                          seed: int = 0) -> ConfigCtor:
#   def make_enn(prior: testbed_base.PriorKnowledge) -> enn_base.EpistemicNetwork:
#   def make_agent_config() -> agents.VanillaEnnConfig:
# def make_bbb_ctor(sigma_0: float,
#                   learning_rate: float,
#                   hidden_size: int = 50,
#                   num_layers: int = 2,
#                   num_index_samples: int = 64,
#                   seed: int = 0) -> ConfigCtor:
#   def make_enn(prior: testbed_base.PriorKnowledge) -> enn_base.EpistemicNetwork:
#   def make_agent_config() -> agents.VanillaEnnConfig:
# def make_ensemble_sweep() -> List[AgentCtorConfig]:
# def make_dropout_sweep() -> List[AgentCtorConfig]:
# def make_hypermodel_sweep() -> List[AgentCtorConfig]:
# def make_bbb_sweep() -> List[AgentCtorConfig]:
# def make_agent_sweep(agent: str = 'all') -> Sequence[AgentCtorConfig]:
# def load_agent_config(agent_id: int,
#                       agent: str = 'all') -> agents.VanillaEnnConfig:
# def xm_agent_sweep(agent: str = 'all') -> Sequence[int]:
# def join_metadata(df: pd.DataFrame) -> pd.DataFrame:
# class AgentCtorConfig:
# 
# Path: enn/experiments/neurips_2021/agents.py
# class VanillaEnnConfig:
# class VanillaEnnAgent(testbed_base.TestbedAgent):
# def extract_enn_sampler(
#     experiment: supervised.Experiment) -> testbed_base.EpistemicSampler:
#   def enn_sampler(x: enn_base.Array, seed: int = 0) -> enn_base.Array:
#   def __call__(
#       self,
#       data: testbed_base.Data,
#       prior: testbed_base.PriorKnowledge,
#   ) -> testbed_base.EpistemicSampler:
# def _round_to_integer(x: float) -> int:
# def logging_freq(num_batches: int,
#                  num_points: int = 100,
#                  log_freq: Optional[int] = None) -> int:
# 
# Path: enn/experiments/neurips_2021/load.py
# class KernelCtor(typing_extensions.Protocol):
# class MLPKernelCtor(KernelCtor):
# class RegressionTestbedConfig:
#   def __call__(self, input_dim: int) -> nt_types.KernelFn:
#   def __post_init__(self):
#   def __call__(self, input_dim: int = 1) -> nt_types.KernelFn:
# def make_benchmark_kernel(input_dim: int = 1) -> nt_types.KernelFn:
# def gaussian_data(key: chex.PRNGKey,
#                   num_train: int,
#                   input_dim: int,
#                   num_test: int) -> Tuple[chex.Array, chex.Array]:
# def regression_load_from_config(
#     config: RegressionTestbedConfig) -> testbed_base.TestbedProblem:
# def regression_load(input_dim: int,
#                     data_ratio: float,
#                     seed: int,
#                     noise_std: float) -> testbed_base.TestbedProblem:


# Path: enn/experiments/neurips_2021/run_testbed.py
from absl import app
from absl import flags
from enn.experiments.neurips_2021 import agent_factories
from enn.experiments.neurips_2021 import agents
from enn.experiments.neurips_2021 import load
from jax.config import config


# Double-precision in JAX helps with numerical stability
config.update('jax_enable_x64', True)

# GP configuration
flags.DEFINE_integer('input_dim', 1, 'Input dimension.')
flags.DEFINE_float('data_ratio', 1., 'Ratio of num_train to input_dim.')
flags.DEFINE_float('noise_std', 0.1, 'Additive noise standard deviation.')
flags.DEFINE_integer('seed', 1, 'Seed for testbed problem.')


# ENN agent
flags.DEFINE_integer('agent_id', 0, 'Which agent id')
flags.DEFINE_enum('agent', 'all',
                  ['all', 'ensemble', 'dropout', 'hypermodel', 'bbb'],
                  'Which agent family.')

FLAGS = flags.FLAGS


def main(_):
  # Load the appropriate testbed problem
  problem = load.regression_load(
      input_dim=FLAGS.input_dim,
      data_ratio=FLAGS.data_ratio,
      seed=FLAGS.seed,
      noise_std=FLAGS.noise_std,
  )

  # Form the appropriate agent for training

  agent_config = agent_factories.load_agent_config(FLAGS.agent_id, FLAGS.agent)