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

'''

from copy import deepcopy
from itertools import combinations
from random import shuffle
import numpy as np
import tensap


class TreeBasedTensorLearning(tensap.TensorLearning):
    '''
    Class TreeBasedTensorLearning.

    See also tensap.TensorLearning.

    Attributes
    ----------
    tree : tensap.DimensionTree
        The dimension tree of the tree-based tensor.
    is_active_node : list or numpy.ndarray
        Booleans indicating if each node is active.

    '''

    def __init__(self, tree, is_active_node, *args):
        '''
        Constructor for the class TreeBasedTensorLearning.

        See also tensap.TensorLearning.

        Parameters
        ----------
        tree : tensap.DimensionTree
            The dimension tree of the tree-based tensor.
        is_active_node : list or numpy.ndarray
            Booleans indicating if each node is active.
        *args : tuple
            Additional parameters.

        Returns
        -------
        None.

        '''
        super().__init__(*args)

        self.tree = tree
        self.is_active_node = list(map(bool, is_active_node))
        self.order = tree.dim2ind.size
        self._number_of_parameters = np.count_nonzero(is_active_node)

        self.initialization_type = 'canonical'
        self.rank_adaptation_options['rank_one_correction'] = True
        self.rank_adaptation_options['theta'] = 0.8
        self.linear_model_learning_parameters[
            'basis_adaptation_internal_nodes'] = False

# %% Standard solver methods
    def initialize(self):
        assert self.tree is not None, \
            'Must provide a DimensionTree object in property tree.'

        if np.isscalar(self.rank) or len(self.rank) == self.order:
            rank = np.zeros(self.tree.nb_nodes, dtype=int)
            rank[self.is_active_node] = self.rank
            self.rank = rank
        self.rank[self.tree.root-1] = self.output_dimension

        shape = [x.shape[1] for x in self.bases_eval]
        if self.initialization_type == 'random':
            f = tensap.TreeBasedTensor.randn(self.tree, self.rank, shape,
                                             self.is_active_node)
        elif self.initialization_type == 'ones':
            f = tensap.TreeBasedTensor.ones(self.tree, self.rank, shape,
                                            self.is_active_node)
        elif self.initialization_type == 'initial_guess':
            f = self.initial_guess
            if not np.all(f.ranks == self.rank):
                tr = tensap.Truncator(tolerance=np.finfo(float).eps,
                                      max_rank=self.rank)
                f = tr.truncate(f)
        elif self.initialization_type == 'mean' or \
                self.initialization_type == 'mean_randomized':
            if not np.all(self.rank == 1):
                raise NotImplementedError('Initialization only implemented ' +
                                          'if np.all(self.rank == 1).')
            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) 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 = [tensap.FullTensor(x, 2, [x.shape[0], 1]) for x in means]

            f = tensap.TreeBasedTensor.ones(self.tree, self.rank, shape)
            f.tensors[self.tree.dim2ind-1] = means
            nb_child = len(self.tree.children(self.tree.root))
            if np.ndim(self.training_data[1]) == 2:
                shape = np.concatenate((np.full(nb_child, 1),
                                       [self.training_data[1].shape[1]]))
            else:
                shape = np.full(nb_child, 1)
            f.tensors[self.tree.root-1] = \
                tensap.FullTensor(np.mean(self.training_data[1]), shape=shape)
            f.update_attributes()

            f = f.inactivate_nodes(
                np.nonzero(np.logical_not(self.is_active_node))[0]+1)
        elif self.initialization_type == 'canonical':
            if self.output_dimension != 1:
                print('Canonical initialization not implemented for ' +
                      'outputDimension > 1, performing a random ' +
                      'initialization.')
                f = tensap.TreeBasedTensor.randn(self.tree, self.rank, shape,
                                                 self.is_active_node)
            else:
                f = self.canonical_initialization(np.max(self.rank))
                if not np.all(f.ranks == self.rank):
                    tr = tensap.Truncator(tolerance=np.finfo(float).eps,
                                          max_rank=self.rank)
                    f = tr.truncate(f)
        else:
            raise ValueError('Wrong initialization type.')

        if not np.all(f.ranks == self.rank):
            f = TreeBasedTensorLearning.\
                enriched_edges_to_ranks_random(f, self.rank)

        # Exploration strategy of the tree by increasing level
        tree = f.tree
        exploration_strategy = np.zeros(self._number_of_parameters, dtype=int)
        active_nodes = f.active_nodes
        rep = 0
        for level in range(np.max(tree.level)+1):
            nodes = np.intersect1d(tree.nodes_with_level(level), active_nodes)
            exploration_strategy[rep:rep+len(nodes)] = nodes
            rep += len(nodes)
        self._exploration_strategy = exploration_strategy
        return self, f

    def pre_processing(self, f):
        if isinstance(self.linear_model_learning, (list, np.ndarray)) and \
                len(self.linear_model_learning) != f.tensor.tree.nb_nodes:
            tmp = np.empty(f.tensor.tree.nb_nodes, dtype=object)
            tmp[f.tensor.is_active_node] = self.linear_model_learning
            self.linear_model_learning = tmp
        return self, f

    def randomize_exploration_strategy(self):
        strategy = np.zeros(self._number_of_parameters)
        for level in np.arange(np.max(self.tree.level), -1, -1):
            active_nodes = np.intersect1d(self.tree.nodes_with_level(level),
                                          np.nonzero(self.is_active_node)[0])
            _, ind = np.intersect1d(self._exploration_strategy,
                                    active_nodes,
                                    return_indices=True)
            strategy[ind] = self._exploration_strategy[
                np.random.permutation(ind)]
        return strategy

    def prepare_alternating_minimization_system(self, f, mu):
        tree = f.tensor.tree
        if self.linear_model_learning[mu-1].basis_adaptation:
            if np.isin(mu, tree.internal_nodes):
                if self.linear_model_learning_parameters[
                        'basis_adaptation_internal_nodes']:
                    tr = tensap.Truncator(tolerance=np.finfo(float).eps,
                                          max_rank=np.max(f.tensor.ranks))
                    f.tensor = tr.hsvd(f.tensor)
                elif np.all(f.tensor.is_active_node[tree.children(mu)-1]):
                    self.linear_model_learning[mu-1].basis_adaptation = False
            f.tensor = f.tensor.orth_at_node(mu)
            self.tree = tree
            self.is_active_node = f.tensor.is_active_node
            self.linear_model_learning[mu-1].basis_adaptation_path = \
                self.create_basis_adaptation_path(f.tensor.ranks, mu)
        else:
            f.tensor = f.tensor.orth_at_node(mu)

        grad = f.parameter_gradient_eval(mu)
        if mu == tree.root:
            A = np.reshape(grad.data, [grad.shape[0], -1], order='F')
            self.linear_model_learning[mu-1].initial_guess = np.reshape(
                    f.tensor.tensors[mu-1].data,
                    [-1, f.tensor.ranks[tree.root-1]], order='F')
        else:
            A = np.reshape(grad.data, [grad.shape[0], -1,
                                       f.tensor.ranks[tree.root-1]], order='F')
            self.linear_model_learning[mu-1].initial_guess = np.reshape(
                    f.tensor.tensors[mu-1].data, -1, order='F')

        if isinstance(self.loss_function, tensap.DensityL2LossFunction):
            if isinstance(self.training_data, list) and \
                    len(self.training_data) == 2:
                y = self.training_data[1]
                if isinstance(y, tensap.FunctionalTensor):
                    y = y.tensor
                y = y.orth()
                if tree.is_leaf[mu-1]:
                    a = deepcopy(y)
                    for nod in tree.internal_nodes:
                        a.tensors[nod-1] = tensap.FullTensor(
                            y.tensors[nod-1].data *
                            f.tensor.tensors[nod-1].data,
                            shape=y.tensors[nod-1].shape)
                    ind = np.setdiff1d(range(self.order),
                                       np.nonzero(tree.dim2ind == mu)[0])
                    C = [x.data for x in f.tensor.tensors[tree.dim2ind-1]]
                    b = a.tensor_vector_product([C[x] for x in ind], ind)
                    b = b.tensors[0].data
                else:
                    b = f.tensor.dot(y) / f.tensor.tensors[mu-1].data
            else:
                b = []
        elif isinstance(self.training_data, list) and \
                len(self.training_data) == 2:
            b = self.training_data[1]
            self.linear_model_learning[mu-1].shared_coefficients = \
                mu != tree.root

        return self, A, b, f

    def set_parameter(self, f, mu, coef):
        f.tensor.tensors[mu-1] = \
            tensap.FullTensor(coef, shape=f.tensor.tensors[mu-1].shape)
        f.tensor.tensors[mu-1].is_orth = False
        return f

    def stagnation_criterion(self, f, f0):
        return (f.tensor - f0.tensor).norm() / f0.tensor.norm()

    def final_display(self, f):
        print('Ranks = [%s]' % ', '.join(map(str, f.tensor.ranks)), end='')

    def canonical_initialization(self, rank):
        '''
        Rank-r canonical initialization.

        Parameters
        ----------
        rank : int
            The rank of the canonical initialization.

        Returns
        -------
        tensap.TreeBasedTensor
            The rank-r canonical initialization..

        '''
        solver = tensap.CanonicalTensorLearning(self.order, self.loss_function)
        if isinstance(self.linear_model_learning, list):
            solver.linear_model_learning = self.linear_model_learning[0]
        else:
            solver.linear_model_learning = self.linear_model_learning
        solver.alternating_minimization_parameters = \
            deepcopy(self.alternating_minimization_parameters)
        solver.tolerance['on_stagnation'] = np.finfo(float).eps
        solver.tolerance['on_error'] = np.finfo(float).eps
        solver.bases = self.bases
        solver.bases_eval = self.bases_eval
        solver.bases_eval_test = self.bases_eval_test
        solver.display = False
        solver.alternating_minimization_parameters['display'] = False
        solver.initialization_type = 'mean'
        solver.rank_adaptation = True
        solver.rank_adaptation_options['max_iterations'] = rank
        solver.bases_adaptation_path = self.bases_adaptation_path
        solver.test_error = self.test_error
        solver.training_data = self.training_data
        solver.test_data = self.test_data
        solver._warnings = self._warnings

        f = solver.solve()[0]
        return f.tensor.tree_based_tensor(self.tree, self.is_active_node)

    def canonical_correction(self, f, rank):
        '''
        Rank-r canonical correction.

        Parameters
        ----------
        f : tensap.FunctionalTensor or None
            The current approximation.
        rank : int
            The rank of the canonical correction.

        Raises
        ------
        NotImplementedError
            If the method is not implemented.

        Returns
        -------
        f : tensap.FunctionalTensor
            The corrected approximation.

        '''
        if isinstance(f, tensap.FunctionalTensor):
            fx = f.tensor.tensor_matrix_product_eval_diag(self.bases_eval).data
        elif f is None:
            fx = 0
        else:
            raise NotImplementedError('Not implemented.')

        solver = deepcopy(self)
        if isinstance(solver.training_data, list) and \
                len(solver.training_data) == 2:
            solver.training_data[1] -= fx
        elif isinstance(solver.loss_function, tensap.DensityL2LossFunction):
            solver.training_data = [solver.training_data, f]

        f_add = solver.canonical_initialization(rank)
        if isinstance(f_add, tensap.FunctionalTensor):
            f_add = f_add.tensor
        if f is not None:
            f = f.tensor + f_add
        else:
            f = f_add
        return f

    def rank_one_correction(self, f):
        '''
        Rank one correction.

        Parameters
        ----------
        f : tensap.FunctionalTensor or None
            The current approximation.

        Raises
        ------
        NotImplementedError
            If the method is not implemented.

        Returns
        -------
        f : tensap.FunctionalTensor
            The corrected approximation.

        '''
        if isinstance(f, tensap.FunctionalTensor):
            fx = f.tensor.tensor_matrix_product_eval_diag(self.bases_eval).data
        elif f is None:
            fx = 0
        else:
            raise NotImplementedError('Not implemented.')

        solver = deepcopy(self)
        solver.model_selection = False
        if isinstance(solver.training_data, list) and \
                len(solver.training_data) == 2:
            solver.training_data[1] -= fx
        elif isinstance(solver.loss_function, tensap.DensityL2LossFunction):
            solver.training_data = [solver.training_data, f]
        solver.rank_adaptation = False
        solver.tree_adaptation = False
        solver.rank = 1
        solver.display = False
        solver.alternating_minimization_parameters['display'] = False
        solver.initialization_type = 'ones'
        solver.alternating_minimization_parameters['max_iterations'] = 1

        f_add = solver.solve()[0]
        if isinstance(f_add, tensap.FunctionalTensor):
            f_add = f_add.tensor
        if f is not None:
            f = f.tensor + f_add
        else:
            f = f_add
        return f

    def create_basis_adaptation_path(self, ranks, alpha):
        '''
        Creation of the basis adaptation path.

        Parameters
        ----------
        rank : list or numpy.ndarray
            The alpha-ranks of the current approximation.
        alpha : int
            The current node.

        Returns
        -------
        path : numpy.ndarray
            The basis adaptation path.

        '''
        tree = self.tree
        ranks = np.array(ranks)
        ranks[tree.root-1] = 1
        if tree.is_leaf[alpha-1]:
            path_alpha = self.bases_adaptation_path[
                np.nonzero(tree.dim2ind == alpha)[0][0]].astype(bool)
            ranks = ranks[alpha-1]
            path = path_alpha[np.newaxis, :, np.newaxis, :]
            path = np.tile(path, [1, 1, ranks, 2])
            path = np.reshape(path, [path.shape[1]*ranks, -1], order='F')
        else:
            assert not self.linear_model_learning_parameters[
                'basis_adaptation_internal_nodes'], \
                'Basis adaptation for internal nodes is not implemented.'
            ch = tree.children(alpha)
            if np.all(self.is_active_node[ch-1]):
                ch_a = ch[self.is_active_node[ch-1]]
                path = np.full(np.prod(ranks[ch_a-1]) * ranks[alpha-1], True)
            else:
                path_alpha = []
                for nod in ch:
                    if self.is_active_node[nod-1]:
                        path_alpha.append(np.full((ranks[nod-1], 1), True))
                    else:
                        ind = np.nonzero(nod == tree.dim2ind)[0][0]
                        path_alpha.append(
                            self.bases_adaptation_path[ind].astype(bool))
                path = path_alpha[-1]
                for ind in np.arange(len(path_alpha)-2, -1, -1):
                    path = np.kron(path, path_alpha[ind])
                path = np.tile(path, [ranks[alpha-1], 1]).astype(bool)
        return path

# %% Rank adaptation solver methods
    def local_solver(self):
        s_local = deepcopy(self)
        s_local.display = False
        s_local.rank_adaptation = False
        s_local.store_iterates = False
        s_local.test_error = False
        s_local.model_selection = False
        return s_local

    def new_rank_selection(self, f):
        if self.rank_adaptation_options['rank_one_correction']:
            s_local = deepcopy(self)
            ranks_add = np.ones(f.tensor.tree.nb_nodes, dtype=int)
            ranks_add[f.tensor.tree.root-1] = 0
            ranks_add[f.tensor.non_active_nodes-1] = 0
            s_local.rank = TreeBasedTensorLearning.make_ranks_admissible(
                f.tensor, f.tensor.ranks + ranks_add)[0]
            s_local.initialization_type = 'initial_guess'
            tr = tensap.Truncator(tolerance=0, max_rank=s_local.rank)
            s_local.initial_guess = tr.truncate(self.rank_one_correction(f))
            s_local.alternating_minimization_parameters['max_iterations'] = 10
            s_local.model_selection = False
            s_local.rank_adaptation = False
            s_local.display = False
            s_local.alternating_minimization_parameters['display'] = False
            tensor_for_selection = s_local.solve()[0].tensor
        else:
            tensor_for_selection = f.tensor

        sin_val = tensor_for_selection.singular_values()

        # Remove from the rank adaptation candidates: the inactive nodes, the
        # root, the leaf nodes with a rank equal to the dimension of the basis
        # associated to it, and the nodes for  which the smallest singular
        # value is almost zero.
        sin_val = np.array([np.nan if x is None else x for x in sin_val])
        sin_val[f.tensor.tree.root-1] = np.nan
        dim2ind = np.intersect1d(f.tensor.tree.dim2ind, f.tensor.active_nodes)
        ind = [len(set(x.size)) == 1 for x in f.tensor.tensors[dim2ind-1]]
        sin_val[dim2ind[ind]-1] = np.nan
        sin_val[s_local.rank != tensor_for_selection.ranks] = np.nan

        sin_val_min = np.array([np.min(x) for x in sin_val])
        sin_val_min[[x / tensor_for_selection.norm() < np.finfo(float).eps if
                     not np.isnan(x) else False for x in sin_val_min]] = np.nan

        # Remove nodes that cannot be enriched because their rank is equal to
        # the product of the ranks of their children, and their children cannot
        # be enriched themselves.
        tree = tensor_for_selection.tree
        rank = f.tensor.ranks
        desc = np.setdiff1d(np.arange(1, tree.nb_nodes+1),
                            np.nonzero(tree.is_leaf)[0]+1)
        cannot_be_increased = np.full(tree.nb_nodes, False)
        cannot_be_increased[tree.root-1] = True
        cannot_be_increased[tree.is_leaf] = np.isnan(sin_val_min[tree.is_leaf])
        for level in np.arange(np.max(tree.level), 0, -1):
            nod_lvl = np.intersect1d(tree.nodes_with_level(level), desc)
            for nod in nod_lvl:
                ch = tree.children(nod)
                if np.all(cannot_be_increased[ch-1]) and \
                        rank[nod-1] == np.prod(rank[ch-1]):
                    cannot_be_increased[nod-1] = True
        cannot_be_increased_nodes = tree.nodes_indices[cannot_be_increased]
        for level in np.arange(1, np.max(tree.level)):
            nod_lvl = np.intersect1d(tree.nodes_with_level(level),
                                     cannot_be_increased_nodes)
            for nod in nod_lvl:
                pa = tree.parent(nod)
                ind = np.setdiff1d(tree.children(pa), nod)
                ind = np.concatenate(([pa], ind))
                if np.all(cannot_be_increased[ind-1]) and \
                        rank[nod-1] == np.prod(rank[ind-1]):
                    cannot_be_increased[nod-1] = True
        sin_val_min[cannot_be_increased] = np.nan

        if np.all(np.isnan(sin_val_min)):
            enriched_nodes = np.array([], dtype=int)
        else:
            theta = self.rank_adaptation_options['theta'] * \
                np.nanmax(sin_val_min)
            enriched_nodes = np.nonzero([not np.isnan(x) and x >= theta for
                                         x in sin_val_min])[0] + 1

        new_rank = np.array(f.tensor.ranks)
        new_rank[enriched_nodes-1] += 1

        if not f.tensor.is_admissible_rank(new_rank):
            # Add to the already enriched nodes nodes one by one in decreasing
            # order of singular value until the rank is admissible.
            enriched_nodes_theta = np.array(enriched_nodes)
            rank_theta = np.array(new_rank)
            sin_val_min[enriched_nodes_theta-1] = np.nan
            sin_val_min_sorted = TreeBasedTensorLearning.unique_tol(
                sin_val_min, 1e-2)
            sin_val_min_sorted = np.flip(sin_val_min_sorted)
            sin_val_min_sorted = sin_val_min_sorted[
                np.isfinite(sin_val_min_sorted)]

            for sv in sin_val_min_sorted:
                new_rank = np.array(rank_theta)
                ind = [x >= sv if not np.isnan(x) else False for
                       x in sin_val_min]
                new_rank[ind] += 1
                if f.tensor.is_admissible_rank(new_rank):
                    enriched_nodes = np.concatenate((enriched_nodes_theta,
                                                     np.nonzero(ind)[0]+1))
                    break
            if not f.tensor.is_admissible_rank(new_rank):
                new_rank = f.tensor.ranks
                enriched_nodes = np.array([])

        return new_rank, enriched_nodes, tensor_for_selection

    def initial_guess_new_rank(self, s_local, f, new_rank):
        s_local.initialization_type = 'initial_guess'
        if not np.all(f.ranks == new_rank):
            tr = tensap.Truncator(tolerance=0, max_rank=new_rank)
            s_local.initial_guess = tr.truncate(f)
        else:
            s_local.initial_guess = f
        return s_local

    def adaptation_display(self, f, enriched_nodes):
        print('\tEnriched nodes: [%s]\n\tRanks = [%s]' %
              (', '.join(map(str, enriched_nodes)),
               ', '.join(map(str, f.tensor.ranks))), flush=True)

    def adapt_tree(self, f, cv_error, test_error, output, *args):
        if not self.tree_adaptation:
            return self, f, output

        output['adapted_tree'] = False

        if np.any(f.tensor.ranks[f.tensor.active_nodes-1] == 0):
            print('Some ranks equal to 0, disabling tree adaptation for ' +
                  'this step.')
            return self, f, output
        if self.tree_adaptation_options['tolerance'] is not None:
            adapt_tree_error = self.tree_adaptation_options['tolerance']
        elif self.loss_function.error_type == 'relative':
            if test_error is None or test_error == 0:
                adapt_tree_error = cv_error
            elif cv_error is None or test_error != 0:
                adapt_tree_error = test_error
        else:
            print('Must provide a tolerance for the tree adaptation in ' +
                  'the treeAdaptationOptions property. Disabling tree ' +
                  'adaptation.')
            self.tree_adaptation = False
            return self, f, output

        f_perm = f.tensor.optimize_dimension_tree(
            adapt_tree_error, self.tree_adaptation_options['max_iterations'])
        if f_perm.storage() < f.tensor.storage():
            f.tensor = f_perm
            self.tree = f.tensor.tree
            self.is_active_node = f.tensor.is_active_node
            output['adapted_tree'] = True
            if self.display:
                print('\tTree adaptation:\n\t\tRanks after permutation ' +
                      '= [%s]' % ', '.join(map(str, f.tensor.ranks)))
        return self, f, output

# %% Inner rank adaptation solver
    def _solve_dmrg_rank_adaptation(self):
        if 'max_rank' not in self.rank_adaptation_options:
            self.rank_adaptation_options['max_rank'] = 100
        if 'post_alternating_minimization' not in self.rank_adaptation_options:
            self.rank_adaptation_options['post_alternating_minimization'] = \
                False
        if self.rank_adaptation_options['type'] == 'dmrg_low_rank' and \
                'model_selection_type' not in self.rank_adaptation_options:
            self.rank_adaptation_options['model_selection_type'] = 'cv_error'

        if self.display:
            self.alternating_minimization_parameters['display'] = True

        output = {'flag': 0}

        self, f = self.initialize()
        f = tensap.FunctionalTensor(f, self.bases_eval)

        # Exploration strategy of the tree by decreasing level
        tree = f.tensor.tree
        exploration_strategy = np.zeros(self._number_of_parameters, dtype=int)
        active_nodes = f.tensor.active_nodes
        rep = 0
        for level in np.arange(np.max(tree.level), -1, -1):
            nodes = np.intersect1d(tree.nodes_with_level(level), active_nodes)
            exploration_strategy[rep:rep+len(nodes)] = nodes
            rep += len(nodes)
        self._exploration_strategy = np.setdiff1d(exploration_strategy,
                                                  tree.root,
                                                  assume_unique=True)

        # 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.')

        if self.error_estimation:
            for x in self.linear_model_learning:
                setattr(x, 'error_estimation', True)

        # 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()

        if self.alternating_minimization_parameters['max_iterations'] == 0:
            return f, output

        output['stagnation_criterion'] = []
        output['iterates'] = []
        output['error_iterations'] = []
        output['test_error_iterations'] = []

        # Alternating minimization loop
        for iteration in range(self.alternating_minimization_parameters
                               ['max_iterations']):
            self, f = self.pre_processing(f)
            f0 = deepcopy(f)

            if self.alternating_minimization_parameters['random']:
                # Randomize the exploration strategy
                alpha_list = self.randomize_exploration_strategy()
            else:
                alpha_list = self._exploration_strategy

            tree = f.tensor.tree
            for alpha in alpha_list:
                if self.linear_model_learning[alpha-1].basis_adaptation:
                    if np.isin(alpha, tree.internal_nodes):
                        if self.linear_model_learning_parameters[
                                'basis_adaptation_internal_nodes']:
                            tr = tensap.Truncator(
                                tolerance=np.finfo(float).eps,
                                max_rank=np.max(f.tensor.ranks))
                            f.tensor = tr.hsvd(f.tensor)
                        elif np.all(f.tensor.is_active_node[
                                tree.children(alpha)-1]):
                            self.linear_model_learning[alpha-1].\
                                basis_adaptation = False
                    f.tensor = f.tensor.orth_at_node(tree.parent(alpha))
                    self.tree = tree
                    self.is_active_node = f.tensor.is_active_node

                    if self.rank_adaptation_options['type'] == 'dmrg':
                        self.linear_model_learning[alpha-1].\
                            basis_adaptation_path = \
                            self.create_basis_adaptation_path_dmrg(
                                f.tensor.ranks, alpha)
                    elif self.rank_adaptation_options['type'] == \
                            'dmrg_low_rank':
                        self.linear_model_learning[alpha-1].\
                            basis_adaptation_path = self.\
                            create_basis_adaptation_path_dmrg_low_rank(
                                f.tensor.ranks, alpha)
                    else:
                        raise ValueError('Wrong rank adaptation type.')
                else:
                    f.tensor = f.tensor.orth_at_node(tree.parent(alpha))

                grad = f.parameter_gradient_eval_dmrg(
                    alpha, dmrg_type=self.rank_adaptation_options['type'])

                if isinstance(self.loss_function,
                              tensap.DensityL2LossFunction):
                    if isinstance(self.training_data, list) and \
                            len(self.training_data) == 2:
                        y = self.training_data[1]
                        if isinstance(y, tensap.FunctionalTensor):
                            y = y.tensor
                        y = y.orth()
                        if tree.is_leaf[alpha-1]:
                            a = deepcopy(y)
                            for nod in tree.internal_nodes:
                                a.tensors[nod-1] = tensap.FullTensor(
                                    y.tensors[nod-1].data *
                                    f.tensor.tensors[nod-1].data,
                                    shape=y.tensors[nod-1].shape)
                            ind = np.setdiff1d(
                                range(self.order),
                                np.nonzero(tree.dim2ind == alpha)[0])
                            C = [x.data for
                                 x in f.tensor.tensors[tree.dim2ind-1]]
                            b = a.tensor_vector_product([C[x] for x in ind],
                                                        ind)
                            b = b.tensors[0].data
                        else:
                            b = f.tensor.dot(y) / \
                                f.tensor.tensors[alpha-1].data
                    else:
                        b = []
                elif isinstance(self.training_data, list) and \
                        len(self.training_data) == 2:
                    b = self.training_data[1]
                    self.linear_model_learning[alpha-1].shared_coefficients = \
                        alpha != tree.root

                gamma = tree.parent(alpha)
                ind = np.setdiff1d(
                    np.arange(1, f.tensor.tensors[gamma-1].order+1),
                    tree.child_number(alpha))

                if self.rank_adaptation_options['type'] == 'dmrg':
                    A = np.reshape(grad.data, (grad.shape[0], -1), order='F')

                    self.linear_model_learning[alpha-1].training_data = [None,
                                                                         b]
                    self.linear_model_learning[alpha-1].basis = None
                    self.linear_model_learning[alpha-1].basis_eval = A

                    C, output_tmp = \
                        self.linear_model_learning[alpha-1].solve()

                    if C is None or np.count_nonzero(C) == 0 or \
                            not np.all(np.isfinite(C)):
                        print('Empty, zero or NaN solution, returning to ' +
                              'the previous iteration.')
                        output['flag'] = -2
                        output['error'] = np.inf
                        break

                    sz_1 = np.prod(f.tensor.tensors[alpha-1].shape[:-1])
                    sz_2 = np.prod(f.tensor.tensors[gamma-1].shape[ind-1])
                    C = tensap.FullTensor(C, 2, [sz_1, sz_2])

                    tr = tensap.Truncator()
                    tr.tolerance = self.tolerance['on_error'] / \
                        np.sqrt(np.count_nonzero(f.tensor.is_active_node)-1)
                    tr.max_rank = self.rank_adaptation_options['max_rank']
                    C = tr.truncate(C)
                    rank = [np.shape(C.space[0])[1]]

                    sz_1 = np.concatenate(
                        (f.tensor.tensors[alpha-1].shape[:-1], rank))
                    sz_2 = np.concatenate(
                        (f.tensor.tensors[gamma-1].shape[ind-1], rank))

                    a_alpha = np.reshape(C.space[0], sz_1, order='F')
                    a_gamma = np.reshape(C.space[1]*np.tile(
                        np.reshape(C.core.data, [1, -1]),
                        (np.shape(C.space[1])[0], 1)), sz_2, order='F')
                    perm = np.concatenate((ind, [tree.child_number(alpha)]))
                    a_gamma = np.transpose(a_gamma, np.argsort(perm-1))
                    sz_2[perm-1] = np.array(sz_2)
                elif self.rank_adaptation_options['type'] == 'dmrg_low_rank':
                    A = [np.reshape(x.data, [x.shape[0], -1], order='F') for
                         x in grad]

                    s_local = TreeBasedTensorLearning(
                        tensap.DimensionTree.trivial(2), [True, True, False],
                        self.loss_function)
                    s_local.rank_adaptation = True
                    s_local.tolerance['on_error'] = \
                        self.tolerance['on_error'] / \
                        np.sqrt(np.count_nonzero(f.tensor.is_active_node)-1)
                    s_local.tolerance['on_stagnation'] = \
                        self.tolerance['on_stagnation']
                    s_local.rank_adaptation_options['max_iterations'] = \
                        self.rank_adaptation_options['max_rank']
                    s_local.alternating_minimization_parameters = \
                        deepcopy(self.alternating_minimization_parameters)
                    s_local.alternating_minimization_parameters['display'] = \
                        False
                    s_local.store_iterates = True
                    s_local.test_error = False
                    s_local.error_estimation = True
                    s_local.display = False
                    s_local.order = 2
                    s_local.linear_model_learning = \
                        self.linear_model_learning[alpha-1]
                    s_local._warnings['orthonormality_warning_display'] = False
                    s_local._warnings['empty_bases_warning_display'] = False
                    s_local.bases_adaptation_path = \
                        self.linear_model_learning[alpha-1].\
                        basis_adaptation_path
                    s_local.training_data = self.training_data
                    s_local.bases_eval = A
                    s_local.model_selection = True
                    s_local.model_selection_options['type'] = \
                        self.rank_adaptation_options['model_selection_type']
                    C, output_tmp = s_local.solve()

                    rank = [C.tensor.ranks[1]]
                    sz_1 = np.concatenate(
                        (f.tensor.tensors[alpha-1].shape[:-1], rank))
                    sz_2 = np.concatenate(
                        (f.tensor.tensors[gamma-1].shape[ind-1], rank))

                    a_alpha = C.tensor.tensors[1].reshape(sz_1)
                    a_gamma = C.tensor.tensors[0].transpose([1, 0]).reshape(
                        sz_2)

                    perm = np.concatenate((ind, [tree.child_number(alpha)]))
                    a_gamma = a_gamma.itranspose(perm-1)
                    sz_2[perm-1] = np.array(sz_2)

                    a_alpha = a_alpha.data
                    a_gamma = a_gamma.data

                    if (a_alpha is None or np.count_nonzero(a_alpha) == 0 or
                            not np.all(np.isfinite(a_alpha))) or \
                        (a_gamma is None or np.count_nonzero(a_gamma) == 0 or
                            not np.all(np.isfinite(a_gamma))):
                        print('Empty, zero or NaN solution, returning to ' +
                              'the previous iteration.')
                        output['flag'] = -2
                        output['error'] = np.inf
                        break
                else:
                    raise ValueError('Wrong rank adaptation type.')
                f.tensor.tensors[alpha-1] = tensap.FullTensor(a_alpha,
                                                              np.size(sz_1),
                                                              sz_1)
                f.tensor.tensors[gamma-1] = tensap.FullTensor(a_gamma,
                                                              np.size(sz_2),
                                                              sz_2)

            if self.rank_adaptation_options['post_alternating_minimization']:
                print('\t\tPost alternating minimization.')
                s_local = TreeBasedTensorLearning(self.tree,
                                                  self.is_active_node,
                                                  self.loss_function)
                s_local.bases_eval = self.bases_eval
                s_local.bases_adaptation_path = self.bases_adaptation_path
                s_local.training_data = self.training_data
                s_local.test_error = self.test_error
                s_local.test_data = self.test_data
                s_local.bases_eval_test = self.bases_eval_test
                s_local.model_selection = False
                s_local.rank_adaptation = False
                s_local.store_iterates = False
                s_local.rank = f.tensor.ranks
                s_local.initialization_type = 'initial_guess'
                s_local.initial_guess = f.tensor
                s_local.tolerance['on_error'] = self.tolerance['on_error'] / \
                    np.sqrt(np.count_nonzero(f.tensor.is_active_node)-1)
                s_local.tolerance['on_stagnation'] = \
                    self.tolerance['on_stagnation']
                s_local.alternating_minimization_parameters = \
                    deepcopy(self.alternating_minimization_parameters)
                s_local.alternating_minimization_parameters['display'] = False
                s_local.error_estimation = True
                s_local.display = False
                s_local.linear_model_learning = \
                    [x for x in self.linear_model_learning if x is not None]
                s_local._warnings['orthonormality_warning_display'] = False
                s_local._warnings['empty_bases_warning_display'] = False
                f, output_tmp = s_local.solve()

            stagnation = self.stagnation_criterion(f, f0)
            output['stagnation_criterion'].append(stagnation)

            if self.store_iterates:
                if isinstance(self.bases, tensap.FunctionalBases):
                    output['iterates'].append(tensap.FunctionalTensor(
                        f.tensor, self.bases))
                else:
                    output['iterates'].append(f)

            if 'error' in output_tmp:
                output['error'] = output_tmp['error']
                output['error_iterations'].append(output['error'])

            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.alternating_minimization_parameters['display']:
                print('\tAlt. min. iteration %s/%i: stagnation = %2.5e' %
                      (str(iteration).
                       zfill(len(str(self.alternating_minimization_parameters
                                     ['max_iterations']-1))),
                       self.alternating_minimization_parameters
                       ['max_iterations']-1,
                       stagnation), end='')
                if 'error' in output:
                    print(', error = %2.5e' % output['error'], end='')
                if self.test_error:
                    if not np.isscalar(output['test_error']):
                        output['test_error'] = output['test_error'].numpy()
                    print(', test error = %2.5e' % output['test_error'],
                          end='')
                print('')

            if iteration > 0 and stagnation < \
                    self.alternating_minimization_parameters['stagnation']:
                output['flag'] = 1
                break

            if self.test_error and \
                    output['test_error'] < self.tolerance['on_error']:
                output['flag'] = 1
                break

            if self.tree_adaptation and iteration > 0:
                C_old = f.tensor.storage()
                self, f, output = self.adapt_tree(f, output['error'],
                                                  None, output, iteration)
                adapted_tree = output['adapted_tree']
                if adapted_tree:
                    if self.display:
                        print('\t\tStorage complexity before permutation ' +
                              '= %i' % C_old)
                        print('\t\tStorage complexity after permutation ' +
                              '= %i' % f.tensor.storage())
                    if self.test_error:
                        f_eval_test = tensap.FunctionalTensor(
                            f, self.bases_eval_test)
                        if self.display:
                            print('\t\tTest error after permutation ' +
                                  '= %2.5e' % self.loss_function.test_error(
                                      f_eval_test, self.test_data))

                    if self.alternating_minimization_parameters['display']:
                        print('')

        if isinstance(self.bases, tensap.FunctionalBases):
            f = tensap.FunctionalTensor(f.tensor, self.bases)
        output['iter'] = iteration

        if 'adapted_tree' in output:
            del output['adapted_tree']

        if self.display and not self.model_selection:
            if self.alternating_minimization_parameters['display']:
                print('')
            self.final_display(f)
            if 'error' in output:
                print(', error = %2.5e' % output['error'], end='')
            if 'test_error' in output:
                print(', test error = %2.5e' % output['test_error'], end='')
            print('')

        return f, output

    def create_basis_adaptation_path_dmrg(self, ranks, alpha):
        tree = self.tree
        ranks = np.array(ranks)
        ranks[tree.root-1] = 1
        gamma = tree.parent(alpha)

        p_alpha = self.create_basis_adaptation_path(ranks, alpha)
        if p_alpha.ndim == 1:
            p_alpha = np.reshape(p_alpha, [-1, 1])
        p_alpha = p_alpha[:int(p_alpha.shape[0]/ranks[alpha-1]), :]

        ch = tree.children(gamma)
        p_gamma = [None]*len(ch)
        for nod in np.setdiff1d(ch, alpha):
            ind = tree.child_number(nod) < tree.child_number(alpha)
            if self.is_active_node[nod-1]:
                p_gamma[tree.child_number(nod)-1+ind] = \
                    np.full((ranks[nod-1], 1), True)
            else:
                rep = np.nonzero(nod == tree.dim2ind)[0][0]
                p_gamma[tree.child_number(nod)-1+ind] = \
                    self.bases_adaptation_path[rep].astype(bool)
        p_gamma[0] = p_alpha

        r_gamma = ranks[gamma-1]
        p = p_gamma[-1]
        for i in np.arange(len(p_gamma)-2, -1, -1):
            p = np.kron(p, p_gamma[i])
        return np.tile(p, (r_gamma, 1)).astype(bool)

    def create_basis_adaptation_path_dmrg_low_rank(self, ranks, alpha):
        tree = self.tree
        ranks = np.array(ranks)
        ranks[tree.root-1] = 1
        gamma = tree.parent(alpha)

        p_alpha = self.create_basis_adaptation_path(ranks, alpha)
        if p_alpha.ndim == 1:
            p_alpha = np.reshape(p_alpha, [-1, 1])
        p_alpha = p_alpha[:int(p_alpha.shape[0]/ranks[alpha-1]), :]

        ch = np.setdiff1d(tree.children(gamma), alpha)
        p_gamma = [None]*len(ch)
        for nod in ch:
            ind = tree.child_number(nod) > tree.child_number(alpha)
            if self.is_active_node[nod-1]:
                p_gamma[tree.child_number(nod)-1-ind] = \
                    np.full((ranks[nod-1], 1), True)
            else:
                rep = np.nonzero(nod == tree.dim2ind)[0][0]
                p_gamma[tree.child_number(nod)-1-ind] = \
                    self.bases_adaptation_path[rep].astype(bool)

        r_gamma = ranks[gamma-1]
        p = p_gamma[-1]
        for i in np.arange(len(p_gamma)-2, -1, -1):
            p = np.kron(p, p_gamma[i])
        p = np.tile(p, (r_gamma, 1)).astype(bool)

        return [p_alpha, p]

# %% Static methods
    @staticmethod
    def tensor_train(order, *args):
        '''
        Call of the constructor of the class TreeBasedTensorLearning, with a
        tree and active nodes corresponding to the Tensor-Train format in
        dimension order.

        See also TreeBasedTensorLearning.

        Parameters
        ----------
        order : int
            The order of the tensor.
        *args : tuple
            Additional parameters (see the constructor of
            TreeBasedTensorLearning).

        Returns
        -------
        TreeBasedTensorLearning
            The solver with a tree and active nodes associated with the
            Tensor-Train format.

        '''
        tree = tensap.DimensionTree.linear(order)
        is_active_node = np.full(tree.nb_nodes, True)
        is_active_node[tree.dim2ind[1:]-1] = False
        return TreeBasedTensorLearning(tree, is_active_node, *args)

    @staticmethod
    def tensor_train_tucker(order, *args):
        '''
        Call of the constructor of the class TreeBasedTensorLearning, with a
        tree and active nodes corresponding to the Tensor-Train Tucker format
        in dimension order.

        See also TreeBasedTensorLearning.

        Parameters
        ----------
        order : int
            The order of the tensor.
        *args : tuple
            Additional parameters (see the constructor of
            TreeBasedTensorLearning).

        Returns
        -------
        TreeBasedTensorLearning
            The solver with a tree and active nodes associated with the
            Tensor-Train Tucker format.

        '''
        tree = tensap.DimensionTree.linear(order)
        is_active_node = np.full(tree.nb_nodes, True)
        return TreeBasedTensorLearning(tree, is_active_node, *args)

    @staticmethod
    def enriched_edges_to_ranks_random(f, new_rank):
        '''
        Enrichment of the ranks of specified edges of the tensor f using random
        additions for each child / parent couple of the enriched edges.

        Parameters
        ----------
        f : tensap.TreeBasedTensor
            The tree-based tensor to enrich.
        new_rank : list or numpy.ndarray
            The new tree-based rank.

        Returns
        -------
        f : tensap.TreeBasedTensor
            The enriched tree-based tensor.

        '''
        f.is_orth = False
        tree = f.tree
        enriched_dims = np.nonzero(new_rank > f.ranks)[0]

        for level in np.arange(1, np.max(tree.level)+1):
            nod_lvl = np.intersect1d(tree.nodes_with_level(level),
                                     enriched_dims)
            for alpha in nod_lvl:
                gamma = tree.parent(alpha)
                rank = new_rank[alpha-1] - f.ranks[alpha-1]

                A = np.reshape(f.tensors[alpha-1].data,
                               [-1, f.ranks[alpha-1]], order='F')
                A = np.hstack((A, np.tile(np.reshape(A[:, -1], [-1, 1]),
                                          [1, rank]) *
                               (1+np.random.randn(A.shape[0], rank))))
                A[:, -1-rank-1:] /= np.sqrt(np.sum(A[:, -1-rank-1:]**2, 0))
                shape = np.array(f.tensors[alpha-1].shape)
                shape[-1] += rank
                f.tensors[alpha-1].data = np.reshape(A, shape, order='F')

                ch = f.tree.child_number(alpha)-1
                ind = np.setdiff1d(range(f.tensors[gamma-1].order), ch)
                ind = np.concatenate((ind, [ch]))
                A = np.transpose(f.tensors[gamma-1].matricize(ch).data)
                A = np.hstack((A, np.tile(np.reshape(A[:, -1], [-1, 1]),
                                          [1, rank]) *
                               (1+np.random.randn(A.shape[0], rank))))
                A[:, -1-rank-1:] /= np.sqrt(np.sum(A[:, -1-rank-1:]**2, 0))
                shape = np.array(f.tensors[gamma-1].shape)
                shape[ch] += rank
                A = np.reshape(A, shape[ind], order='F')
                f.tensors[gamma-1].data = np.transpose(A, np.argsort(ind))

                f = f.update_attributes()
        return f

    @staticmethod
    def make_ranks_admissible(f, rank):
        '''
        Adjustment of the ranks to make the associated tree-based tensor f
        rank-admissible, by enriching new edges associated with nodes of the
        tree until all the rank admissibility conditions are met.

        Parameters
        ----------
        f : tensap.TreeBasedTensor
            The current approximation in tree-based tensor format.
        rank : numpy.ndarray
            The proposed tree-based rank.

        Returns
        -------
        rank : numpy.ndarray
            The admissible tree-based rank (if possible).
        d : numpy.ndarray
            The enriched nodes.

        '''
        # Do not increase the ranks of leaf nodes with rank equal to the
        # dimension of the approximation space.
        nodes = f.active_nodes
        ind = [f.tree.is_leaf[x] and rank[x] > f.tensors[x].shape[0] for
               x in nodes-1]
        rank[nodes[ind]-1] = [x.shape[0] for x in f.tensors[nodes[ind]-1]]
        rank[np.logical_not(f.is_active_node)] = 0

        delta = rank - f.ranks
        if f.is_admissible_rank(f.ranks + delta):
            rank = f.ranks + delta
            d = np.nonzero(delta)[0] + 1
            return rank, d

        ind = np.nonzero(delta)[0]
        for i in np.arange(1, np.count_nonzero(delta)+1):
            pos = list(combinations(np.arange(np.count_nonzero(delta)), i))
            shuffle(pos)
            for pos_loc in pos:
                delta_loc = np.array(delta)
                delta_loc[ind[list(pos_loc)]] = 0
                if f.is_admissible_rank(f.ranks + delta_loc):
                    rank = f.ranks + delta_loc
                    d = np.nonzero(delta_loc)[0] + 1
                    return rank, d
        return rank, []

    @staticmethod
    def unique_tol(inp, tol):
        '''
        Unique values within tolerance, with sorted output.

        Parameters
        ----------
        inp : list of numpy.ndarray
            The values to be checked.
        tol : float
            The tolerance defining the uniqueness.

        Returns
        -------
        bout : numpy.ndarray
            The unique values of inp, within the tolerance tol.

        '''
        inp = np.sort(inp)
        out = [inp[0]]

        for ind in np.arange(1, inp.size):
            if np.abs(out[-1] - inp[ind]) / out[-1] > tol:
                out = np.concatenate((out, [inp[ind]]))
        return out
