"""Generic trainer for TensorFlow models."""
#  Copyright 2015-present The Scikit Flow 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.

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

from six.moves import xrange   # pylint: disable=redefined-builtin

from tensorflow.python.framework import ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import clip_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gradients
from tensorflow.python.ops import variables
from tensorflow.python.ops import variable_scope as vs
from tensorflow.contrib.layers import optimizers


def train(session, train_op, loss, global_step, feed_dict_fn, steps, monitor,
          summary_writer=None, summaries=None,
          feed_params_fn=None):
    """Trains a model for given number of steps, given feed_dict function.

    Args:
        session: Session object.
        train: Tensor, trains model.
        loss: Tensor, loss value.
        global_step: Tensor, global step of the model.
        feed_dict_fn: Function that will return a feed dictionary.
        summary_writer: SummaryWriter object to use for writing summaries.
        steps: Number of steps to run.
        monitor: Monitor object to track training progress and induce early stopping
        summaries: Joined object of all summaries that should be ran.
    """
    for step in xrange(steps):
        feed_dict = feed_dict_fn()
        if summaries is not None:
            global_step_value, loss_value, summ, _ = session.run(
                [global_step, loss, summaries, train_op],
                feed_dict=feed_dict)
        else:
            global_step_value, loss_value, _ = session.run(
                [global_step, loss, train_op],
                feed_dict=feed_dict)
        monitor.update(step, global_step_value, loss_value, session,
                       feed_params_fn, loss_expression_tensor=loss)
        if summaries is not None and summary_writer and summ is not None:
            summary_writer.add_summary(summ, global_step_value)
        if monitor.monitor_inducing_stop():
            break

