# 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 canonical_tensor_learning.

"""

from copy import deepcopy
import numpy as np
import tensap


class CanonicalTensorLearning(tensap.TensorLearning):
    """
    Class CanonicalTensorLearning.

    See also tensap.TensorLearning.

    Attributes
    ----------
    order : int
        The order of the tensor.

    """

    def __init__(self, order, *args):
        """
        Constructor for the class CanonicalTensorLearning.

        Parameters
        ----------
        order : int
            The order of the tensor.
        *args : tuple
            Additional parameters.

        Returns
        -------
        None.

        """
        super().__init__(*args)
        self.order = order
        self.initialization_type = "mean"
        self.truncate_initialization = True

    # %% Standard solver methods
    def initialize(self):
        if self.tree_adaptation:
            print("tree_adaptation not defined for CanonicalTensorLearning.")
            self.tree_adaptation = False

        if (
            "one_by_one_factor" in self.alternating_minimization_parameters
            and self.alternating_minimization_parameters["one_by_one_factor"]
            and self.rank != 1
        ):
            self._exploration_strategy = np.arange(
                1,
                self.alternating_minimization_parameters["inner_loops"]
                * self.rank
                * self.order
                + 2,
            )
            self._number_of_parameters = self._exploration_strategy.size
        else:
            self.alternating_minimization_parameters["one_by_one_factor"] = False
            self._exploration_strategy = np.arange(1, self.order + 2)
            self._number_of_parameters = self.order + 1

        shape = [x.shape[1] for x in self.bases_eval]
        if self.initialization_type == "random":
            f = tensap.CanonicalTensor.randn(self.rank, shape)
        elif self.initialization_type == "ones":
            f = tensap.CanonicalTensor.ones(self.rank, shape)
        elif self.initialization_type == "initial_guess":
            f = self.initial_guess
        elif (
            self.initialization_type == "mean"
            or self.initialization_type == "mean_randomized"
        ):
            if not isinstance(self.training_data, list) or (
                isinstance(self.training_data, list) and len(self.training_data) == 1
            ):
                raise NotImplementedError(
                    "Initialization type not "
                    + "implemented in unsupervised "
                    + "learning."
                )
            if isinstance(self.bases, tensap.FunctionalBases):
                means = self.bases.mean()
            else:
                means = [np.mean(x, 0) for x in self.bases_eval]
            if self.initialization_type == "mean_randomized":
                means = [x + 0.01 * np.random.randn(*x.shape) for x in means]
            means = [np.reshape(x, [-1, 1]) for x in means]

            f = tensap.CanonicalTensor(
                means, np.atleast_1d(np.mean(self.training_data[1]))
            )
        elif self.initialization_type == "greedy":
            s_ini = deepcopy(self)
            s_ini.rank_adaptation = False
            s_ini.algorithm = "greedy"
            s_ini.initialization_type = "mean"
            s_ini.alternating_minimization_parameters["display"] = False
            s_ini.linear_model_learning.error_estimation = False
            s_ini.test_error = False
            s_ini.display = False
            f, output_ini = s_ini.solve()

            if self.display and "error" in output_ini:
                print(
                    "Greedy initialization: rank = %i, error = %2.5e"
                    % (len(f.tensor.core.data), output_ini["error"])
                )
        else:
            raise ValueError("Wrong initialization type.")

        if isinstance(f, tensap.FunctionalTensor):
            f = f.tensor

        if self.rank > len(f.core.data):
            fx = f.tensor_matrix_product(self.bases_eval).eval_diag().data

            s_ini = deepcopy(self)
            s_ini.rank_adaptation = False
            s_ini.algorithm = "standard"
            s_ini.initialization_type = "greedy"
            s_ini.rank = self.rank - len(f.core.data)
            s_ini.alternating_minimization_parameters["display"] = False
            s_ini.linear_model_learning.error_estimation = False
            s_ini.test_error = False
            s_ini.display = False
            if isinstance(s_ini.training_data, list) and len(s_ini.training_data) == 2:
                s_ini.training_data[1] -= fx
            elif isinstance(s_ini.loss_function, tensap.DensityL2LossFunction):
                s_ini.training_data = [s_ini.training_data, fx]

            f_add = s_ini.solve()[0]
            if isinstance(f_add, tensap.FunctionalTensor):
                f += f_add.tensor
            else:
                f += f_add

        if self.truncate_initialization:
            if self.order == 2:
                tr = tensap.Truncator()
                f = tr.truncate(f)
                f = tensap.CanonicalTensor(f.space, f.core.data)
                self.rank = len(f.core.data)
        return self, f

    def pre_processing(self, f):
        return self, f

    def randomize_exploration_strategy(self):
        if (
            "one_by_one_factor" not in self.alternating_minimization_parameters
            or not self.alternating_minimization_parameters["one_by_one_factor"]
        ):
            strategy = np.concatenate(
                (
                    np.random.permutation(self._number_of_parameters - 1) + 1,
                    [self._number_of_parameters],
                )
            )
        else:
            strat_mu = np.random.permutation(self.order)
            strat_i = np.random.permutation(self.rank)
            strategy = np.reshape(
                self._exploration_strategy[:-1], [-1, self.order], order="F"
            )
            strategy[:, strat_mu] = np.array(strategy)
            strategy = np.reshape(
                np.reshape(strategy, [1, -1], order="F"), [self.rank, -1], order="F"
            )
            strategy[strat_i, :] = np.array(strategy)
            strategy = np.concatenate(
                (np.reshape(strategy, -1, order="F"), [self._number_of_parameters])
            )
        return strategy

    def prepare_alternating_minimization_system(self, f, mu):
        assert isinstance(
            self.loss_function, tensap.SquareLossFunction
        ), "Method not implemented for this loss function."
        y = self.training_data[1]
        N = self.bases_eval[0].shape[0]

        if mu != self._number_of_parameters:
            if self.alternating_minimization_parameters["one_by_one_factor"]:
                ind = mu
                mu = int(
                    np.ceil(
                        ind
                        / self.alternating_minimization_parameters["inner_loops"]
                        / self.rank
                    )
                )
                ind = int(ind - self.rank * (np.ceil(ind / self.rank) - 1))

                coef = f.tensor.space[mu - 1][:, ind - 1]

                fH = f.tensor.tensor_matrix_product(self.bases_eval)

                fH_mu = np.ones((N, self.rank))
                no_mu = np.setdiff1d(np.arange(1, f.tensor.order + 1), mu)
                for nu in no_mu:
                    fH_mu *= fH.space[nu - 1]

                B = fH_mu * fH.space[mu - 1]
                no_i = np.setdiff1d(np.arange(1, self.rank + 1), ind)
                b = y - np.matmul(B[:, no_i - 1], fH.core.data[no_i - 1])

                A = self.bases_eval[mu - 1] * np.transpose(
                    np.tile(fH_mu[:, ind - 1], (f.tensor.shape[mu - 1], 1))
                )
                if self.linear_model_learning[mu - 1].basis_adaptation:
                    self.linear_model_learning[mu - 1].basis_adaptation_path = (
                        self.bases_adaptation_path[mu - 1]
                    )
            else:
                coef = f.tensor.space[mu - 1]
                grad = f.parameter_gradient_eval(mu).transpose([0, 2, 1])
                A = np.reshape(grad.data, [grad.shape[0], -1], order="F")

                if self.linear_model_learning[mu - 1].basis_adaptation:
                    self.linear_model_learning[mu - 1].basis_adaptation_path = np.tile(
                        self.bases_adaptation_path[mu - 1], (self.rank, 1)
                    )
                elif self.rank > 1:
                    self.linear_model_learning[mu - 1].options["non_zero_blocks"] = (
                        np.empty(self.rank, dtype=object)
                    )
                    for kk in range(self.rank):
                        shape_mu = f.tensor.shape[mu - 1]
                        self.linear_model_learning[mu - 1].options["non_zero_blocks"][
                            kk
                        ] = kk * shape_mu + np.arange(shape_mu)
                b = y
        else:
            coef = f.tensor.core.data
            if self.alternating_minimization_parameters["one_by_one_factor"]:
                mu = self.order + 1
            grad = f.parameter_gradient_eval(mu)
            A = np.reshape(grad.data, [grad.shape[0], -1], order="F")
            b = y
            self.linear_model_learning[mu - 1].basis_adaptation = False

        self.linear_model_learning[mu - 1].initial_guess = np.reshape(coef, -1)

        return self, A, b, f

    def set_parameter(self, f, mu, coef):
        if mu != self._number_of_parameters:
            if not self.alternating_minimization_parameters["one_by_one_factor"]:
                coef = np.reshape(coef, [f.tensor.shape[mu - 1], self.rank], order="F")
                norm_coef = np.sqrt(np.sum(coef**2, 0))
                ind = norm_coef != 0
                if not np.all(ind):
                    print("Degenerate case: one factor is zero.")
                coef[:, ind] = np.matmul(coef[:, ind], np.diag(1 / norm_coef[ind]))
                f.tensor.space[mu - 1] = coef
                f.tensor.core.data = np.ravel(norm_coef)

                if len(f.tensor.space) == 2 and mu == 1:
                    f.tensor.space[0] = np.linalg.qr(f.tensor.space[0])[0]
            else:
                ind = mu
                mu = int(
                    np.ceil(
                        ind
                        / self.alternating_minimization_parameters["inner_loops"]
                        / self.rank
                    )
                )
                ind = int(ind - self.rank * (np.ceil(ind / self.rank) - 1))
                norm_coef = np.linalg.norm(coef)
                coef /= norm_coef
                f.tensor.space[mu - 1][:, ind - 1] = coef
                f.tensor.core.data[ind - 1] = norm_coef
        else:
            f.tensor.core.data = coef
        return f

    def stagnation_criterion(self, f, f0):
        norm_f = f.tensor.norm()
        norm_f0 = f0.tensor.norm()
        return 2 * np.abs(norm_f - norm_f0) / (norm_f + norm_f0)

    def final_display(self, f):
        print("Rank = %i" % len(f.tensor.core.data), end="")

    # %% Rank adaptation solver methods
    def local_solver(self):
        s_local = deepcopy(self)
        s_local.display = False
        s_local.rank_adaptation = False
        s_local.test_error = False
        s_local.algorithm = "standard"
        s_local.initialization_type = "mean"
        s_local.model_selection = False
        return s_local

    def new_rank_selection(self, f):
        return len(f.tensor.core.data) + 1, 1, deepcopy(f)

    def initial_guess_new_rank(self, s_local, f, *args):
        s_local.initialization_type = "initial_guess"
        s_local.initial_guess = f
        return s_local

    def adaptation_display(self, f, *args):
        print("\tRank = %i" % len(f.tensor.core.data), flush=True)

    # %% Greedy solver
    def _solve_greedy(self):
        """
        Greedy solver in canonical tensor format.

        Raises
        ------
        ValueError
            If the number of LinearModelLearning objects is not equal to
            _numberOfParameters.

        Returns
        -------
        f : tensap.FunctionalTensor
            The learned approximation.
        output : dict
            The outputs of the solver.

        """
        assert isinstance(
            self.loss_function, tensap.SquareLossFunction
        ), "Method not implemented for this loss function."

        bases_eval = self.bases_eval
        output = {}
        output["sequence"] = np.empty(self.rank, dtype=object)

        # Replication of the LinearModelLearning objects
        if self.linear_model_learning_parameters[
            "identical_for_all_parameters"
        ] and not isinstance(self.linear_model_learning, (list, np.ndarray)):
            self.linear_model_learning = list(
                map(deepcopy, [self.linear_model_learning] * self._number_of_parameters)
            )
        elif (
            isinstance(self.linear_model_learning, (list, np.ndarray))
            and len(self.linear_model_learning) != self._number_of_parameters
        ):
            raise ValueError(
                "Must provide self._numberOfParameters "
                + "LinearModelLearning objects."
            )

        # Working set paths
        if (
            isinstance(self.linear_model_learning, (list, np.ndarray))
            and np.any([x.basis_adaptation for x in self.linear_model_learning])
            and self.bases_adaptation_path is None
        ):
            self.bases_adaptation_path = self.bases.adaptation_path()

        y = self.training_data[1]

        s_local = deepcopy(self)
        s_local.algorithm = "standard"
        s_local.rank = 1
        s_local.display = False
        s_local.test_error = False
        s_local.model_selection = False

        f = tensap.CanonicalTensor.zeros(0, [x.shape[1] for x in bases_eval])
        f_0 = deepcopy(f)
        stagnation = np.zeros((1, self.rank))

        ls_local = deepcopy(
            s_local.linear_model_learning[self._number_of_parameters - 1]
        )

        is_error = False
        error = np.ones((1, self.rank))
        pure_greedy = False
        if not pure_greedy:
            for linear_solver in self.linear_model_learning:
                setattr(linear_solver, "error_estimation", False)

        for k in np.arange(1, self.rank + 1):
            s_local.training_data[1] = (
                y - f.tensor_matrix_product(bases_eval).eval_diag().data
            )
            f_add, output_greedy = s_local.solve()
            if isinstance(f_add, tensap.FunctionalTensor):
                f += f_add.tensor
            else:
                f += f_add

            if not pure_greedy:
                f_H = f.tensor_matrix_product(bases_eval)
                A = np.ones((bases_eval[0].shape[0], len(f_H.core.data)))
                for space in f_H.space:
                    A *= space
                ls_local.basis_adaptation = False
                ls_local.basis = None
                ls_local.basis_eval = A
                ls_local.training_data = [None, y]
                coef, output_greedy = ls_local.solve()
                f.core.data = np.ravel(coef)

            stagnation[k - 1] = 2 * (f - f_0).norm() / (f.norm() + f_0.norm())
            current_rank = len(f.core.data)
            output["sequence"][k - 1] = f

            if "error" in output_greedy:
                error[k - 1] = output_greedy["error"]
                is_error = True
                if self.alternating_minimization_parameters["display"]:
                    print(
                        "Alternating minimization (greedy): rank = %i, "
                        + "error = %2.5e, stagnation = %2.5e",
                        current_rank,
                        error[k - 1],
                        stagnation[k - 1],
                    )
            else:
                if self.alternating_minimization_parameters["display"]:
                    print(
                        "Alternating minimization (greedy): rank = %i, "
                        + "stagnation = %2.5e",
                        current_rank,
                        stagnation[k - 1],
                    )

            if (
                error[k - 1] < self.tolerance["on_error"]
                or stagnation[k - 1] < self.tolerance["on_stagnation"]
                or (
                    k > 2
                    and error[k - 1] > error[k - 2]
                    and error[k - 2] > error[k - 3]
                )
            ):
                break

            f_0 = deepcopy(f)
            if self.test_error:
                f_eval_test = tensap.FunctionalTensor(f, self.bases_eval_test)
                output["test_error"] = self.loss_function.test_error(
                    f_eval_test, self.test_data
                )
                output["test_error_iterations"].append(output["test_error"])
                if self.display:
                    print(
                        "Greedy: iteration %i, test error = %2.5e",
                        k,
                        output["test_error"],
                    )

        output["stagnation"] = stagnation[: k - 1]
        if is_error:
            output["errors"] = error[: k - 1]
            K = np.argmin(output["errors"])
            f = output["sequence"][K]
            output["selected_iterate"] = K
            output["error"] = output["errors"][K]

        if isinstance(self.bases, tensap.FunctionalBases):
            f = tensap.FunctionalTensor(f, self.bases)

        return f, output
