# Lint as: python3
# Copyright 2020 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.
"""A phase that automatically ensembles models."""

import abc
import random

from typing import Iterable, Iterator, List
from adanet.experimental.keras.ensemble_model import EnsembleModel
from adanet.experimental.keras.ensemble_model import MeanEnsemble
from adanet.experimental.phases.phase import DatasetProvider
from adanet.experimental.phases.phase import ModelProvider
from adanet.experimental.storages.in_memory_storage import InMemoryStorage
from adanet.experimental.storages.storage import Storage
from adanet.experimental.work_units.keras_trainer_work_unit import KerasTrainerWorkUnit
from adanet.experimental.work_units.work_unit import WorkUnit
import tensorflow.compat.v2 as tf


class EnsembleStrategy(abc.ABC):
  """An abstract ensemble strategy."""

  @abc.abstractmethod
  def __call__(
      self, candidates: List[tf.keras.Model]) -> Iterable[List[tf.keras.Model]]:
    pass


class Ensembler(abc.ABC):
  """An abstract ensembler."""

  def __init__(self, loss, optimizer, metrics):
    self._loss = loss
    self._optimizer = optimizer
    self._metrics = metrics

  @abc.abstractmethod
  def __call__(self, submodels: List[tf.keras.Model]) -> EnsembleModel:
    pass


class MeanEnsembler(Ensembler):
  """An ensembler that averages the weights of submodel outputs."""

  def __init__(self, loss, optimizer, metrics, freeze_submodels=True):
    super().__init__(loss, optimizer, metrics)
    self._freeze_submodels = freeze_submodels

  def __call__(self, submodels: List[tf.keras.Model]) -> EnsembleModel:
    ensemble = MeanEnsemble(submodels, freeze_submodels=self._freeze_submodels)
    if self._freeze_submodels:
      for layer in ensemble.layers:
        layer.trainable = False
    # Compile SGD with learning rate set to 0 for no weight updates.
    ensemble.compile(
        loss=self._loss, optimizer=tf.keras.optimizers.SGD(0),
        metrics=self._metrics)
    return ensemble


class GrowStrategy(EnsembleStrategy):
  """An ensemble strategy that adds one candidate to the ensemble at a time."""

  def __call__(
      self, candidates: List[tf.keras.Model]) -> Iterable[List[tf.keras.Model]]:
    return [[candidate] for candidate in candidates]


class AllStrategy(EnsembleStrategy):
  """An ensemble strategy that adds all candidates to the ensemble."""

  def __call__(
      self, candidates: List[tf.keras.Model]) -> Iterable[List[tf.keras.Model]]:
    return [candidates]


class RandomKStrategy(EnsembleStrategy):
  """An ensemble strategy that adds k random candidates (with replacement)."""

  def __init__(self, k, seed=None):
    """Initializes a RandomKStrategy ensemble strategy.

    Args:
      k: Number of candidates to sample.
      seed: Random seed.
    """
    self._k = k
    self._seed = seed

  def __call__(
      self, candidates: List[tf.keras.Model]) -> Iterable[List[tf.keras.Model]]:
    if self._seed:
      random_state = random.getstate()
      random.seed(self._seed)
      candidates = [random.choices(candidates, k=self._k)]
      random_state = random.setstate(random_state)
    else:
      candidates = [random.choices(candidates, k=self._k)]
    return [candidates]


class AutoEnsemblePhase(DatasetProvider, ModelProvider):
  """A phase that automatically ensembles models from a prior phase."""

  def __init__(self,
               ensemblers: List[Ensembler],
               ensemble_strategies: List[EnsembleStrategy],
               storage: Storage = InMemoryStorage(),
               num_candidates: int = None):
    """Initializes an AutoEnsemblePhase.

    Args:
      ensemblers: A list of `Ensembler` instances to determine how to combine
        subnetworks.
      ensemble_strategies: A list of `EnsembleStrategy` instances to determine
        which subnetworks compose an ensemble.
      storage: A `Storage` instance to store models and model metadata.
      num_candidates: The number of subnetwork candidates to consider from the
        previous phase. If `None` then all of the subnetworks generated in the
        previous phase will be considered.
    """

    super().__init__(storage)
    self._ensemblers = ensemblers
    self._ensemble_strategies = ensemble_strategies
    self._num_candidates = num_candidates

  def work_units(self, previous_phase) -> Iterator[WorkUnit]:
    self._train_dataset = previous_phase.get_train_dataset()
    self._eval_dataset = previous_phase.get_eval_dataset()
    if self._num_candidates:
      candidates = previous_phase.get_best_models(
          num_models=self._num_candidates)
    else:
      candidates = previous_phase.get_models()
    if self.get_best_models():
      current_best_ensemble = list(self.get_best_models())[0]
    else:
      current_best_ensemble = None

    for ensemble_strategy in self._ensemble_strategies:
      for submodels in ensemble_strategy(candidates):
        for ensembler in self._ensemblers:
          if current_best_ensemble:
            previous_ensemble = current_best_ensemble.submodels
          else:
            previous_ensemble = []
          ensemble = ensembler(previous_ensemble + submodels)
          yield KerasTrainerWorkUnit(ensemble,
                                     previous_phase.get_train_dataset(),
                                     previous_phase.get_eval_dataset(),
                                     self._storage)

  def get_models(self) -> Iterable[tf.keras.Model]:
    return self._storage.get_models()

  def get_best_models(self, num_models=1) -> Iterable[tf.keras.Model]:
    return self._storage.get_best_models(num_models)

  # TODO: Add some way to check that work_units has to be called
  # before accessing these methods.
  def get_train_dataset(self) -> tf.data.Dataset:
    return self._train_dataset

  def get_eval_dataset(self) -> tf.data.Dataset:
    return self._eval_dataset
