# Author: Mengye Ren (mren@cs.toronto.edu).
#
# Modified from Tensorflow original code.
# Original Tensorflow license shown below.
# =============================================================================
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# =============================================================================
"""ResNet model.

Related papers:
https://arxiv.org/pdf/1603.05027v2.pdf
https://arxiv.org/pdf/1512.03385v1.pdf
https://arxiv.org/pdf/1605.07146v1.pdf
"""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import numpy as np
import tensorflow as tf

from resnet.models import nnlib as nn
from resnet.utils import logger

log = logger.get()


class ResNetModel(object):
  """ResNet model."""

  def __init__(self,
               config,
               is_training=True,
               inference_only=False,
               inp=None,
               label=None):
    """ResNet constructor.

    Args:
      config: Hyperparameters.
      is_training: One of "train" and "eval".
      inference_only: Do not build optimizer.
    """
    self._config = config
    self._l1_collection = []
    self.is_training = is_training

    # Input.
    if inp is None:
      x = tf.placeholder(
          tf.float32, [None, config.height, config.width, config.num_channel])
    else:
      x = inp

    if label is None:
      y = tf.placeholder(tf.int32, [None])
    else:
      y = label

    logits = self.build_inference_network(x)
    predictions = tf.nn.softmax(logits)

    with tf.variable_scope("costs"):
      xent = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,labels= y)
      xent = tf.reduce_sum(xent, name="xent") / tf.to_float(tf.shape(x)[0])
      cost = xent
      cost += self._decay()
      cost += self._l1_loss()

    self._cost = cost
    self._input = x
    self._label = y
    self._cross_ent = xent
    self._output = predictions

    if not is_training or inference_only:
      return

    global_step = tf.Variable(0.0, name="global_step", trainable=False)
    lr = tf.Variable(0.0, name="learn_rate", trainable=False)
    trainable_variables = tf.trainable_variables()
    grads = tf.gradients(cost, trainable_variables)
    if config.optimizer == "sgd":
      optimizer = tf.train.GradientDescentOptimizer(lr)
    elif config.optimizer == "mom":
      optimizer = tf.train.MomentumOptimizer(lr, 0.9)
    train_op = optimizer.apply_gradients(
        zip(grads, trainable_variables),
        global_step=global_step,
        name="train_step")
    self._train_op = train_op
    self._global_step = global_step
    self._lr = lr
    self._new_lr = tf.placeholder(
        tf.float32, shape=[], name="new_learning_rate")
    self._lr_update = tf.assign(self._lr, self._new_lr)

  def assign_lr(self, session, lr_value):
    """Assigns new learning rate."""
    log.info("Adjusting learning rate to {}".format(lr_value))
    session.run(self._lr_update, feed_dict={self._new_lr: lr_value})

  @property
  def cost(self):
    return self._cost

  @property
  def train_op(self):
    return self._train_op

  @property
  def config(self):
    return self._config

  @property
  def lr(self):
    return self._lr

  @property
  def input(self):
    return self._input

  @property
  def output(self):
    return self._output

  @property
  def label(self):
    return self._label

  @property
  def cross_ent(self):
    return self._cross_ent

  @property
  def global_step(self):
    return self._global_step

  @property
  def l1_collection(self):
    return self._l1_collection

  def build_inference_network(self, x):
    config = self.config
    is_training = self.is_training
    num_stages = len(self.config.num_residual_units)
    strides = config.strides
    activate_before_residual = config.activate_before_residual
    filters = [ff for ff in config.filters]  # Copy filter config.
    init_filter = config.init_filter

    with tf.variable_scope("init"):
      h = self._conv("init_conv", x, init_filter,
                     int(x.get_shape()[-1]), filters[0],
                     self._stride_arr(config.init_stride))

      # Max-pooling is used in ImageNet experiments to further reduce
      # dimensionality.
      if config.init_max_pool:
        h = tf.nn.max_pool(h, [1, 3, 3, 1], [1, 2, 2, 1], "SAME")

    if config.use_bottleneck:
      res_func = self._bottleneck_residual
      # For CIFAR-10 it's [16, 16, 32, 64] => [16, 64, 128, 256]
      for ii in range(1, len(filters)):
        filters[ii] *= 4
    else:
      res_func = self._residual

    for ss in range(num_stages):
      with tf.variable_scope("unit_{}_0".format(ss + 1)):
        h = res_func(
            h,
            filters[ss],
            filters[ss + 1],
            self._stride_arr(strides[ss]),
            activate_before_residual=activate_before_residual[ss])
      for ii in range(1, config.num_residual_units[ss]):
        with tf.variable_scope("unit_{}_{}".format(ss + 1, ii)):
          h = res_func(
              h,
              filters[ss + 1],
              filters[ss + 1],
              self._stride_arr(1),
              activate_before_residual=False)

    with tf.variable_scope("unit_last"):
      h = self._batch_norm("final_bn", h)
      h = self._relu(h, config.relu_leakiness)

    h = self._global_avg_pool(h)

    with tf.variable_scope("logit"):
      logits = self._fully_connected(h, config.num_classes)

    return logits

  def _stride_arr(self, stride):
    """Map a stride scalar to the stride array for tf.nn.conv2d."""