# 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.
# ==============================================================================

"""Utility functions for setting up the CMP graph.
"""

import os, numpy as np
import matplotlib.pyplot as plt


import tensorflow as tf

from tensorflow.contrib import slim
from tensorflow.contrib.slim import arg_scope
import logging
from src import utils
import src.file_utils as fu
from tfcode import tf_utils

resnet_v2 = tf_utils.resnet_v2
custom_residual_block = tf_utils.custom_residual_block

def value_iteration_network(
    fr, num_iters, val_neurons, action_neurons, kernel_size, share_wts=False,
    name='vin', wt_decay=0.0001, activation_fn=None, shape_aware=False):
  """
  Constructs a Value Iteration Network, convolutions and max pooling across
  channels.
  Input:
    fr:             NxWxHxC
    val_neurons:    Number of channels for maintaining the value.
    action_neurons: Computes action_neurons * val_neurons at each iteration to
                    max pool over.
  Output:
    value image:  NxHxWx(val_neurons)
  """
  init_var = np.sqrt(2.0/(kernel_size**2)/(val_neurons*action_neurons))
  vals = []
  with tf.variable_scope(name) as varscope:
    if shape_aware == False:
      fr_shape = tf.unstack(tf.shape(fr))
      val_shape = tf.stack(fr_shape[:-1] + [val_neurons])
      val = tf.zeros(val_shape, name='val_init')
    else:
      val = tf.expand_dims(tf.zeros_like(fr[:,:,:,0]), dim=-1) * \
          tf.constant(0., dtype=tf.float32, shape=[1,1,1,val_neurons])
      val_shape = tf.shape(val)
    vals.append(val)
    for i in range(num_iters):
      if share_wts:
        # The first Value Iteration maybe special, so it can have its own
        # paramterss.
        scope = 'conv'
        if i == 0: scope = 'conv_0'
        if i > 1: varscope.reuse_variables()
      else:
        scope = 'conv_{:d}'.format(i)
      val = slim.conv2d(tf.concat([val, fr], 3, name='concat_{:d}'.format(i)),
                        num_outputs=action_neurons*val_neurons,
                        kernel_size=kernel_size, stride=1, activation_fn=activation_fn,
                        scope=scope, normalizer_fn=None,
                        weights_regularizer=slim.l2_regularizer(wt_decay),
                        weights_initializer=tf.random_normal_initializer(stddev=init_var),
                        biases_initializer=tf.zeros_initializer())
      val = tf.reshape(val, [-1, action_neurons*val_neurons, 1, 1],
                       name='re_{:d}'.format(i))
      val = slim.max_pool2d(val, kernel_size=[action_neurons,1],
                            stride=[action_neurons,1], padding='VALID',
                            scope='val_{:d}'.format(i))
      val = tf.reshape(val, val_shape, name='unre_{:d}'.format(i))
      vals.append(val)
  return val, vals


def rotate_preds(loc_on_map, relative_theta, map_size, preds,
                 output_valid_mask):
  with tf.name_scope('rotate'):
    flow_op = tf_utils.get_flow(loc_on_map, relative_theta, map_size=map_size)
    if type(preds) != list:
      rotated_preds, valid_mask_warps = tf_utils.dense_resample(preds, flow_op,
                                                                output_valid_mask)
    else:
      rotated_preds = [] ;valid_mask_warps = []
      for pred in preds:
        rotated_pred, valid_mask_warp = tf_utils.dense_resample(pred, flow_op,
                                                                output_valid_mask)
        rotated_preds.append(rotated_pred)
        valid_mask_warps.append(valid_mask_warp)
  return rotated_preds, valid_mask_warps

def get_visual_frustum(map_size, shape_like, expand_dims=[0,0]):
  with tf.name_scope('visual_frustum'):
    l = np.tril(np.ones(map_size)) ;l = l + l[:,::-1]
    l = (l == 2).astype(np.float32)
    for e in expand_dims:
      l = np.expand_dims(l, axis=e)
    confs_probs = tf.constant(l, dtype=tf.float32)
    confs_probs = tf.ones_like(shape_like, dtype=tf.float32) * confs_probs
  return confs_probs

def deconv(x, is_training, wt_decay, neurons, strides, layers_per_block,
            kernel_size, conv_fn, name, offset=0):
  """Generates a up sampling network with residual connections. 
  """
  batch_norm_param = {'center': True, 'scale': True,
                      'activation_fn': tf.nn.relu,
                      'is_training': is_training}
  outs = []
  for i, (neuron, stride) in enumerate(zip(neurons, strides)):
    for s in range(layers_per_block):
      scope = '{:s}_{:d}_{:d}'.format(name, i+1+offset,s+1)
      x = custom_residual_block(x, neuron, kernel_size, stride, scope,
                                is_training, wt_decay, use_residual=True,
                                residual_stride_conv=True, conv_fn=conv_fn,
                                batch_norm_param=batch_norm_param)
      stride = 1
    outs.append((x,True))
  return x, outs

def fr_v2(x, output_neurons, inside_neurons, is_training, name='fr',
          wt_decay=0.0001, stride=1, updates_collections=tf.GraphKeys.UPDATE_OPS):
  """Performs fusion of information between the map and the reward map.
  Inputs
    x:   NxHxWxC1

  Outputs
    fr map:     NxHxWx(output_neurons)
  """
  if type(stride) != list:
    stride = [stride]
  with slim.arg_scope(resnet_v2.resnet_utils.resnet_arg_scope(
      is_training=is_training, weight_decay=wt_decay)):
    with slim.arg_scope([slim.batch_norm], updates_collections=updates_collections) as arg_sc:
      # Change the updates_collections for the conv normalizer_params to None
      for i in range(len(arg_sc.keys())):
        if 'convolution' in arg_sc.keys()[i]:
          arg_sc.values()[i]['normalizer_params']['updates_collections'] = updates_collections
      with slim.arg_scope(arg_sc):
        bottleneck = resnet_v2.bottleneck
        blocks = []
        for i, s in enumerate(stride):
          b = resnet_v2.resnet_utils.Block(
              'block{:d}'.format(i + 1), bottleneck, [{
                  'depth': output_neurons,
                  'depth_bottleneck': inside_neurons,
                  'stride': stride[i]
              }])
          blocks.append(b)
        x, outs = resnet_v2.resnet_v2(x, blocks, num_classes=None, global_pool=False,
                                     output_stride=None, include_root_block=False,
                                     reuse=False, scope=name)
  return x, outs