"""Approximation of a function using several optimization strategies."""

import matplotlib.pyplot as plt
import torch

from scimba_torch.approximation_space.nn_space import NNxSpace
from scimba_torch.domain.meshless_domain.domain_2d import Square2D
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,
    LinearProjector,
    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, x2 = x.get_components()
    mu1 = mu.get_components()
    return torch.sin(x1) * mu1 * torch.cos(x2)


def bc_test(x: LabelTensor, xn: LabelTensor, mu: LabelTensor):
    return func_test(x, mu)


domain_x = Square2D([(-1, 1), (-1, 1)], is_main_domain=True)
sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler([(1.0, 2.0)])]
)
opt = {
    "name": "adam",
    "optimizer_args": {"lr": 1.5e-2},
}

print("\n\n")
print(" ################################################# ")
print(" # CollocationProjector, no preconditioning      # ")
print(" ################################################# ")

torch.manual_seed(0)

space = NNxSpace(1, 1, GenericMLP, domain_x, sampler, layer_sizes=[20])
print("ndof with Adam: ", space.ndof)

p1 = CollocationProjector(space, func_test, optimizers=opt, bool_preconditioner=False)

new_solve = True
if new_solve or not p1.load(__file__, "projector_no_precond"):
    p1.solve(epochs=800, n_collocation=1000)
    p1.save(__file__, "projector_no_precond")

print("\n\n")
print(" ################################################# ")
print(" # NaturalGradientProjector                      # ")
print(" ################################################# ")

torch.manual_seed(0)

space2 = NNxSpace(1, 1, GenericMLP, domain_x, sampler, layer_sizes=[20])
print("ndof with natural gradient: ", space2.ndof)

p2 = NaturalGradientProjector(space2, func_test, adaptive_matrix_regularization=True)

new_solve = True
if new_solve or not p2.load(__file__, "natural"):
    # load from p1.space
    # p2.space.load_from_dict(p1.space.dict_for_save())
    p2.solve(epochs=50, n_collocation=1000)
    p2.save(__file__, "natural")
print("\n\n")
print(" ################################################# ")
print(" # AnagramProjector                              # ")
print(" ################################################# ")

torch.manual_seed(0)

space3 = NNxSpace(1, 1, GenericMLP, domain_x, sampler, layer_sizes=[20])
print("ndof with anagram: ", space3.ndof)

p3 = NaturalGradientProjector(space3, func_test, ng_algo="ANaGRAM", svd_threshold=1e-3)

new_solve = True
if new_solve or not p3.load(__file__, "anagram"):
    p3.solve(epochs=60, n_collocation=1000)
    p3.save(__file__, "anagram")

print("\n\n")
print(" ################################################# ")
print(" # LinearProjector                               # ")
print(" ################################################# ")

torch.manual_seed(0)

space4 = NNxSpace(
    1, 1, GenericMLP, domain_x, sampler, layer_sizes=[20], last_layer_has_bias=True
)
print("ndof for linear projection: ", space4.ndof)

p4 = LinearProjector(space4, func_test)

p4.solve(n_collocation=100000)

plot_abstract_approx_spaces(
    (
        p1.space,
        p2.space,
        p3.space,
        p4.space,
    ),  # the approximation spaces
    (domain_x),  # the spatial domain
    ([[1.0, 2.0]]),  # the parameter's domain
    loss=(
        p1.losses,
        p2.losses,
        p3.losses,
        p4.losses,
    ),  # for plot of the loss: the losses
    solution=(func_test),  # for plot of the exact sol: sol
    error=(func_test),  # for plot of the error with respect to a func: the func
    draw_contours=True,
    n_drawn_contours=20,
    parameters_values="mean",
)

plt.show()