# 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 learning in tree-based tensor format a tensorized function, with
an example of model selection using a slope heuristic instead of the error on a
test sample.

'''

from time import time
import numpy as np
import matplotlib.pyplot as plt
import tensap

# %% Function to approximate: identification of a function f(x) with a function
# g(i1,...,id,y) (see also tutorial_TensorizedFunction)
R = 4  # Resolution
B = 5  # Scaling factor
ORDER = R+1

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

CHOICE = 1
if CHOICE == 1:
    def FUN(x):
        return np.sin(10*np.pi*(2*x+0.5))/(4*x+1) + (2*x-0.5)**4
elif CHOICE == 2:
    def FUN(x):
        return (np.sin(4*np.pi*x) + 0.2*np.cos(16*np.pi*x))**(x < 0.5) + \
            (2*x-1)*(x >= 0.5)
else:
    raise ValueError('Bad function choice.')

T = tensap.Tensorizer(B, R, 1, X, Y)
TENSORIZED_FUN = T.tensorize(FUN)
TENSORIZED_FUN.fun.evaluation_at_multiple_points = True

# %% Approximation basis
DEGREE = 5
H = tensap.PolynomialFunctionalBasis(Y.orthonormal_polynomials(),
                                     range(DEGREE+1))
BASES = T.tensorized_function_functional_bases(H)

# %% Training and test samples
NUM_TRAIN = 200
X_TRAIN = X.random(NUM_TRAIN)
Y_TRAIN = FUN(X_TRAIN)
X_TRAIN = T.map(X_TRAIN)  # Identification of X_TRAIN with (i_1,...,i_d,y)

NUM_TEST = 1000
X_TEST = X.random(NUM_TEST)
Y_TEST = FUN(X_TEST)
X_TEST = T.map(X_TEST)

# %% Tree-based tensor format
# Tensor format
# 1 - Random tree and active nodes
# 2 - Tensor-Train
# 3 - Hierarchial Tensor-Train
# 4 - Binary tree
CHOICE = 3
if CHOICE == 1:
    print('Random tree with active nodes')
    ARITY = [2, 4]
    TREE = tensap.DimensionTree.random(ORDER, ARITY)
    IS_ACTIVE_NODE = np.full(TREE.nb_nodes, True)
    SOLVER = tensap.TreeBasedTensorLearning(TREE, IS_ACTIVE_NODE,
                                            tensap.SquareLossFunction())
elif CHOICE == 2:
    print('Tensor-train format')
    SOLVER = tensap.TreeBasedTensorLearning.tensor_train(
        ORDER, tensap.SquareLossFunction())
elif CHOICE == 3:
    print('Tensor Train Tucker')
    SOLVER = tensap.TreeBasedTensorLearning.tensor_train_tucker(
        ORDER, tensap.SquareLossFunction())
elif CHOICE == 4:
    print('Binary tree')
    TREE = tensap.DimensionTree.balanced(ORDER)
    IS_ACTIVE_NODE = np.full(TREE.nb_nodes, True)
    SOLVER = tensap.TreeBasedTensorLearning(TREE, IS_ACTIVE_NODE,
                                            tensap.SquareLossFunction())
else:
    raise NotImplementedError('Not implemented.')

# %% Random shuffling of the dimensions associated to the leaves
RANDOMIZE = True
if RANDOMIZE:
    SOLVER.tree.dim2ind = np.random.permutation(SOLVER.tree.dim2ind)
    SOLVER.tree = SOLVER.tree.update_dims_from_leaves()

# %% Learning in tree-based tensor format
SOLVER.bases = BASES
SOLVER.bases_eval = BASES.eval(X_TRAIN)
SOLVER.training_data = [None, Y_TRAIN]

# Let the solver go to the maximum number of iterations, to perform model
# selection a posteriori using the test error and a slope heuristic
SOLVER.tolerance['on_stagnation'] = 0
SOLVER.tolerance['on_error'] = 0
SOLVER.rank_adaptation_options['early_stopping'] = False

SOLVER.initialization_type = 'canonical'

SOLVER.linear_model_learning.regularization = False
SOLVER.linear_model_learning.basis_adaptation = True
SOLVER.linear_model_learning.error_estimation = True

SOLVER.test_error = True
SOLVER.test_data = [X_TEST, Y_TEST]
# SOLVER.bases_eval_test = BASES.eval(X_TEST)

SOLVER.rank_adaptation = True
SOLVER.rank_adaptation_options['max_iterations'] = 50
SOLVER.rank_adaptation_options['theta'] = 0.8

SOLVER.tree_adaptation = True
SOLVER.tree_adaptation_options['max_iterations'] = 100
# SOLVER.tree_adaptation_options['force_rank_adaptation'] = True

SOLVER.alternating_minimization_parameters['stagnation'] = 1e-10
SOLVER.alternating_minimization_parameters['max_iterations'] = 50

SOLVER.display = True
SOLVER.alternating_minimization_parameters['display'] = False

SOLVER.model_selection = True
SOLVER.model_selection_options['type'] = 'test_error'

T0 = time()
F, OUTPUT = SOLVER.solve()
T1 = time()
print(T1-T0)

# %% Displays
TEST_ERROR = SOLVER.loss_function.test_error(F, [X_TEST, Y_TEST])
print('\nRanks: %s' % F.tensor.ranks)
print('Loo error = %2.5e' % OUTPUT['error'])
print('Test error = %2.5e' % TEST_ERROR)

F.tensor.plot(title='Active nodes')
F.tensor.tree.plot_dims(title='Dimensions associated to the leaf nodes')
F.tensor.tree.plot_with_labels_at_nodes(F.tensor.representation_rank,
                                        title='Representation ranks')

plt.figure()
x_lin = np.linspace(0, 1, 1000)
plt.plot(x_lin, FUN(x_lin), x_lin, F(T.map(x_lin)))
plt.legend(('True function', 'Approximation'))
plt.show()

# %% Model selection using a slope heuristic instead of the test error
SEL = tensap.ModelSelection()
# Complexity of each model: standard complexity based on the method storage
SEL.data['complexity'] = SEL.complexity(OUTPUT['iterates'],
                                        fun='storage',
                                        c_type='standard')
# Empirical risk associated with each model
SEL.data['empirical_risk'] = [SOLVER.loss_function.risk_estimation(
    x, (X_TRAIN, Y_TRAIN)) for x in OUTPUT['iterates']]
# Shape of the penalization function
SEL.pen_shape = lambda x: x / NUM_TRAIN
# Gap factor used in the slope heuristic
SEL.gap_factor = 2

LAMBDA_HAT, M_HAT, LAMBDA_PATH, M_PATH = SEL.slope_heuristic()

print('Model selected using a test sample:         %i, test error = %2.5e.' %
      (OUTPUT['selected_model_number'], OUTPUT['test_error']))
print('Model selected using the slope heuristic:   %i, test error = %2.5e.' %
      (M_HAT, OUTPUT['test_error_iterations'][M_HAT]))
