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

'''
Tutorial on functional PCA for low rank approximation of multivariate
functions.

'''

import numpy as np
import tensap

# %% Choice of the function to approximate
CHOICE = 2
if CHOICE == 1:
    print('Henon-Heiles')
    D = 5
    FUN, X = tensap.multivariate_functions_benchmark('henon_heiles', D)

    DEGREE = 4
    BASES = [tensap.PolynomialFunctionalBasis(x.orthonormal_polynomials(),
                                              range(DEGREE+1)) for
             x in X.random_variables]

elif CHOICE == 2:
    print('Anisotropic function')
    D = 6
    X = tensap.RandomVector(tensap.UniformRandomVariable(-1, 1), D)

    def FUN(x):
        return 1/(10 + 2*x[:, 0] + x[:, 2] + 2*x[:, 3] - x[:, 4])**2

    DEGREE = 13
    BASES = [tensap.PolynomialFunctionalBasis(x.orthonormal_polynomials(),
                                              range(DEGREE+1)) for
             x in X.random_variables]

elif CHOICE == 3:
    print('Sinus of a sum')
    D = 10
    X = tensap.RandomVector(tensap.UniformRandomVariable(-1, 1), D)

    def FUN(x):
        return np.sin(np.sum(x[:, :10], 1))

    DEGREE = 17
    BASES = [tensap.PolynomialFunctionalBasis(x.orthonormal_polynomials(),
                                              range(DEGREE+1)) for
             x in X.random_variables]

elif CHOICE == 4:
    print('Composition of functions')
    D = 6
    X = tensap.RandomVector(tensap.UniformRandomVariable(-1, 1), D)

    TREE = tensap.DimensionTree.balanced(D)

    def FUN(x1, x2):
        return 1 / (1 + x1**2 + x2**2)
    FUN = tensap.CompositionalModelFunction(TREE, FUN, X)

    DEGREE = 4
    BASES = [tensap.PolynomialFunctionalBasis(x.orthonormal_polynomials(),
                                              range(DEGREE+1)) for
             x in X.random_variables]

elif CHOICE == 5:
    print('Borehole function')
    FUN, X = tensap.multivariate_functions_benchmark('borehole')
    D = X.size

    DEGREE = 14
    BASES = [tensap.PolynomialFunctionalBasis(x.orthonormal_polynomials(),
                                              range(DEGREE+1)) for
             x in X.random_variables]

elif CHOICE == 6:
    print('Tensorized function')
    R = 11  # Resolution
    B = 2  # Scaling factor
    D = R+1

    X = tensap.UniformRandomVariable(0, 1)
    Y = tensap.UniformRandomVariable(0, 1)

    def IFUN(x):
        return 1/(1+x)

    T = tensap.Tensorizer(B, R, 1, X, Y)
    FUN = T.tensorize(IFUN)
    FUN.fun.evaluation_at_multiple_points = True
    DEGREE = 1
    H = tensap.PolynomialFunctionalBasis(Y.orthonormal_polynomials(),
                                         range(DEGREE+1))
    BASES = T.tensorized_function_functional_bases(H)
    X = tensap.RandomVector(BASES.measure)

else:
    raise ValueError('Bad function choice.')

# %% HOPCA (PCA for each dimension, provides reduced spaces)
# PCA for each dimension to get principal subspaces
print('--- Higher order PCA ---')
FPCA = tensap.FunctionalTensorPrincipalComponentAnalysis()
FPCA.bases = BASES
FPCA.pca_sampling_factor = 1
FPCA.tol = 1e-10
SUB_BASES, OUTPUTS = FPCA.hopca(FUN)
print('Number of evaluations = \n%s' % [x['number_of_evaluations'] for x
                                        in OUTPUTS])
print('Ranks {1, ..., d} = \n%s' % [x.cardinal() for x in SUB_BASES])

# %% Approximation in Tucker Format
# PCA for each dimension to get principal subspaces and interpolation on the
# tensor product of principal subspaces
print('\n--- Approximation in Tucker format ---')
FPCA = tensap.FunctionalTensorPrincipalComponentAnalysis()
FPCA.pca_sampling_factor = 1
FPCA.pca_adaptive_sampling = True
FPCA.bases = BASES

print('\nPrescribed ranks')
FPCA.tol = np.inf
FPCA.max_rank = np.random.randint(1, 5, D)
F, OUTPUT = FPCA.tucker_approximation(FUN)

print('Number of evaluations = %i' % OUTPUT['number_of_evaluations'])
print('Storage = %i' % F.storage())
print('Prescribed ranks = %s' % FPCA.max_rank)
print('Ranks = %s' % F.tensor.ranks[F.tensor.tree.dim2ind-1])
X_TEST = X.random(1e4)
F_X_TEST = F(X_TEST)
Y_TEST = FUN(X_TEST)
print('Error = %2.5e' % (np.linalg.norm(Y_TEST - F_X_TEST) /
                         np.linalg.norm(Y_TEST)))

print('\nPrescribed tolerance')
FPCA.tol = 1e-10
FPCA.max_rank = np.inf
F, OUTPUT = FPCA.tucker_approximation(FUN)

print('Number of evaluations = %i' % OUTPUT['number_of_evaluations'])
print('Ranks = %s' % F.tensor.ranks[F.tensor.tree.dim2ind-1])
X_TEST = X.random(1e4)
F_X_TEST = F(X_TEST)
Y_TEST = FUN(X_TEST)
print('Error = %2.5e' % (np.linalg.norm(Y_TEST - F_X_TEST) /
                         np.linalg.norm(Y_TEST)))

# %% Approximation in tree based format
print('\n--- Approximation in tree based format ---')
FPCA = tensap.FunctionalTensorPrincipalComponentAnalysis()
FPCA.pca_sampling_factor = 20
FPCA.bases = BASES
TREE = tensap.DimensionTree.balanced(D)

print('\nPrescribed ranks')
FPCA.tol = np.inf
FPCA.max_rank = np.random.randint(1, 9, TREE.nb_nodes)
FPCA.max_rank[TREE.root-1] = 1
F, OUTPUT = FPCA.tree_based_approximation(FUN, TREE)

print('Number of evaluations = %i' % OUTPUT['number_of_evaluations'])
print('Storage = %i' % F.storage())
print('Prescribed ranks = %s' % FPCA.max_rank)
print('Ranks = %s' % F.tensor.ranks)
X_TEST = X.random(1e3)
F_X_TEST = F(X_TEST)
Y_TEST = FUN(X_TEST)
print('Error = %2.5e' % (np.linalg.norm(Y_TEST - F_X_TEST) /
                         np.linalg.norm(Y_TEST)))

print('\nPrescribed tolerance')
FPCA.tol = 1e-10
FPCA.max_rank = np.inf
F, OUTPUT = FPCA.tree_based_approximation(FUN, TREE)

print('Number of evaluations = %i' % OUTPUT['number_of_evaluations'])
print('Storage = %i' % F.storage())
print('Ranks = %s' % F.tensor.ranks)
X_TEST = X.random(1e3)
F_X_TEST = F(X_TEST)
Y_TEST = FUN(X_TEST)
print('Error = %2.5e' % (np.linalg.norm(Y_TEST - F_X_TEST) /
                         np.linalg.norm(Y_TEST)))

# %% Approximation in Tensor Train format
print('\n--- Approximation in tensor train format ---')
FPCA = tensap.FunctionalTensorPrincipalComponentAnalysis()
FPCA.projection_type = 'interpolation'
FPCA.pca_sampling_factor = 3
FPCA.pca_adaptive_sampling = True
FPCA.bases = BASES

print('\nPrescribed ranks')
FPCA.tol = np.inf
FPCA.max_rank = np.random.randint(1, 9, D-1)
F, OUTPUT = FPCA.tt_approximation(FUN)

print('Number of evaluations = %i' % OUTPUT['number_of_evaluations'])
print('Storage = %i' % F.storage())
print('Prescribed ranks = %s' % FPCA.max_rank)
TT_RANKS = np.flip(F.tensor.ranks[F.tensor.is_active_node])
print('TT-ranks = %s' % TT_RANKS[:-1])
X_TEST = X.random(1e3)
F_X_TEST = F(X_TEST)
Y_TEST = FUN(X_TEST)
print('Error = %2.5e' % (np.linalg.norm(Y_TEST - F_X_TEST) /
                         np.linalg.norm(Y_TEST)))

print('\nPrescribed tolerance')
FPCA.tol = 1e-4
FPCA.max_rank = np.inf
F, OUTPUT = FPCA.tt_approximation(FUN)

print('Number of evaluations = %i' % OUTPUT['number_of_evaluations'])
print('Storage = %i' % F.storage())
TT_RANKS = np.flip(F.tensor.ranks[F.tensor.is_active_node])
print('TT-ranks = %s' % TT_RANKS[:-1])
X_TEST = X.random(1e3)
F_X_TEST = F(X_TEST)
Y_TEST = FUN(X_TEST)
print('Error = %2.5e' % (np.linalg.norm(Y_TEST - F_X_TEST) /
                         np.linalg.norm(Y_TEST)))
