# 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 FunctionalBasis, projection, interpolation, least-squares
approximation.

"""

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


# %% Interpolation on a polynomial space using Chebyshev Points
def FUN(x):
    # Function to approximate
    return np.cos(10 * x)


FUN = tensap.UserDefinedFunction(FUN, 1)
FUN.evaluation_at_multiple_points = True
X = tensap.UniformRandomVariable(-1, 1)

# Interpolation basis and points
P = 30
H = tensap.PolynomialFunctionalBasis(X.orthonormal_polynomials(), range(P + 1))
X_CP = tensap.chebyshev_points(H.cardinal(), [-1, 1])

# Interpolation of the function
F = H.interpolate(FUN, X_CP)

# Displays and error
print("Interpolation on a polynomial space using Chebyshev Points")
plt.figure()
X_PLOT = np.linspace(-1, 1, 100)
plt.plot(X_PLOT, FUN(X_PLOT))
plt.plot(X_PLOT, F(X_PLOT))
plt.legend(("True function", "Interpolation"))
plt.show()
N = 100
ERR_L2, ERR_L_INF = F.test_error(FUN, N, X)
print("Mean squared error = %2.5e" % ERR_L2)


# %% Interpolation on a polynomial space using magic Points
def FUN(x):
    # Function to approximate
    return np.cos(10 * x)


FUN = tensap.UserDefinedFunction(FUN, 1)
FUN.evaluation_at_multiple_points = True
X = tensap.UniformRandomVariable(-1, 1)

# Interpolation basis and points
P = 30
H = tensap.PolynomialFunctionalBasis(X.orthonormal_polynomials(), range(P + 1))
X_MP = H.magic_points(X.random(10000))[0]

# Interpolation of the function
F = H.interpolate(FUN, X_MP)

# Displays and error
print("\nInterpolation on a polynomial space using magic points")
plt.figure()
X_PLOT = np.linspace(-1, 1, 100)
plt.plot(X_PLOT, FUN(X_PLOT))
plt.plot(X_PLOT, F(X_PLOT))
plt.legend(("True function", "Interpolation"))
plt.show()
N = 100
ERR_L2, ERR_L_INF = F.test_error(FUN, N, X)
print("Mean squared error = %2.5e" % ERR_L2)


# %% Projection on polynomial space through quadrature
def FUN(x):
    # Function to approximate
    return x**2 / 2


FUN = tensap.UserDefinedFunction(FUN, 1)
FUN.evaluation_at_multiple_points = True
X = tensap.NormalRandomVariable()

# Integration rule
I0 = X.gauss_integration_rule(5)

# Approximation basis
P = 3
H = tensap.PolynomialFunctionalBasis(X.orthonormal_polynomials(), range(P + 1))

# Computation of the projection
F = H.projection(FUN, I0)

# Displays and error
print("\nProjection on polynomial space through quadrature")
N = 100
ERR_L2, ERR_L_INF = F.test_error(FUN, N, X)
print("Mean squared error = %2.5e" % ERR_L2)


# %% Interpolation on UserDefinedFunctionalBasis
def FUN(x):
    # Function to approximate
    return np.exp(-((x - 1 / 2) ** 2))


FUN = tensap.UserDefinedFunction(FUN, 1)
FUN.evaluation_at_multiple_points = True
X = tensap.UniformRandomVariable(0, 1)

# Fourier Basis
N = 15
H = np.empty(2 * N + 1, dtype=object)
H[0] = lambda x: np.ones(np.size(x))
for i in np.arange(1, N + 1):
    H[2 * i - 1] = lambda x, i=i: np.sqrt(2) * np.cos(2 * np.pi * i * x)
    H[2 * i] = lambda x, i=i: np.sqrt(2) * np.sin(2 * np.pi * i * x)
H = tensap.UserDefinedFunctionalBasis(H, X)

# Display of the Fourier basis
plt.figure()
X_PLOT = np.linspace(0, 1, 100)
plt.plot(X_PLOT, H.eval(X_PLOT))
plt.title("Fourier basis")
plt.show()

# Computation of the interpolation
X_MP = H.magic_points(X.random(10000))[0]
F = H.interpolate(FUN, X_MP)

# Displays and error
print("\nInterpolation on UserDefinedFunctionalBasis")
plt.figure()
X_PLOT = np.linspace(0, 1, 100)
plt.plot(X_PLOT, FUN(X_PLOT))
plt.plot(X_PLOT, F(X_PLOT))
plt.legend(("True function", "Interpolation"))
plt.show()
N = 100
ERR_L2, ERR_L_INF = F.test_error(FUN, N, X)
print("Mean squared error = %2.5e" % ERR_L2)


# %% Interpolation with a radial basis
def FUN(x):
    # Function to approximate
    return np.exp(-(x**2))


FUN = tensap.UserDefinedFunction(FUN, 1)
FUN.evaluation_at_multiple_points = True
X = tensap.UniformRandomVariable(-1, 1)

# Radial basis
N = 20
X_LIN = np.linspace(-1, 1, N)
S = 10 / N


def k(x, y, s=S):
    return np.exp(-((x - y) ** 2) / s**2)


H = np.empty(N, dtype=object)
for i in range(N):
    H[i] = lambda y, x=X_LIN[i], k=k: k(y, x)
H = tensap.UserDefinedFunctionalBasis(H, X)

# Computation of the interpolation
X_MP = H.magic_points(X.random(100000))[0]
F = H.interpolate(FUN, X_MP)

# Displays and error
print("\nInterpolation with a radial basis")
plt.figure()
X_PLOT = np.linspace(-1, 1, 100)
plt.plot(X_PLOT, FUN(X_PLOT))
plt.plot(X_PLOT, F(X_PLOT))
plt.legend(("True function", "Interpolation"))
plt.show()
N = 100
ERR_L2, ERR_L_INF = F.test_error(FUN, N, X)
print("Mean squared error = %2.5e" % ERR_L2)


# %% Projection on polynomial space through quadrature
def FUN(x):
    # Function to approximate
    return x**2 / 2


FUN = tensap.UserDefinedFunction(FUN, 1)
FUN.evaluation_at_multiple_points = True
X = tensap.NormalRandomVariable()

# Integration rule
I0 = X.gauss_integration_rule(5)

# Approximation basis
P = 3
H = tensap.PolynomialFunctionalBasis(tensap.HermitePolynomials(), range(P + 1))

# Computation of the approximation
F = H.projection(FUN, I0)

# Derivative and second derivative of F through projection
DF = H.projection(lambda x, F=F: F.eval_derivative(1, x), I0)
DDF = H.projection(lambda x, DF=DF: DF.eval_derivative(1, x), I0)

# Displays and error
print("\nProjection on polynomial space through quadrature")
N = 100
ERR_L2, ERR_L_INF = F.test_error(FUN, N, X)
print("Mean squared error = %2.5e" % ERR_L2)
plt.figure()
X_PLOT = np.linspace(-1, 1, 100)
plt.plot(X_PLOT, F(X_PLOT))
plt.plot(X_PLOT, DF(X_PLOT))
plt.plot(X_PLOT, DDF(X_PLOT))
plt.legend(("F", "df", "ddf"))
plt.show()


# %% Least-squares approximation
def FUN(x):
    # Function to approximate
    return np.exp(-((x - 1 / 2) ** 2))


FUN = tensap.UserDefinedFunction(FUN, 1)
FUN.evaluation_at_multiple_points = True
X = tensap.UniformRandomVariable(0, 1)

# Approximation basis: Hermite polynomials of maximal degree P
P = 10
H = tensap.PolynomialFunctionalBasis(X.orthonormal_polynomials(), range(P + 1))

# Solver
SOLVER = tensap.LinearModelLearningSquareLoss()
SOLVER.regularization = True
SOLVER.regularization_type = "l1"
SOLVER.basis_adaptation = False
SOLVER.model_selection = True

# Training sample
X_TRAIN = X.random(100)
Y_TRAIN = FUN(X_TRAIN)

# Computation of the approximation
SOLVER.basis = H
SOLVER.training_data = [X_TRAIN, Y_TRAIN]
F, OUTPUT = SOLVER.solve()

# Displays and error
print("\nLeast-squares approximation")
plt.figure()
X_PLOT = np.linspace(0, 1, 100)
plt.plot(X_PLOT, FUN(X_PLOT))
plt.plot(X_PLOT, F(X_PLOT))
plt.legend(("True function", "Interpolation"))
plt.show()
X_TEST = X.random(100)
F_X_TEST = F(X_TEST)
Y_TEST = FUN(X_TEST)
print(
    "Mean squared error = %2.5e"
    % (np.linalg.norm(Y_TEST - F_X_TEST) / np.linalg.norm(Y_TEST))
)
