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

'''

from copy import deepcopy
import numpy as np
import tensap


class SubFunctionalBasis(tensap.FunctionalBasis):
    '''
    Class SubFunctionalBasis.

    Attributes
    ----------
    underlying_basis : tensap.FunctionalBasis
        The underlying basis.

    basis : numpy.ndarray
        Array of shape (n, m), where n is the number of elements in
        underlying_basis, which defines a set of m basis functions in the space
        generated by underlying_basis.

    '''

    def __init__(self, underlying_basis=None, basis=None):
        '''
        Constructor for the class SubFunctionalBasis.

        Parameters
        ----------
        underlying_basis : tensap.FunctionalBasis
            The underlying basis.
        basis : numpy.ndarray
            Array of shape (n, m), where n is the number of elements in
            underlying_basis, which defines a set of m basis functions in the
            space generated by underlying_basis.

        Returns
        -------
        None.

        '''
        tensap.FunctionalBasis.__init__(self)
        self.underlying_basis = deepcopy(underlying_basis)
        self.basis = None
        if underlying_basis is not None:
            if basis is None:
                self.basis = np.zeros((self.underlying_basis.cardinal(), 0))
            else:
                self.basis = basis

            self.measure = underlying_basis.measure

            if underlying_basis.is_orthonormal and \
                np.linalg.norm(np.matmul(np.transpose(basis), basis) -
                               np.eye(basis.shape[1])) / \
                    np.sqrt(basis.shape[1]) < 1e-15:
                self.is_orthonormal = True

    def storage(self):
        return self.basis.size

    def cardinal(self):
        return self.basis.shape[1]

    def ndim(self):
        return self.underlying_basis.ndim()

    def domain(self):
        return self.underlying_basis.domain()

    def orthonormalize(self):
        '''
        Orthonormalize the SubFunctionalBasis.

        Returns
        -------
        tensap.SubFunctionalBasis
            The orthonormalized SubFunctionalBasis.

        '''
        assert self.underlying_basis.is_orthonormal, \
            'Method not implemented.'

        new_basis = np.array(self.basis)
        new_basis[:, 0] /= np.linalg.norm(new_basis[:, 0])
        for i in np.arange(1, new_basis.shape[1]):
            new_basis[:, i] -= np.matmul(new_basis[:, :i], np.matmul(
                np.transpose(new_basis[:, :i]), new_basis[:, i]))
            new_basis[:, i] /= np.linalg.norm(new_basis[:, i])

        return SubFunctionalBasis(self.underlying_basis, new_basis)

    def eval(self, x, indices=None):
        if indices is None:
            indices = np.arange(self.cardinal())

        return np.matmul(self.underlying_basis.eval(x), self.basis[:, indices])

    def eval_derivative(self, n, x):
        '''
        Evaluate the k-th order derivative of the functions of the basis at the
        points x.

        Parameters
        ----------
        k : int
            The order of the derivative.
        x : list or numpy.ndarray
            The points at which the k-th derivative of the basis functions are
            to be evaluated.

        Returns
        -------
        numpy.ndarray
            The evaluations of the k-th derivative of the functions of
            self.basis of degrees in self.indices at points x.

        '''
        return np.matmul(self.underlying_basis.eval_derivative(x, n),
                         self.basis)

    def derivative(self, n):
        '''
        Compute the k-th order derivative of the functions of the basis
        projected on itself.

        Parameters
        ----------
        k : int
            The order of the derivative.

        Returns
        -------
        tensap.SubFunctionalBasis
            The k-th order derivative of the functions of the basis projected
            on itself.

        '''
        d_f = deepcopy(self)
        d_f.underlying_basis = self.underlying_basis.derivative(n)
        return d_f

    def mean(self):
        return np.matmul(self.underlying_basis.mean(), self.basis)
