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

"""
from copy import deepcopy
import numpy as np
import tensap


class CanonicalTensor:
    """
    Class CanonicalTensor.

    Attributes
    ----------
    core : tensap.DiagonalTensor
        The core of the canonical tensor.
    space : numpy.ndarray
        The tensor space.
    order : int
        The order of the tensor.
    shape : numpy.ndarray
        The shape of the tensor.
    is_orth : bool
        Boolean indicating, if is_orth = True, the dimension mu for which
        the mu-matricization of the tensor is orthogonal.

    """

    def __init__(self, space, core):
        """
        Constructor for the class CanonicalTensor

        Parameters
        ----------
        space : numpy.ndarray or list
        The tensor space.
        core : tensap.DiagonalTensor or list or numpy.ndarray
        The core of the canonical tensor.

        Returns
        -------
        None.

        """
        assert isinstance(
            space, (list, np.ndarray)
        ), "The input space must be a list or a numpy.ndarray."

        assert isinstance(core, (list, np.ndarray, tensap.DiagonalTensor)), (
            "The input core must be a list or a numpy.ndarray or a "
            + "tensap.DiagonalTensor."
        )

        if not isinstance(core, tensap.DiagonalTensor):
            core = tensap.DiagonalTensor(core, len(space))

        self.core = core
        self.space = space
        self.order = core.order
        self.shape = np.array([x.shape[0] for x in self.space])
        self.is_orth = False

    def __add__(self, arg):
        core = np.concatenate((self.core.data, arg.core.data))
        space = [np.hstack((x, y)) for x, y in zip(list(self.space), list(arg.space))]
        return CanonicalTensor(space, core)

    def __radd__(self, arg):
        return self + arg

    def __sub__(self, arg):
        return self + (-arg)

    def __rsub__(self, arg):
        return -self + arg

    def __neg__(self):
        return CanonicalTensor(list(self.space), -self.core.data)

    def __repr__(self):
        return (
            "<{} CanonicalTensor:{n}"
            + "{t}order = {},{n}"
            + "{t}shape = {},{n}"
            + "{t}is_orth = {}>"
        ).format(
            "x".join(map(str, self.shape)),
            self.order,
            self.shape,
            self.is_orth,
            t="\t",
            n="\n",
        )

    @property
    def ndim(self):
        """
        Compute the order of the tensor. Equivalent to self.order.

        Returns
        -------
        int
            The order of the tensor.

        """
        return self.order

    def orth(self):
        """
        Return an orthonormalized representation of the tensor.

        Returns
        -------
        out : CanonicalTensor
            The orthonormalized representation of the tensor.

        """
        dims = range(self.order)
        qr = [np.linalg.qr(x) for x in self.space]
        space = [x[0] for x in qr]
        M = [x[1] for x in qr]
        core = self.core.tensor_matrix_product(M, dims)
        core = core.orth()[0]
        out = CanonicalTensor(space, self.core)
        out.core = core
        return out

    def full(self):
        """
        Convert the object to a tensap.FullTensor.

        Returns
        -------
        tensap.FullTensor
            The canonical tensor as a tensap.FullTensor.

        """
        return self.core.full().tensor_matrix_product(self.space)

    def numpy(self):
        """
        Convert the CanonicalTensor to a numpy.ndarray.

        Returns
        -------
        numpy.ndarray
            The CanonicalTensor as a numpy.ndarray.

        """
        return self.full().data

    def tensor_vector_product(self, vectors, dims=None):
        """
        Compute the contraction of the tensor with vectors.

        Compute the contraction of self with each vector contained in the list
        vectors along dimensions specified by dims. The operation is such that
        vectors[k] is contracted with the dims[k]-th dimension of self.

        Parameters
        ----------
        vectors : numpy.ndarray or list of numpy.ndarray
            The vectors to use in the product.
        dims : list or numpy.ndarray, optional
            Indices of the contractions. The default is None, indicating all
            the dimensions.

        Returns
        -------
        CanonicalTensor
            The tensor after the contractions with the vectors.

        """
        if dims is None:
            assert isinstance(vectors, list), "vectors should be a list."
            assert len(vectors) == self.order, "len(vectors) must be self.order."
            dims = np.arange(self.order)
        else:
            dims = np.atleast_1d(dims)
            if not isinstance(vectors, list):
                vectors = [vectors]
            assert len(vectors) == dims.size, "len(vectors) must be equal to dims.size."

        vectors = [np.reshape(x, [1, -1]) for x in vectors]
        return self.tensor_matrix_product(vectors, dims).squeeze(dims.tolist())

    def squeeze(self, dims=None):
        """
        Remove the singleton dimensions of the tensor.

        Parameters
        ----------
        dims : list or numpy.ndarray, optional
            Dimensions to squeeze. The default is None, indicating all the
            singleton dimensions.

        Returns
        -------
        CanonicalTensor or scalar
            The squeezed tensor.

        """
        if dims is None:
            dims = np.arange(self.order)
            dims = dims[self.shape == 1]
        else:
            assert np.all(
                self.shape[dims] == 1
            ), "dimensions to squeeze should have a size 1."

        dims = np.atleast_1d(dims)
        dims = np.sort(dims)
        remaining_dims = tensap.fast_setdiff(np.arange(self.order), dims)

        if remaining_dims.size == 0:
            tensor = self.full().squeeze()
            if isinstance(tensor, tensap.FullTensor):
                tensor = tensor.data
        else:
            core = self.core.data
            space = list(self.space)
            for i, dim in enumerate(dims):
                core = core * space[dim]
            space = [space[i] for i in remaining_dims.tolist()]
            tensor = CanonicalTensor(space, core)

        return tensor

    def tensor_matrix_product(self, matrices, dims=None):
        """
        Contract a tensor with matrices.

        The second dimension of the matrix matrices[k] is contracted with the
        k-th dimension of self, with the indices k given in dims (if provided).

        Parameters
        ----------
        matrices : numpy.ndarray or list of numpy.ndarray
            The matrices to use in the product.
        dims : list or numpy.ndarray, optional
            Indices of the contractions. The default is None, indicating all
            the dimensions.

        Returns
        -------
        CanonicalTensor
            The tensor after the contractions with the matrices.

        """
        if dims is None:
            assert isinstance(
                matrices, (list, np.ndarray)
            ), "matrices should be a list or a numpy.ndarray."
            assert len(matrices) == self.order, "len(matrices) must be self.order."
            dims = range(self.order)
        else:
            dims = np.atleast_1d(dims)
            if not isinstance(matrices, list):
                matrices = [matrices]
            assert (
                len(matrices) == dims.size
            ), "len(matrices) must be equal to dims.size."

        space = list(self.space)
        for i, dim in enumerate(dims):
            space[dim] = np.matmul(matrices[i], space[dim])
        return CanonicalTensor(space, np.array(self.core.data))

    def tensor_matrix_product_eval_diag(self, matrices, dims=None):
        """
        Evaluate the diagonal of a tensor obtained by contraction with
        matrices.

        Provides the diagonal of the tensor obtained by contracting the tensor
        with matrices H[k] along dimensions dims(k)+1, for k = 0, ...,
        dims.size-1.

        Parameters
        ----------
        matrices : list
            The matrices to use in the product.
        dims : list or numpy.ndarray, optional
            Indices of the contractions. The default is None, indicating all
            the dimensions.

        Returns
        -------
        out : CanonicalTensor or tensap.FullTensor
            The diagonal of the contractions of the tensor with the matrices.

        """
        return self.tensor_matrix_product(matrices, dims).eval_diag(dims)

    def eval_diag(self, dims=None):
        """
        Extract the diagonal of the tensor.

        The tensor must be such that self.shape[mu] = n for all mu (in dims if
        provided).

        Parameters
        ----------
        dims : list of numpy.ndarray, optional
            The dimensions associated with the indices of the diagonal. The
            default is None,indicating that the indices refer to all the
            dimensions.

        Returns
        -------
        data : CanonicalTensor or tensap.FullTensor
            The evaluations of the diagonal of the tensor.

        """
        if dims is None:
            is_none = True
            dims = np.arange(self.order)
        else:
            is_none = False
            dims = np.atleast_1d(dims)
            if dims.size == 1:
                print(
                    "Only one dimension: degenerate case for eval_diag, "
                    + "returning the tensor itself."
                )
                return deepcopy(self)
            dims = np.sort(dims)
            new_dims = np.setdiff1d(range(self.order), dims[1:])

        out = self.space[dims[0]]
        for k in dims[1:]:
            out *= self.space[k]

        if is_none or dims.size == self.order:
            out = tensap.FullTensor(np.matmul(out, self.core.data))
        else:
            space = list(self.space)
            core = deepcopy(self.core)

            space[dims[0]] = out
            space = space[new_dims]
            core.shape = core.shape[new_dims]
            core.order = new_dims.size
            out = CanonicalTensor(space, core)
        return out

    def dot(self, tensor_2):
        """
        Return the inner product of two tensors.

        Parameters
        ----------
        tensor2 : CanonicalTensor
            The second tensor of the inner products.

        Returns
        -------
        numpy scalar
            The inner product of the two tensors.

        """
        assert isinstance(
            tensor_2, tensap.CanonicalTensor
        ), "The second argument must be a tensap.CanonicalTensor."
        matrices = [
            np.matmul(np.transpose(x), y) for x, y in zip(self.space, tensor_2.space)
        ]
        return self.core.dot_with_rank_one_metric(tensor_2.core, matrices)

    def norm(self, matrix=None):
        """
        Compute the canonical norm of the CanonicalTensor.

        Returns
        -------
        numpy scalar
            The norm of the tensor.

        """
        if matrix is None:
            if self.is_orth:
                norm = np.linalg.norm(self.core.data)
            else:
                norm = np.sqrt(np.abs(self.dot(self)))
        else:
            raise NotImplementedError("Method not implemented.")
        return norm

    def storage(self):
        """
        Return the storage complexity of the CanonicalTensor.

        Returns
        -------
        int
            The storage complexity of the CanonicalTensor.

        """
        return self.core.data.size + np.sum([x.size for x in self.space])

    def sparse_storage(self):
        """
        Return the sparse storage complexity of the CanonicalTensor.

        Returns
        -------
        int
            The sparse storage complexity of the CanonicalTensor.

        """
        return np.count_nonzero(self.core.data) + np.sum(
            [np.count_nonzero(x) for x in self.space]
        )

    def representation_rank(self):
        """
        Return the representation rank of the tensor.

        Returns
        -------
        int
            The representation rank of the tensor.

        """
        return self.core.data.size

    def parameter_gradient_eval_diag(self, mu, matrices=None):
        """
        Compute the diagonal of the gradient of the tensor with respect to a
        given parameter.

        Parameters
        ----------
        mu : int
            Index of the parameter.
        matrices : list or numpy.array, optional
            Matrices with which to compute outer_product_eval_diag if alpha is
            associated with some dimensions. Useful for evaluation the gradient
            of a tensap.FunctionalTensor. The default is None, indicating
            identity matrices.

        Returns
        -------
        out : tensap.FullTensor
            The diagonal of the gradient of the tensor with respect to
            the parameter with index mu.

        """
        rank = len(self.core.data)
        if mu == self.order + 1:
            N = self.space[0].shape[0]
            out = np.ones((N, rank))
            for nu in range(self.order):
                out *= self.space[nu]
            out = tensap.FullTensor(out)
        else:
            no_mu = np.setdiff1d(np.arange(1, self.order + 1), mu)
            N = self.space[no_mu[0] - 1].shape[0]
            f_mu = np.ones((N, rank))
            for nu in no_mu:
                f_mu *= self.space[nu - 1]
            if matrices is not None:
                out = tensap.FullTensor(f_mu).outer_product_eval_diag(
                    tensap.FullTensor(matrices[mu - 1]), 0, 0
                )
            else:
                out = tensap.FullTensor(f_mu).outer_product_eval_diag(
                    tensap.FullTensor(np.eye(self.shape[mu - 1])), [], [], True
                )
        return out

    def reduce_sum(self, dims=None):
        """
        Compute the sum of elements across dimensions dims of a tensor.

        Similar to tensorflow.reduce_sum.

        Parameters
        ----------
        dims : list or numpy.ndarray, optional
            The dimensions to be reduced. The default is None, indicating all
            the dimensions.

        Returns
        -------
        CanonicalTensor or a scalar if the sum is over all dimensions
            The reduced tensor.

        """

        if dims is None:
            dims = np.arange(self.order)
        else:
            dims = np.atleast_1d(dims)
        a = [np.ones(self.shape[i]) for i in dims]

        return self.tensor_vector_product(a, dims)

    def tree_based_tensor(self, tree, is_active_node=None):
        """
        Convert a CanonicalTensor to a tensap.TreeBasedTensor with given
        dimension tree and active nodes.

        Parameters
        ----------
        tree : tensap.DimensionTree
            The dimension tree.
        is_active_node : list or numpy.ndarray, optional
            Booleans indicating if the nodes are active. The default is None,
            settings all the nodes active.

        Returns
        -------
        out : tensap.TreeBasedTensor
            The canonical tensor converted into a tree-based tensor.

        """

        rank = self.core.shape[0]
        tensors = np.empty(tree.nb_nodes, dtype=object)
        for nod in np.arange(1, tree.nb_nodes + 1):
            ch = tree.children(nod)
            if tree.parent(nod) == 0:
                order = ch.size
                tensors[nod - 1] = tensap.FullTensor.diag(self.core.data, order)
            else:
                order = ch.size + 1
                tensors[nod - 1] = tensap.FullTensor.diag(np.ones(rank), order)

        tensors[tree.dim2ind - 1] = self.space
        out = tensap.TreeBasedTensor(tensors, tree)
        if is_active_node is not None:
            out = out.inactivate_nodes(
                np.nonzero(np.logical_not(is_active_node))[0] + 1
            )
        return out

    @staticmethod
    def create(generator, rank, shape):
        """
        Create a FullTensor of rank rank and shape shape using a given
        generator.

        Parameters
        ----------
        generator : function
            Function generating a numpy.ndarray, given a shape.
        rank : int
            The rank of the tensor.
        shape : numpy.ndarray or list
            The shape of the tensor.

        Returns
        -------
        CanonicalTensor
            The created tensor.

        """
        space = [generator([x, rank]) for x in shape]
        return CanonicalTensor(space, np.ones(rank))

    @staticmethod
    def zeros(rank, shape):
        """
        Create a FullTensor of rank rank and shape shape with entries equal to
        0.

        Parameters
        ----------
        rank : int
            The rank of the tensor.
        shape : numpy.ndarray or list
            The shape of the tensor.

        Returns
        -------
        CanonicalTensor
            The created tensor.

        """
        return CanonicalTensor.create(np.zeros, rank, shape)

    @staticmethod
    def ones(rank, shape):
        """
        Create a FullTensor of rank rank and shape shape with entries equal to
        1.

        Parameters
        ----------
        rank : int
            The rank of the tensor.
        shape : numpy.ndarray or list
            The shape of the tensor.

        Returns
        -------
        CanonicalTensor
            The created tensor.

        """
        return CanonicalTensor.create(np.ones, rank, shape)

    @staticmethod
    def rand(rank, shape):
        """
        Create a FullTensor of rank rank and shape shape with i.i.d. entries
        drawn according to the uniform distribution on [0, 1].

        Parameters
        ----------
        rank : int
            The rank of the tensor.
        shape : numpy.ndarray or list
            The shape of the tensor.

        Returns
        -------
        CanonicalTensor
            The created tensor.

        """
        return CanonicalTensor.create(lambda x: np.random.rand(*x), rank, shape)

    @staticmethod
    def randn(rank, shape):
        """
        Create a FullTensor of rank rank and shape shape with i.i.d. entries
        drawn according to the standard gaussian distribution.

        Parameters
        ----------
        rank : int
            The rank of the tensor.
        shape : numpy.ndarray or list
            The shape of the tensor.

        Returns
        -------
        CanonicalTensor
            The created tensor.

        """
        return CanonicalTensor.create(lambda x: np.random.randn(*x), rank, shape)
