# Copyright 2020 The Trieste Contributors
#
# 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.

r"""
This module contains common optimizers based on :class:`~tf.optimizers.Optimizer` that can be used
with models. Specific models can also sub-class these optimizers or implement their own, and should
register their loss functions using a :func:`create_loss_function`.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from functools import singledispatch
from typing import Any, Callable, Iterable, Optional, Tuple, Union

import scipy
import tensorflow as tf
import tensorflow_probability as tfp

from ..data import Dataset
from ..types import TensorType
from ..utils import jit

TrainingData = Union[Tuple[TensorType, TensorType], Iterable[Tuple[TensorType, TensorType]]]
""" Type alias for a batch, or batches, of training data. """

DatasetTransformer = Callable[[Dataset, Optional[int]], TrainingData]
"""
Type alias for a function that converts a :class:`~trieste.data.Dataset` to batches of training
data.
"""

LossClosure = Callable[[], TensorType]
""" Type alias for a loss closure, typically used in optimization. """

OptimizeResult = Union[scipy.optimize.OptimizeResult, None]
"""
Optimization result. TensorFlow optimizer doesn't return any result. For scipy optimizer that is
also commonly used, it is :class:`~scipy.optimize.OptimizeResult`.
"""


@dataclass
class Optimizer:
    """Optimizer for training models with all the training data at once."""

    optimizer: Any
    """
    The underlying optimizer to use. For example, one of the subclasses of
    :class:`~tensorflow.optimizers.Optimizer` could be used. Note that we use a flexible type `Any`
    to allow for various optimizers that specific models might need to use.
    """

    minimize_args: dict[str, Any] = field(default_factory=lambda: {})
    """ The keyword arguments to pass to the :meth:`minimize` method of the :attr:`optimizer`. """

    compile: bool = False
    """ If `True`, the optimization process will be compiled with :func:`~tf.function`. """

    def create_loss(self, model: tf.Module, dataset: Dataset) -> LossClosure:
        """
        Build a loss function for the specified `model` with the `dataset` using a
        :func:`create_loss_function`.

        :param model: The model to build a loss function for.
        :param dataset: The data with which to build the loss function.
        :return: The loss function.
        """
        x = tf.convert_to_tensor(dataset.query_points)
        y = tf.convert_to_tensor(dataset.observations)
        data = (x, y)
        return create_loss_function(model, data, self.compile)

    def optimize(self, model: tf.Module, dataset: Dataset) -> OptimizeResult:
        """
        Optimize the specified `model` with the `dataset`.

        :param model: The model to optimize.
        :param dataset: The data with which to optimize the `model`.
        :return: The return value of the optimizer's :meth:`minimize` method.
        """
        loss_fn = self.create_loss(model, dataset)
        variables = model.trainable_variables
        return self.optimizer.minimize(loss_fn, variables, **self.minimize_args)


@dataclass
class BatchOptimizer(Optimizer):
    """Optimizer for training models with mini-batches of training data."""

    max_iter: int = 100
    """ The number of iterations over which to optimize the model. """

    batch_size: int = 100
    """ The size of the mini-batches. """

    dataset_builder: DatasetTransformer | None = None
    """ A mapping from :class:`~trieste.observer.Observer` data to mini-batches. """

    def create_loss(self, model: tf.Module, dataset: Dataset) -> LossClosure:
        """
        Build a loss function for the specified `model` with the `dataset`.

        :param model: The model to build a loss function for.
        :param dataset: The data with which to build the loss function.
        :return: The loss function.
        """

        def creator_fn(data: TrainingData) -> LossClosure:
            return create_loss_function(model, data, self.compile)

        if self.dataset_builder is None:
            return creator_fn(
                iter(
                    tf.data.Dataset.from_tensor_slices(dataset.astuple())
                    .shuffle(len(dataset))
                    .batch(self.batch_size)
                    .prefetch(tf.data.experimental.AUTOTUNE)
                    .repeat()
                )
            )

        return creator_fn(self.dataset_builder(dataset, self.batch_size))

    def optimize(self, model: tf.Module, dataset: Dataset) -> None:
        """
        Optimize the specified `model` with the `dataset`.

        :param model: The model to optimize.
        :param dataset: The data with which to optimize the `model`.
        """
        loss_fn = self.create_loss(model, dataset)
        variables = model.trainable_variables

        @jit(apply=self.compile)
        def train_fn() -> None:
            self.optimizer.minimize(loss_fn, variables, **self.minimize_args)

        for _ in range(self.max_iter):
            train_fn()


@dataclass
class KerasOptimizer:
    """Optimizer wrapper for training models implemented with Keras."""

    optimizer: tf.keras.optimizers.Optimizer
    """ The underlying optimizer to use for training the model. """

    fit_args: dict[str, Any] = field(default_factory=lambda: {})
    """
    The keyword arguments to pass to the ``fit`` method of a :class:`~tf.keras.Model` instance.
    See https://keras.io/api/models/model_training_apis/#fit-method for a list of possible
    arguments in the dictionary.
    """

    loss: Optional[
        Union[
            tf.keras.losses.Loss, Callable[[TensorType, tfp.distributions.Distribution], TensorType]
        ]
    ] = None
    """ Optional loss function for training the model. """

    metrics: Optional[list[tf.keras.metrics.Metric]] = None
    """ Optional metrics for monitoring the performance of the network. """


@singledispatch
def create_loss_function(model: Any, dataset: TrainingData, compile: bool = False) -> LossClosure:
    """
    Generic function for building a loss function for a specified `model` and `dataset`.
    The implementations depends on the type of the model, which should use this function as a
    decorator together with its register method to make a model-specific loss function available.

    :param model: The model to build a loss function for.
    :param dataset: The data with which to build the loss function.
    :param compile: Whether to compile with :func:`tf.function`.
    :return: The loss function.
    """
    raise NotImplementedError(f"Unknown model {model} passed for loss function extraction")
