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

"""

import numpy as np
import tensap

# %% Definitions
ORDER = 6
NK = 8


def fun(inputs):
    """
    Example function used for the tutorial.

    Parameters
    ----------
    inputs : numpy.ndarray
        The inputs points at which the function is to be evaluated.

    Returns
    -------
    numpy.ndarray
        The evaluations of the function at the input points.

    """
    return (
        1 / (3 + inputs[:, 0] + inputs[:, 2])
        + inputs[:, 1]
        + np.cos(inputs[:, 3] + inputs[:, 4])
    )


GRID = tensap.FullTensorGrid(np.linspace(0, 1, NK), ORDER)
ARRAY = GRID.array()

U = fun(ARRAY)
U = tensap.FullTensor(U, ORDER, np.full(ORDER, NK))

# %% Higher Order SVD - truncation in Tucker format
TR = tensap.Truncator()
TR.tolerance = 1e-8
UR = TR.hosvd(U)
print("Tucker format - Compression at tolerance %2.5e" % TR.tolerance)
print("Error = %2.5e" % ((UR.full() - U).norm() / U.norm()))
print("Storage = %d" % UR.storage())
print("Dimension of spaces = %s\n" % UR.tensors[0].shape)

# %% Tree-based format
ARITY_INTERVAL = [2, 3]
TREE = tensap.DimensionTree.random(ORDER, ARITY_INTERVAL)
TR = tensap.Truncator()
TR.tolerance = 1e-15
UR = TR.hsvd(U, TREE)

print("Tree-based format - Compression at tolerance %2.5e" % TR.tolerance)
print("Error = %2.5e" % ((UR.full() - U).norm() / U.norm()))
print("Storage = %d" % UR.storage())
print("Ranks = %s\n" % UR.ranks)

TR.tolerance = 1e-6
URR = TR.hsvd(UR, TREE)
print("Rounding at tolerance %2.5e" % TR.tolerance)
print("Error = %2.5e" % ((URR.full() - U).norm() / U.norm()))
print("Storage = %d" % URR.storage())
print("Ranks = %s\n" % URR.ranks)


# %% Truncation in tensor-train format
TR = tensap.Truncator()
TR.tolerance = 1e-8
UR = TR.ttsvd(U)
print(UR)

SIN_VAL = UR.singular_values()

print("Tensor Train format - Compression at tolerance %2.5e" % TR.tolerance)
print("Error = %2.5e" % ((UR.full() - U).norm() / U.norm()))
print("Storage = %d" % UR.storage())
print("TT-rank = %s\n" % UR.ranks)

# %% Truncation in tensor-train format after permutation
PERM = np.random.permutation(ORDER)
U_PERM = U.transpose(PERM)

TR = tensap.Truncator()
TR.tolerance = 1e-10
TR.max_rank = 200

UR = TR.ttsvd(U_PERM)
print("With permutation...")
print("\tOrdering = %s" % PERM)
print("\tError = %2.5e" % ((UR.full() - U_PERM).norm() / U_PERM.norm()))
print("\tStorage = %d" % UR.storage())
print("\tRanks = %s\n" % UR.ranks)

# %% Tensor-train tree optimization
print("Tree optimization (leaves permutation)...")
UR_PERM = UR.optimize_leaves_permutations(1e-10, 10)
print("\tInitial storage = %d" % UR.storage())
print("\tError = %2.5e" % ((UR_PERM.full() - U_PERM).norm() / U.norm()))
print("\tFinal storage = %d" % UR_PERM.storage())
print("\tRanks = %s" % UR_PERM.ranks)

UR.tree.plot_dims(title="Dimensions before tree optimization")
UR_PERM.tree.plot_dims(title="Dimensions before after optimization")

# %% Global tree optimization
print("Tree optimization...")
UR_PERM = UR.optimize_dimension_tree(1e-10, 100)
print("\tInitial storage = %d" % UR.storage())
print("\tError = %2.5e" % ((UR_PERM.full() - U_PERM).norm() / U.norm()))
print("\tFinal storage = %d" % UR_PERM.storage())
print("\tRanks = %s" % UR_PERM.ranks)

UR.tree.plot_dims(title="Dimensions before tree optimization")
UR_PERM.tree.plot_dims(title="Dimensions before after optimization")
