# coding=utf-8
# Copyright 2022 The Edward2 Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# 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.

# Lint as: python3
"""ResNet-32x4 with rank-1 distributions."""
import functools
import os
import time
from absl import app
from absl import flags
from absl import logging

from experimental.rank1_bnns import resnet_cifar_model  # local file import
import robustness_metrics as rm
import tensorflow as tf
import tensorflow_datasets as tfds
import uncertainty_baselines as ub
from uncertainty_baselines.baselines.cifar import utils
import uncertainty_metrics as um

# ~24.4 steps per epoch for 4x4 TPU; per_core_batch_size=64; 300 epochs;
# so 2/3 of training time.
flags.DEFINE_integer('kl_annealing_steps', int(24.4 * 200),
                     'Number of steps over which to anneal the KL term to 1.')
flags.DEFINE_string('alpha_initializer', 'trainable_deterministic',
                    'Initializer name for the alpha parameters.')
flags.DEFINE_string('gamma_initializer', 'trainable_deterministic',
                    'Initializer name for the gamma parameters.')
flags.DEFINE_string('alpha_regularizer', None,
                    'Regularizer name for the alpha parameters.')
flags.DEFINE_string('gamma_regularizer', None,
                    'Regularizer name for the gamma parameters.')
flags.DEFINE_boolean('use_additive_perturbation', False,
                     'Use additive perturbations instead of multiplicative.')
flags.DEFINE_integer('num_train_samples', 1,
                     'Number of samples per example during training.')
flags.DEFINE_integer('num_eval_samples', 1,
                     'Number of samples per example during evaluation.')

# General model flags
flags.DEFINE_integer('ensemble_size', 4, 'Size of ensemble.')
flags.DEFINE_bool(
    'member_sampling', default=False,
    help=('Whether or not to sample a single ensemble member per step with '
          'which to compute the loss and derivatives.'))
flags.DEFINE_integer('per_core_batch_size', 64, 'Batch size per TPU core/GPU.')
flags.DEFINE_float('random_sign_init', -0.5,
                   'Use random sign init for fast weights.')
flags.DEFINE_integer('seed', 0, 'Random seed.')
flags.DEFINE_float('fast_weight_lr_multiplier', 0.5,
                   'fast weights lr multiplier.')
flags.DEFINE_bool('version2', True, 'Use ensemble version2.')
flags.DEFINE_bool(
    'expected_probs', default=False,
    help=('Whether or not to compute the loss over the per-example average of '
          'the predicted probabilities across the ensemble members.'))
flags.DEFINE_float('base_learning_rate', 0.1,
                   'Base learning rate when total training batch size is 128.')
flags.DEFINE_integer('lr_warmup_epochs', 1,
                     'Number of epochs for a linear warmup to the initial '
                     'learning rate. Use 0 to do no warmup.')
flags.DEFINE_float('lr_decay_ratio', 0.1, 'Amount to decay learning rate.')
flags.DEFINE_list('lr_decay_epochs', ['80', '160', '180'],
                  'Epochs to decay learning rate by.')
flags.DEFINE_float('dropout_rate', 0.,
                   'Dropout rate. Only used if alpha/gamma initializers are, '
                   'e.g., trainable normal with a fixed stddev.')
flags.DEFINE_float('l2', 2e-4, 'L2 coefficient.')
flags.DEFINE_enum('dataset', 'cifar10',
                  enum_values=['cifar10', 'cifar100'],
                  help='Dataset.')
# TODO(ghassen): consider adding CIFAR-100-C to TFDS.
flags.DEFINE_string('cifar100_c_path',
                    '',
                    'Path to the TFRecords files for CIFAR-100-C. Only valid '
                    '(and required) if dataset is cifar100 and corruptions.')
flags.DEFINE_integer('corruptions_interval', 250,
                     'Number of epochs between evaluating on the corrupted '
                     'test data. Use -1 to never evaluate.')
flags.DEFINE_integer('checkpoint_interval', 25,
                     'Number of epochs between saving checkpoints. Use -1 to '
                     'never save checkpoints.')
flags.DEFINE_integer('num_bins', 15, 'Number of bins for ECE.')
flags.DEFINE_string('output_dir', '/tmp/cifar',
                    'The directory where the model weights and '
                    'training/evaluation summaries are stored.')
flags.DEFINE_integer('train_epochs', 250, 'Number of training epochs.')

# Accelerator flags.
flags.DEFINE_bool('use_gpu', False, 'Whether to run on GPU or otherwise TPU.')
flags.DEFINE_bool('use_bfloat16', False, 'Whether to use mixed precision.')
flags.DEFINE_integer('num_cores', 8, 'Number of TPU cores or number of GPUs.')
flags.DEFINE_string('tpu', None,
                    'Name of the TPU. Only used if use_gpu is False.')
FLAGS = flags.FLAGS


def main(argv):
  del argv  # Unused arg.

  tf.random.set_seed(FLAGS.seed)

  if FLAGS.version2:
    per_core_bs_train = FLAGS.per_core_batch_size // (FLAGS.ensemble_size *
                                                      FLAGS.num_train_samples)
    per_core_bs_eval = FLAGS.per_core_batch_size // (FLAGS.ensemble_size *
                                                     FLAGS.num_eval_samples)
  else:
    per_core_bs_train = FLAGS.per_core_batch_size // FLAGS.num_train_samples
    per_core_bs_eval = FLAGS.per_core_batch_size // FLAGS.num_eval_samples
  batch_size_train = per_core_bs_train * FLAGS.num_cores
  batch_size_eval = per_core_bs_eval * FLAGS.num_cores

  logging.info('Saving checkpoints at %s', FLAGS.output_dir)

  if FLAGS.use_gpu:
    logging.info('Use GPU')
    strategy = tf.distribute.MirroredStrategy()
  else:
    logging.info('Use TPU at %s',
                 FLAGS.tpu if FLAGS.tpu is not None else 'local')