# Copyright 2015 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.
# ==============================================================================
"""Keras training and evaluation routines.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import copy
import multiprocessing
import threading
import time

import numpy as np
import six

from tensorflow.contrib.keras.python.keras import backend as K
from tensorflow.contrib.keras.python.keras import callbacks as cbks
from tensorflow.contrib.keras.python.keras import losses
from tensorflow.contrib.keras.python.keras import metrics as metrics_module
from tensorflow.contrib.keras.python.keras import optimizers
from tensorflow.contrib.keras.python.keras.engine.topology import Container
from tensorflow.contrib.keras.python.keras.utils.generic_utils import Progbar
from tensorflow.python.platform import tf_logging as logging


# pylint: disable=g-import-not-at-top
try:
  import queue
except ImportError:
  import Queue as queue
# pylint: enable=g-import-not-at-top


def _standardize_input_data(data,
                            names,
                            shapes=None,
                            check_batch_axis=True,
                            exception_prefix=''):
  """Normalizes inputs and targets provided by users.

  Users may pass data as a list of arrays, dictionary of arrays,
  or as a single array. We normalize this to an ordered list of
  arrays (same order as `names`), while checking that the provided
  arrays have shapes that match the network's expectations.

  Arguments:
      data: User-provided input data (polymorphic).
      names: List of expected array names.
      shapes: Optional list of expected array shapes.
      check_batch_axis: B
oolean; whether to check that
          the batch axis of the arrays matches the expected
          value found in `shapes`.
      exception_prefix:
 String prefix used for exception formatting.

  Returns:
      List of standardized input arrays (one array per model input).

  Raises:
      ValueError: in case of improperly formatted user-provided data.
  """
  if not names:
    return []
  if data is None:
    return [None for _ in range(len(names))]
  if isinstance(data, dict):
    arrays = []
    for name in names:
      if name not in data:
        raise ValueError('No data provided for "' + name +
                         '". Need data for each key in: ' + str(names))
      arrays.append(data[name])
  elif isinstance(data, list):
    if len(data) != len(names):
      if data and hasattr(data[0], 'shape'):
        raise ValueError(
            'Error when checking model ' + exception_prefix +
            ': the list of Numpy arrays '
            'that you are passing to your model '
            'is not the size the model expected. '
            'Expected to see ' + str(len(names)) + ' arrays but instead got '
            'the following list of ' + str(len(data)) + ' arrays: ' +
            str(data)[:200] + '...')
      else:
        if len(names) == 1:
          data = [np.asarray(data)]
        else:
          raise ValueError('Error when checking model ' + exception_prefix +
                           ': you are passing a list as '
                           'input to your model, '
                           'but the model expects '
                           'a list of ' + str(len(names)) +
                           ' Numpy arrays instead. '
                           'The list you passed was: ' + str(data)[:200])
    arrays = data
  else:
    if not hasattr(data, 'shape'):
      raise TypeError('Error when checking model ' + exception_prefix +
                      ': data should be a Numpy array, '
                      'or list/dict of Numpy arrays. '
                      'Found: ' + str(data)[:200] + '...')
    if len(names) > 1:
      # Case: model expects multiple inputs but only received
      # a single Numpy array.
      raise ValueError('The model expects ' + str(len(names)) + exception_prefix
                       + ' arrays, but only received one array. '
                       'Found: array with shape ' + str(data.shape))
    arrays = [data]

  # Make arrays at least 2D.
  for i in range(len(names)):
    array = arrays[i]
    if len(array.shape) == 1:
      array = np.expand_dims(array, 1)
      arrays[i] = array

  # Check shapes compatibility.
  if shapes:
    for i in range(len(names)):
      if shapes[i] is None:
        continue
      array = arrays[i]
      if len(array.shape) != len(shapes[i]):
        raise ValueError(
            'Error when checking ' + exception_prefix + ': expected ' + names[i]
            + ' to have ' + str(len(shapes[i])) +
            ' dimensions, but got array with shape ' + str(array.shape))
      for j, (dim, ref_dim) in enumerate(zip(array.shape, shapes[i])):
        if not j and not check_batch_axis:
          # skip the first axis
          continue
        if ref_dim:
          if ref_dim != dim:
            raise ValueError('Error when checking ' + exception_prefix +
                             ': expected ' + names[i] + ' to have shape ' +
                             str(shapes[i]) + ' but got array with shape ' +
                             str(array.shape))
  return arrays


def _standardize_sample_or_class_weights(x_weight, output_names, weight_type):
  """Maps `sample_weight` or `class_weight` to model outputs.

  Arguments:
      x_weight: User-provided `sample_weight` or `class_weight` argument.
      output_names: List of output names (strings) in the model.
      weight_type: A string used purely for exception printing.

  Returns:
      A list of `sample_weight` or `class_weight` where there are exactly
          one element per model output.

  Raises:
      ValueError: In case of invalid user-provided argument.
  """
  if x_weight is None or len(x_weight) == 0:  # pylint: disable=g-explicit-length-test
    return [None for _ in output_names]
  if len(output_names) == 1:
    if isinstance(x_weight, list) and len(x_weight) == 1:
      return x_weight
    if isinstance(x_weight, dict) and output_names[0] in x_weight:
      return [x_weight[output_names[0]]]
    else:
      return [x_weight]
  if isinstance(x_weight, list):
    if len(x_weight) != len(output_names):
      raise ValueError('Provided `' + weight_type + '` was a list of ' +
                       str(len(x_weight)) + ' elements, but the model has ' +
                       str(len(output_names)) + ' outputs. '
                       'You should provide one `' + weight_type + '`'
                       'array per model output.')
    return x_weight
  if isinstance(x_weight, dict):
    x_weights = []
    for name in output_names:
      x_weights.append(x_weight.get(name))
    return x_weights
  else:
    raise TypeError('The model has multiple outputs, so `' + weight_type + '` '
                    'should be either a list of a dict. '
                    'Provided `' + weight_type + '` type not understood: ' +
                    str(x_weight))


def _standardize_class_weights(class_weight, output_names):
  return _standardize_sample_or_class_weights(class_weight, output_names,
                                              'class_weight')


def _standardize_sample_weights(sample_weight, output_names):
  return _standardize_sample_or_class_weights(sample_weight, output_names,
                                              'sample_weight')


def _check_array_lengths(inputs, targets, weights):
  """Does user input validation for numpy arrays.

  Arguments:
      inputs: list of Numpy arrays of inputs