# Copyright 2018 Guanshuo Wang. 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.
# ==============================================================================

# Reference:
# @inproceedings{shufflenetv2,
#   title={ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design},
#   author={Ningning Ma and Xiangyu Zhang and Hai-Tao Zheng and Jian Sun},
#   booktitle={ECCV},
#   year={2018}
# }

from collections import OrderedDict

import tensorflow as tf
from tensorflow.contrib import layers
from tensorflow.contrib.framework import arg_scope
from model_parallel import *

import net_base

class ShuffleNet_v2_small(net_base.Network):
  def __init__(self, 
               alpha=1.0, 
               se=False,
               residual=False,
               weight_decay=0.0005, 
               data_format='NCHW', 
               name='ShuffleNet_v2_small'):
    super(ShuffleNet_v2_small, self).__init__(weight_decay, data_format, name)
    if alpha == 0.5:
      self.num_outputs = [24, 48, 96, 1024]
      self.name += '_x0_5'
    elif alpha == 1.0:
      self.num_outputs = [58, 116, 232, 1024]
    elif alpha == 1.5:
      self.num_outputs = [88, 176, 352, 1024]
      self.name += '_x1_5'
    elif alpha == 2.0:
      self.num_outputs = [122, 244, 488, 2048]
      self.name += '_x2'
    self.se = se
    if se:
      self.name += '_se'
    self.residual = residual
    if residual:
      self.name += '_res'

  def _channel_split(self, x):
    num_channels = x.get_shape().as_list()[self.channel_axis]
    split_size = int(0.5*num_channels)
    shortcut, x = tf.split(x, [split_size, num_channels-split_size], axis=self.channel_axis)
    return shortcut, x

  def _channel_shuffle(self, x):
    input_height, input_width = x.get_shape().as_list()[2:4] if self.data_format=='NCHW' else x.get_shape().as_list()[1:3]
    num_channels = x.get_shape().as_list()[self.channel_axis]
    if self.data_format == 'NCHW':
      x = tf.reshape(x, [-1, 2, num_channels/2, input_height, input_width])
      x = tf.transpose(x, [0, 2, 1, 3, 4])
      x = tf.reshape(x, [-1, num_channels, input_height, input_width])
    else:
      x = tf.reshape(x, [-1, input_height, input_width, num_channels/2, 2])
      x = tf.transpose(x, [0, 1, 2, 4, 3])
      x = tf.reshape(x, [-1, input_height, input_width, num_channels])
    return x

  def _squeeze_excitation(self, x):
    num_channels = x.get_shape().as_list()[self.channel_axis]
    w = tf.reduce_mean(x, axis=self.spatial_axis, keepdims=True)
    w = layers.conv2d(w, num_channels/2, kernel_size=1, normalizer_fn=None)
    w = layers.conv2d(w, num_channels, kernel_size=1, normalizer_fn=None, activation_fn=tf.nn.sigmoid)
    x = x*w
    return x

  def separable_resBlock(self, x, 
                         num_outputs, 
                         stride=1, 
                         activation_fn=tf.nn.relu, 
                         normalizer_fn=layers.batch_norm,
                         scope=None):
    residual_flag = self.residual and (stride == 1 and num_outputs == x.get_shape().as_list()[self.channel_axis]) 
    with tf.variable_scope(scope, 'resBlock'):
      # channel_split
      shortcut, x = self._channel_split(x)
      if stride != 1:
        shortcut = layers.separable_conv2d(shortcut, num_outputs, kernel_size=3, stride=stride, 
                                           scope='separable_conv_shortcut_3x3')
        shortcut = layers.conv2d(shortcut, num_outputs, kernel_size=1, stride=1, scope='conv_shortcut_1x1')
      if residual_flag:
        res_shortcut = x
      x = layers.conv2d(x, num_outputs, kernel_size=1, stride=1, scope='conv1_1x1',)
      x = layers.separable_conv2d(x, num_outputs, kernel_size=3, stride=stride, scope='separable_conv2_3x3')
      x = layers.conv2d(x, num_outputs, kernel_size=1, stride=1, scope='conv3_1x1')
      if self.se:
        x = self._squeeze_excitation(x)      