# 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/thompson.py
# class ThompsonEnnBandit:
#   def __init__(
#       self,
#       enn_config: agents.VanillaEnnConfig,
#       input_dim: int,
#       noise_std: float,
#       num_actions: int,
#       steps_per_obs: int = 1,
#       kernel_ctor: load.KernelCtor = load.make_benchmark_kernel,
#       kernel_ridge: float = 1e-3,
#       logger: Optional[loggers.Logger] = None,
#       batch_size: int = 8,
#       seed: int = 0,
#   ):
#     def forward(params: hk.Params,
#                 inputs: enn_base.Array,
#                 key: enn_base.RngKey) -> enn_base.Array:
#     def sgd_step(
#         params: hk.Params,
#         opt_state: optax.OptState,
#         batch: enn_base.Batch,
#         key: enn_base.RngKey
#     ) -> Tuple[hk.Params, optax.OptState]:
#     def select_action(params: hk.Params,
#                       key: enn_base.RngKey) -> Dict[str, enn_base.Array]:
#     def make_batch(
#         replay: Sequence[Dict[str, enn_base.Array]])-> enn_base.Batch:
#   def step(self) -> float:
#   def run(self, num_steps: int, log_freq: int = 1):


from absl import app
from absl import flags
from enn.experiments.neurips_2021 import agent_factories
from enn.experiments.neurips_2021 import thompson
from jax.config import config
# 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.
# ============================================================================
"""Example running an ENN on Thompson bandit task."""


# 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('action_ratio', 10., 'Ratio of actions to input_dim.')
flags.DEFINE_float('noise_std', 0.1, 'Additive noise standard deviation.')
flags.DEFINE_integer('seed', 1, 'Seed for testbed problem.')
flags.DEFINE_integer('num_steps', 100_000, 'Number of steps to run experiment.')


# 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(_):
  # 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/thompson.py
# class ThompsonEnnBandit:
#   def __init__(
#       self,
#       enn_config: agents.VanillaEnnConfig,
#       input_dim: int,
#       noise_std: float,
#       num_actions: int,
#       steps_per_obs: int = 1,
#       kernel_ctor: load.KernelCtor = load.make_benchmark_kernel,
#       kernel_ridge: float = 1e-3,
#       logger: Optional[loggers.Logger] = None,
#       batch_size: int = 8,
#       seed: int = 0,
#   ):
#     def forward(params: hk.Params,
#                 inputs: enn_base.Array,
#                 key: enn_base.RngKey) -> enn_base.Array:
#     def sgd_step(
#         params: hk.Params,
#         opt_state: optax.OptState,
#         batch: enn_base.Batch,
#         key: enn_base.RngKey
#     ) -> Tuple[hk.Params, optax.OptState]:
#     def select_action(params: hk.Params,
#                       key: enn_base.RngKey) -> Dict[str, enn_base.Array]:
#     def make_batch(
#         replay: Sequence[Dict[str, enn_base.Array]])-> enn_base.Batch:
#   def step(self) -> float:
#   def run(self, num_steps: int, log_freq: int = 1):


# Path: enn/experiments/neurips_2021/run_thompson.py
from absl import app
from absl import flags
from enn.experiments.neurips_2021 import agent_factories
from enn.experiments.neurips_2021 import thompson
from jax.config import config

# 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.
# ============================================================================
"""Example running an ENN on Thompson bandit task."""


# 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('action_ratio', 10., 'Ratio of actions to input_dim.')
flags.DEFINE_float('noise_std', 0.1, 'Additive noise standard deviation.')
flags.DEFINE_integer('seed', 1, 'Seed for testbed problem.')
flags.DEFINE_integer('num_steps', 100_000, 'Number of steps to run experiment.')


# 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(_):
  # Form the appropriate agent for training

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