# Copyright (c) 2020, Anthony Nouy, Erwan Grelier
# This file is part of tensap (tensor approximation package).

# tensap is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# tensap is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.

# You should have received a copy of the GNU Lesser General Public License
# along with tensap.  If not, see <https://www.gnu.org/licenses/>.

"""
Module linear_model_learning_custom_loss.

"""

from copy import deepcopy
from numpy import arange, finfo
import numpy as np
import tensap

try:
    import tensorflow as tf

    cond = True
except ImportError:
    cond = False


class LinearModelLearningCustomLoss(tensap.LinearModelLearning):
    """
    Class LinearModelLearningCustomLoss.

    Attributes
    ----------
    optimizer : tensorflow.keras.optimizers.Optimizer
          The optimizer used to solve the learning problem. The default is Adam.
    initial_guess : numpy.ndarray or tensorflow.Tensor
          The initial guess used as a starting point of the optimization
          algorithm. The default is a tensor with components drawn according
          to a standard normal random variable.
    options : dict
        Options for the optimizer:

        - max_iter: the maximum number of iterations of an iterative
          minimization algorithm,
        - stagnation: the value of a stopping criterion based on the
          relative stagnation between two iterates

    """

    def __init__(self, custom_loss):
        """
        Constructor for the class LinearModelLearningCustomLoss.

        Parameters
        ----------
        custom_loss : tap.CustomLossFunction
            The loss function.
        """
        if not cond:
            raise ImportError(
                "Package tensorflow must be installed to "
                + "use LinearModelLearningCustomLoss."
            )
        super().__init__(custom_loss)

        self.optimizer = tf.keras.optimizers.Adam()

        self.initial_guess = None

        self.options = {"max_iterations": 1e3, "stagnation": finfo(float).eps}
        self.var = None

    def __deepcopy__(self, memo):
        # https://github.com/tensorflow/tensorflow/issues/58973
        strategy = self.optimizer._distribution_strategy
        self.optimizer._distribution_strategy = None
        cls = self.__class__
        result = cls.__new__(cls)
        memo[id(self)] = result
        for k, v in self.__dict__.items():
            setattr(result, k, deepcopy(v, memo))
        self.optimizer._distribution_strategy = strategy
        result.optimizer._distribution_strategy = strategy
        return result

    def solve(self):
        """
        Solution of the minimization problem.

        Returns
        -------
        sol : numpy.ndarray or tensap.FunctionalBasisArray
            The solution of the minimization problem.
        output : dict
            Outputs of the algorithm.
        """
        self.initialize()

        if self.initial_guess is None:
            self.initial_guess = tf.random.normal(
                [self.basis_eval.shape[1]], dtype=tf.float64
            )
        else:
            self.initial_guess = tf.convert_to_tensor(
                self.initial_guess, dtype=tf.float64
            )

        basis_eval = self.basis_eval
        training_data = self.training_data

        def risk():
            fun_eval = tf.squeeze(tf.tensordot(basis_eval, self.var, [1, 0]))
            out = self.loss_function.risk_estimation(fun_eval, training_data)
            return out

        if self.var is None:
            self.var = tf.Variable(self.initial_guess, dtype=tf.float64)
            self.optimizer.build([self.var])
        else:
            self.var.assign(self.initial_guess)

        for it in arange(self.options["max_iterations"]):
            var0 = self.var.numpy()
            self.optimizer.minimize(risk, var_list=[self.var])

            stagnation = (
                tf.linalg.norm(var0 - self.var) / tf.linalg.norm(var0)
            ).numpy()
            if stagnation < self.options["stagnation"]:
                break

        sol = self.var.numpy()

        output = {}
        if self.test_error:
            f_eval = np.matmul(self.basis_eval_test, sol)
            test_error = self.loss_function.test_error(f_eval, self.test_data)
            if isinstance(test_error, tf.Tensor):
                test_error = test_error.numpy()
            output["test_error"] = test_error

        if self.basis is not None:
            if np.ndim(sol) == 1:
                sol = tensap.FunctionalBasisArray(sol, self.basis)
            else:
                sol = tensap.FunctionalBasisArray(sol, self.basis, sol.shape[1])

        return sol, output
