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

"""Various high level TF models."""

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

from tensorflow.contrib.learn.python.learn.ops import autoencoder_ops
from tensorflow.contrib.learn.python.learn.ops import dnn_ops
from tensorflow.contrib.learn.python.learn.ops import losses_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops as array_ops_
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import logging_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import variable_scope as vs


def linear_regression_zero_init(X, y):
  # pylint: disable=invalid-name
  """Linear regression subgraph with zero-value initial weights and bias.

  Args:
    X: tensor or placeholder for input features.
    y: tensor or placeholder for target.

  Returns:
    Predictions and loss tensors.
  """
  return linear_regression(X, y, init_mean=0.0, init_stddev=0.0)


def logistic_regression_zero_init(X, y):
  # pylint: disable=invalid-name
  """Logistic regression subgraph with zero-value initial weights and bias.

  Args:
    X: tensor or placeholder for input features.
    y: tensor or placeholder for target.

  Returns:
    Predictions and loss tensors.
  """
  return logistic_regression(X, y, init_mean=0.0, init_stddev=0.0)


def linear_regression(X, y, init_mean=None, init_stddev=1.0):
  # pylint: disable=invalid-name
  """Creates linear regression TensorFlow subgraph.

  Args:
    X: tensor or placeholder for input features.
    y: tensor or placeholder for target.
    init_mean: the mean value to use for initialization.
    init_stddev: the standard devation to use for initialization.

  Returns:
    Predictions and loss tensors.

  Side effects:
    The variables linear_regression.weights and linear_regression.bias are
    initialized as follows.  If init_mean is not None, then initialization
    will be done using a random normal initializer with the given init_mean
    and init_stddv.  (These may be set to 0.0 each if a zero initialization
    is desirable for convex use cases.)  If init_mean is None, then the
    uniform_unit_scaling_initialzer will be used.
  """
  with vs.variable_scope('linear_regression'):
    logging_ops.histogram_summary('linear_regression.X', X)
    logging_ops.histogram_summary('linear_regression.y', y)
    y_shape = y.get_shape()
    if len(y_shape) == 1:
      output_shape = 1
    else:
      output_shape = y_shape[1]
    # Set up the requested initialization.
    if init_mean is None:
      weights = vs.get_variable('weights', [X.get_shape()[1], output_shape])
      bias = vs.get_variable('bias', [output_shape])
    else:
      weights = vs.get_variable('weights', [X.get_shape()[1], output_shape],
                                initializer=init_ops.random_normal_initializer(
                                    init_mean, init_stddev))
      bias = vs.get_variable('bias', [output_shape],
                             initializer=init_ops.random_normal_initializer(
                                 init_mean, init_stddev))
    logging_ops.histogram_summary('linear_regression.weights', weights)
    logging_ops.histogram_summary('linear_regression.bias', bias)
    return losses_ops.mean_squared_error_regressor(X, y, weights, bias)


def logistic_regression(X,
                        y,
                        class_weight=None,
                        init_mean=None,
                        init_stddev=1.0):
  # pylint: disable=invalid-name
  """Creates logistic regression TensorFlow subgraph.

  Args:
    X: tensor or placeholder for input features,
       shape should be [batch_size, n_features].
    y: tensor or placeholder for target,
       shape should be [batch_size, n_classes].
    class_weight: tensor, [n_classes], where for each class
                  it has weight of the class. If not provided
                  will check if graph contains tensor `class_weight:0`.
                  If that is not provided either all ones are used.
    init_mean: the mean value to use for initialization.
    init_stddev: the standard devation to use for initialization.

  Returns:
    Predictions and loss tensors.

  Side effects:
    The variables linear_regression.weights and linear_regression.bias are
    initialized as follows.  If init_mean is not None, then initialization
    will be done using a random normal initializer with the given init_mean
    and init_stddv.  (These may be set to 0.0 each if a zero initialization
    is desirable for convex use cases.)  If init_mean is None, then the
    uniform_unit_scaling_initialzer will be used.
  """
  with vs.variable_scope('logistic_regression'):
    logging_ops.histogram_summary('%s.X' % vs.get_variable_scope().name, X)
    logging_ops.histogram_summary('%s.y' % vs.get_variable_scope().name, y)
    # Set up the requested initialization.
    if init_mean is None:
      weights = vs.get_variable('weights',
                                [X.get_shape()[1], y.get_shape()[-1]])
      bias = vs.get_variable('bias', [y.get_shape()[-1]])
    else:
      weights = vs.get_variable('weights',
                                [X.get_shape()[1], y.get_shape()[-1]],
                                initializer=init_ops.random_normal_initializer(
                                    init_mean, init_stddev))
      bias = vs.get_variable('bias', [y.get_shape()[-1]],
                             initializer=init_ops.random_normal_initializer(
                                 init_mean, init_stddev))
    logging_ops.histogram_summary('%s.weights' % vs.get_variable_scope().name,
                                  weights)
    logging_ops.histogram_summary('%s.bias' % vs.get_variable_scope().name,
                                  bias)
    # If no class weight provided, try to retrieve one from pre-defined
    # tensor name in the graph.
    if not class_weight:
      try:
        class_weight = ops.get_default_graph().get_tensor_by_name(
            'class_weight:0')
      except KeyError:
        pass

    return losses_ops.softmax_classifier(X,
                                         y,
                                         weights,
                                         bias,
                                         class_weight=class_weight)


def get_dnn_model(hidden_units, target_predictor_fn, dropout=None):
  """Returns a function that creates a DNN TensorFlow subgraph.

  Args:
    hidden_units: List of values of hidden units for layers.
    target_predictor_fn: Function that will predict target from input
                         features. This can be logistic regression,
                         linear regression or any other model,
                         that takes X, y and returns predictions and loss
                         tensors.
    dropout: When not none, causes dropout regularization to be used,
             with the specified probability of removing a given coordinate.

  Returns:
    A function that creates the subgraph.
  """

  def dnn_estimator(X, y):
    # pylint: disable=invalid-name
    """DNN estimator with target predictor function on top."""
    layers = dnn_ops.dnn(X, hidden_units, dropout=dropout)
    return target_predictor_fn(layers, y)

  return dnn_estimator


def get_autoencoder_model(hidden_units, target_predictor_fn,