from typing import Optional
from typing import MutableMapping
from typing import Any


def _get_classifier_parameters(
    num_gpus: int = 0,
    builder: str = 'records',
    skip_eval: bool = False,
    distribution_strategy: str = 'mirrored',
    per_replica_batch_size: int = 128,
    epochs: int = 90,
    steps: int = 0,
    epochs_between_evals: int = 1,
    dtype: str = 'float32',
    enable_xla: bool = False,
    run_eagerly: bool = False,
    gpu_thread_mode: Optional[str] = None,
    dataset_num_private_threads: Optional[int] = None,
    loss_scale: Optional[str] = None,
    batchnorm_spatial_persistent: bool = False) -> MutableMapping[str, Any]:
  """Gets classifier trainer's ResNet parameters."""
  return {
      'runtime': {
          'num_gpus': num_gpus,
          'distribution_strategy': distribution_strategy,
          'run_eagerly': run_eagerly,
          'enable_xla': enable_xla,
          'dataset_num_private_threads': dataset_num_private_threads,
          'gpu_thread_mode': gpu_thread_mode,
          'loss_scale': loss_scale,
          'batchnorm_spatial_persistent': batchnorm_spatial_persistent,
      },
      'train_dataset': {
          'builder': builder,
          'use_per_replica_batch_size': True,
          'batch_size': per_replica_batch_size,
          'image_size': 224,
          'dtype': dtype,
      },
      'validation_dataset': {
          'builder': builder,
          'batch_size': per_replica_batch_size,
          'use_per_replica_batch_size': True,
          'image_size': 224,
          'dtype': dtype,
      },
      'train': {
          'epochs': epochs,
          'steps': steps,
          'callbacks': {
              'enable_tensorboard': False,
              'enable_checkpoint_and_export': False,
              'enable_time_history': True,
          },
      },
      'model': {
          'loss': {
              'label_smoothing': 0.1,
          },
      },
      'evaluation': {
          'epochs_between_evals': epochs_between_evals,
          'skip_eval': skip_eval,
      },
  }