"""An estimator that learns to ensemble.

Copyright 2018 The AdaNet 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

    https://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 adanet import core
from adanet.autoensemble.common import _GeneratorFromCandidatePool

import tensorflow.compat.v2 as tf


class AutoEnsembleEstimator(core.Estimator):  # pylint: disable=g-classes-have-attributes
  # pyformat: disable
  """A :class:`tf.estimator.Estimator` that learns to ensemble models.

  Specifically, it learns to ensemble models from a candidate pool using the
  Adanet algorithm.

  .. code-block:: python

      # A simple example of learning to ensemble linear and neural network
      # models.

      import adanet
      import tensorflow as tf

      feature_columns = ...

      head = MultiClassHead(n_classes=10)

      # Learn to ensemble linear and DNN models.
      estimator = adanet.AutoEnsembleEstimator(
          head=head,
          candidate_pool=lambda config: {
              "linear":
                  tf.estimator.LinearEstimator(
                      head=head,
                      feature_columns=feature_columns,
                      config=config,
                      optimizer=...),
              "dnn":
                  tf.estimator.DNNEstimator(
                      head=head,
                      feature_columns=feature_columns,
                      config=config,
                      optimizer=...,
                      hidden_units=[1000, 500, 100])},
          max_iteration_steps=50)

      # Input builders
      def input_fn_train:
        # Returns tf.data.Dataset of (x, y) tuple where y represents label's
        # class index.
        pass
      def input_fn_eval:
        # Returns tf.data.Dataset of (x, y) tuple where y represents label's
        # class index.
        pass
      def input_fn_predict:
        # Returns tf.data.Dataset of (x, None) tuple.
        pass
      estimator.train(input_fn=input_fn_train, steps=100)
      metrics = estimator.evaluate(input_fn=input_fn_eval, steps=10)
      predictions = estimator.predict(input_fn=input_fn_predict)

  Or to train candidate subestimators on different training data subsets:

  .. code-block:: python

      train_data_files = [...]

      # Learn to ensemble linear and DNN models.
      estimator = adanet.AutoEnsembleEstimator(
          head=head,
          candidate_pool=lambda config: {
              "linear":
                  adanet.AutoEnsembleSubestimator(
                      tf.estimator.LinearEstimator(
                          head=head,
                          feature_columns=feature_columns,
                          config=config,
                          optimizer=...),
                      make_train_input_fn(train_data_files[:-1])),
              "dnn":
                  adanet.AutoEnsembleSubestimator(
                      tf.estimator.DNNEstimator(
                          head=head,
                          feature_columns=feature_columns,
                          config=config,
                          optimizer=...,
                          hidden_units=[1000, 500, 100]),
                      make_train_input_fn(train_data_files[0:]))},
          max_iteration_steps=50)

      estimator.train(input_fn=make_train_input_fn(train_data_files), steps=100)


  Args:
    head: A :class:`tf.contrib.estimator.Head` instance for computing loss and
      evaluation metrics for every candidate.
    candidate_pool: List of :class:`tf.estimator.Estimator` and
      :class:`AutoEnsembleSubestimator` objects, or dict of string name to
      :class:`tf.estimator.Estimator` and :class:`AutoEnsembleSubestimator`
      objects that are candidate subestimators to ensemble at each iteration.
      The order does not directly affect which candidates will be included in
      the final ensemble, but will affect the name of the candidate. When using
      a dict, the string key becomes the candidate subestimator's name.
      Alternatively, this argument can be a function that takes a `config`
      argument and returns the aforementioned values in case the
      objects need to be re-instantiated at each adanet iteration.
    max_iteration_steps: Total number of steps for which to train candidates per
      iteration. If `OutOfRange` or `StopIteration` occurs in the middle,
      training stops before `max_iteration_steps` steps.
    logits_fn: A function for fetching the subnetwork logits from a
      :class:`tf.estimator.EstimatorSpec`, which should obey the following
      signature:
        - `Args`: Can only have following argument:
          - estimator_spec: The candidate's :class:`tf.estimator.EstimatorSpec`.
        - `Returns`: Logits :class:`tf.Tensor` or dict of string to logits
          :class:`tf.Tensor` (for multi-head) for the candidate subnetwork
          extracted from the given `estimator_spec`. When `None`, it will
          default to returning `estimator_spec.predictions` when they are a
          :class:`tf.Tensor` or the :class:`tf.Tensor` for the key 'logits' when
          they are a dict of string to :class:`tf.Tensor`.
    last_layer_fn: An optional function for fetching the subnetwork last_layer
      from a :class:`tf.estimator.EstimatorSpec`, which should obey the
      following signature:
        - `Args`: Can only have following argument:
          - estimator_spec: The candidate's :class:`tf.estimator.EstimatorSpec`.
        - `Returns`: Last layer :class:`tf.Tensor` or dict of string to last
          layer :class:`tf.Tensor` (for multi-head) for the candidate subnetwork
          extracted from the given `estimator_spec`. The last_layer can be used
          for learning ensembles or exporting them as embeddings.
      When `None`, it will default to using the logits as the last_layer.
    ensemblers: See :class:`adanet.Estimator`.
    ensemble_strategies: See :class:`adanet.Estimator`.
    evaluator:  See :class:`adanet.Estimator`.
    metric_fn:  See :class:`adanet.Estimator`.
    force_grow:  See :class:`adanet.Estimator`.
    adanet_loss_decay: See :class:`adanet.Estimator`.
    worker_wait_timeout_secs: See :class:`adanet.Estimator`.
    model_dir: See :class:`adanet.Estimator`.
    config: See :class:`adanet.Estimator`.
    debug: See :class:`adanet.Estimator`.
    enable_ensemble_summaries: See :class:`adanet.Estimator`.
    enable_subnetwork_summaries: See :class:`adanet.Estimator`.
    global_step_combiner_fn: See :class:`adanet.Estimator`.
    max_iterations: See :class:`adanet.Estimator`.
    replay_config: See :class:`adanet.Estimator`.
    **kwargs: Extra keyword args passed to the parent.

  Returns:
    An :class:`adanet.AutoEnsembleEstimator` instance.

  Raises:
    ValueError: If any of the candidates in `candidate_pool` are not
      :class:`tf.estimator.Estimator` instances.
  """
  # pyformat: enable

  def __init__(self,
               head,
               candidate_pool,
               max_iteration_steps,
               ensemblers=None,
               ensemble_strategies=None,
               logits_fn=None,
               last_layer_fn=None,
               evaluator=None,
               metric_fn=None,
               force_grow=False,
               adanet_loss_decay=.9,
               worker_wait_timeout_secs=7200,
               model_dir=None,
               config=None,
               debug=False,
               enable_ensemble_summaries=True,
               enable_subnetwork_summaries=True,
               global_step_combiner_fn=tf.math.reduce_mean,
               max_iterations=None,
               replay_config=None,
               **kwargs):
    subnetwork_generator = _GeneratorFromCandidatePool(candidate_pool,
                                                       logits_fn, last_layer_fn)
    super(AutoEnsembleEstimator, self).__init__(
        head=head,
        subnetwork_generator=subnetwork_generator,
        max_iteration_steps=max_iteration_steps,
        ensemblers=ensemblers,
        ensemble_strategies=ensemble_strategies,
        evaluator=evaluator,
        metric_fn=metric_fn,
        force_grow=force_grow,
        adanet_loss_decay=adanet_loss_decay,
        worker_wait_timeout_secs=worker_wait_timeout_secs,
        model_dir=model_dir,
        config=config,
        debug=debug,
        enable_ensemble_summaries=enable_ensemble_summaries,
        enable_subnetwork_summaries=enable_subnetwork_summaries,
        global_step_combiner_fn=global_step_combiner_fn,
        max_iterations=max_iterations,
        replay_config=replay_config,
        **kwargs)


class AutoEnsembleTPUEstimator(core.TPUEstimator):  # pylint: disable=g-classes-have-attributes
  # pyformat: disable
  """A :class:`tf.estimator.tpu.TPUEstimator` that learns to ensemble models.

  Specifically, it learns to ensemble models from a candidate pool using the
  Adanet algorithm.

  This estimator is capable of training and evaluating on TPU. It can ensemble
  both :class:`tf.estimator.tpu.TPUEstimator` candidates as well as regular
  :class:`tf.estimator.Estimator` candidates, as long as these candidates are
  TPU compatible.

  Note the following restrictions compared to AutoEnsembleEstimator:
    * All candidates must wrap their optimizers with a
      :class:`tf.tpu.CrossShardOptimizer`.
    * The `input_fn` must expose a `params` argument.
    * The `model_fn` of :class:`tf.estimator.tpu.TPUEstimator` candidates must
      also expose a `params` argument.

  WARNING: This Estimator is a work in progress and the API could change at any
  moment. May not support all AutoEnsembleEstimator features.

    .. code-block:: python

      # A simple example of learning to ensemble linear and neural network
      # models on TPU.

      import adanet
      import tensorflow as tf

      feature_columns = ...

      head = MultiClassHead(n_classes=10)

      # Learn to ensemble linear and DNN models.
      estimator = adanet.AutoEnsembleTPUEstimator(
          head=head,
          candidate_pool=lambda config: {
              "linear":
                  tf.estimator.LinearEstimator(
                      head=head,
                      feature_columns=feature_columns,
                      config=config,
                      optimizer=tf.tpu.CrossShardOptimizer(...)),
              "dnn":
                  tf.estimator.DNNEstimator(
                      head=head,
                      feature_columns=feature_columns,
                      config=config,
                      optimizer=tf.tpu.CrossShardOptimzier(...),
                      hidden_units=[1000, 500, 100])},
          max_iteration_steps=50)

      # Input builders
      def input_fn_train(params):
        # Returns tf.data.Dataset of (x, y) tuple where y represents label's
        # class index.
        pass
      def input_fn_eval(params):
        # Returns tf.data.Dataset of (x, y) tuple where y represents label's
        # class index.
        pass
      def input_fn_predict():
        # Returns tf.data.Dataset of (x, None) tuple.
        pass
      estimator.train(input_fn=input_fn_train, steps=100)
      metrics = estimator.evaluate(input_fn=input_fn_eval, steps=10)
      predictions = estimator.predict(input_fn=input_fn_predict)

  Args:
    head: A :class:`tf.contrib.estimator.Head` instance for computing loss and
      evaluation metrics for every candidate.
    candidate_pool: List of :class:`tf.estimator.tpu.TPUEstimator` and
      :class:`AutoEnsembleSubestimator` objects, or dict of string name to
      :class:`tf.estimator.tpu.TPUEstimator` and
      :class:`AutoEnsembleSubestimator` objects that are candidate subestimators
      to ensemble at each iteration. The order does not directly affect which
      candidates will be included in the final ensemble, but will affect the
      name of the candidate. When using a dict, the string key becomes the
      candidate subestimator's name. Alternatively, this argument can be a
      function that takes a `config` argument and returns the aforementioned
      values in case the objects need to be re-instantiated at each adanet
      iteration.
    max_iteration_steps: See :class:`adanet.Estimator`.
    logits_fn: A function for fetching the subnetwork logits from a
      :class:`tf.estimator.EstimatorSpec`, which should obey the following
      signature:
        - `Args`: Can only have following argument:
          - estimator_spec: The candidate's :class:`tf.estimator.EstimatorSpec`.
        - `Returns`: Logits :class:`tf.Tensor` or dict of string to logits
          :class:`tf.Tensor` (for multi-head) for the candidate subnetwork
          extracted from the given `estimator_spec`. When `None`, it will
          default to returning `estimator_spec.predictions` when they are a
          :class:`tf.Tensor` or the :class:`tf.Tensor` for the key 'logits' when
          they are a dict of string to :class:`tf.Tensor`.
    last_layer_fn: An optional function for fetching the subnetwork last_layer
      from a :class:`tf.estimator.EstimatorSpec`, which should obey the
      following signature:
        - `Args`: Can only have following argument:
          - estimator_spec: The candidate's :class:`tf.estimator.EstimatorSpec`.
        - `Returns`: Last layer :class:`tf.Tensor` or dict of string to last
          layer :class:`tf.Tensor` (for multi-head) for the candidate subnetwork
          extracted from the given `estimator_spec`. The last_layer can be used
          for learning ensembles or exporting them as embeddings.
      When `None`, it will default to using the logits as the last_layer.
    ensemblers: See :class:`adanet.Estimator`.
    ensemble_strategies: See :class:`adanet.Estimator`.
    evaluator:  See :class:`adanet.Estimator`.
    metric_fn:  See :class:`adanet.Estimator`.
    force_grow:  See :class:`adanet.Estimator`.
    adanet_loss_decay: See :class:`adanet.Estimator`.
    model_dir: See :class:`adanet.Estimator`.
    config: See :class:`adanet.Estimator`.
    use_tpu: See :class:`adanet.Estimator`.
    eval_on_tpu: See :class:`adanet.Estimator`.
    export_to_tpu: See :class:`adanet.Estimator`.
    train_batch_size: See :class:`adanet.Estimator`.
    eval_batch_size: See :class:`adanet.Estimator`.
    embedding_config_spec: See :class:`adanet.Estimator`.
    debug: See :class:`adanet.Estimator`.
    enable_ensemble_summaries: See :class:`adanet.Estimator`.
    enable_subnetwork_summaries: See :class:`adanet.Estimator`.
    global_step_combiner_fn: See :class:`adanet.Estimator`.
    max_iterations: See :class:`adanet.Estimator`.
    replay_config: See :class:`adanet.Estimator`.
    **kwargs: Extra keyword args passed to the parent.

  Returns:
    An :class:`adanet.AutoEnsembleTPUEstimator` instance.

  Raises:
    ValueError: If any of the candidates in `candidate_pool` are not
      :class:`tf.estimator.Estimator` instances.
  """
  # pyformat: enable

  def __init__(self,
               head,
               candidate_pool,
               max_iteration_steps,
               ensemblers=None,
               ensemble_strategies=None,
               logits_fn=None,
               last_layer_fn=None,
               evaluator=None,
               metric_fn=None,
               force_grow=False,
               adanet_loss_decay=.9,
               model_dir=None,
               config=None,
               use_tpu=True,
               eval_on_tpu=True,
               export_to_tpu=True,
               train_batch_size=None,
               eval_batch_size=None,
               predict_batch_size=None,
               embedding_config_spec=None,
               debug=False,
               enable_ensemble_summaries=True,
               enable_subnetwork_summaries=True,
               global_step_combiner_fn=tf.math.reduce_mean,
               max_iterations=None,
               replay_config=None,
               **kwargs):
    subnetwork_generator = _GeneratorFromCandidatePool(candidate_pool,
                                                       logits_fn, last_layer_fn)
    super(AutoEnsembleTPUEstimator, self).__init__(
        head=head,
        subnetwork_generator=subnetwork_generator,
        max_iteration_steps=max_iteration_steps,
        ensemblers=ensemblers,
        ensemble_strategies=ensemble_strategies,
        evaluator=evaluator,
        metric_fn=metric_fn,
        force_grow=force_grow,
        adanet_loss_decay=adanet_loss_decay,
        model_dir=model_dir,
        config=config,
        use_tpu=use_tpu,
        eval_on_tpu=eval_on_tpu,
        export_to_tpu=export_to_tpu,
        train_batch_size=train_batch_size,
        eval_batch_size=eval_batch_size,
        predict_batch_size=predict_batch_size,
        embedding_config_spec=embedding_config_spec,
        debug=debug,
        enable_ensemble_summaries=enable_ensemble_summaries,
        enable_subnetwork_summaries=enable_subnetwork_summaries,
        global_step_combiner_fn=global_step_combiner_fn,
        max_iterations=max_iterations,
        replay_config=replay_config,
        **kwargs)
