"""Approximation and plot of a 1D parametric function using a neural network.

This example also compares two different optimization methods: the classical Adam method, and a natural gradient preconditioning.
"""

import matplotlib.pyplot as plt
import torch

from scimba_torch.approximation_space.nn_space import (
    NNxSpace,
)
from scimba_torch.domain.meshless_domain.domain_1d import Segment1D
from scimba_torch.integration.monte_carlo import DomainSampler, TensorizedSampler
from scimba_torch.integration.monte_carlo_parameters import UniformParametricSampler
from scimba_torch.neural_nets.coordinates_based_nets.mlp import GenericMLP
from scimba_torch.numerical_solvers.collocation_projector import (
    CollocationProjector,
    NaturalGradientProjector,
)
from scimba_torch.plots.plots_nd import plot_abstract_approx_spaces
from scimba_torch.utils.scimba_tensors import LabelTensor


def func_test(x: LabelTensor, mu: LabelTensor):
    x1 = x.get_components()
    mu1 = mu.get_components()
    return mu1 * torch.sin(x1 * 2 * torch.pi)


domain_x = Segment1D((-1.0, 1.0), is_main_domain=True)

sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler([(1.0, 2.0)])]
)

space = NNxSpace(1, 1, GenericMLP, domain_x, sampler, layer_sizes=[24] * 3)

p1 = NaturalGradientProjector(
    space,
    func_test,
    type_linesearch="logarithmic_grid",
    data_linesearch={"M": 10, "interval": [0.0, 2.0]},
)
p1.solve(epochs=90, n_collocation=2000, verbose=True)
p1.save("proj1d_p1")

# p1.load("proj1d_p1")
# p1.space.load_from_best_approx()

space2 = NNxSpace(1, 1, GenericMLP, domain_x, sampler, layers_sizes=[24] * 3)

p2 = CollocationProjector(space2, func_test, bool_preconditioner=True)
p2.solve(epochs=200, n_collocation=2000, verbose=True)
p2.save("proj1d_p2")

# p2.load("proj1d_p2")
# p2.space.load_from_best_approx()

plot_abstract_approx_spaces(
    (p1.space, p2.space),
    domain_x,
    [[1.0, 2.0]],
    loss=(p1.losses, p2.losses),
    solution=func_test,
    error=func_test,
    derivatives=(["ux", "uxx"], [None, "uxx"]),
    # parameters_values=([1.5], [1.6])
)

plt.show()