r"""Learns an oscillatory function using a neural network or a spectral approximation."""

import matplotlib.pyplot as plt
import torch

from scimba_torch.approximation_space.spectral_space import SpectralxSpace
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.numerical_solvers.collocation_projector import (
    CollocationProjector,
    LinearProjector,
    NaturalGradientProjector,
)
from scimba_torch.optimizers.optimizers_data import OptimizerData
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()
    sigma = 0.4
    return (
        torch.sin(2.5 * torch.pi * x1)
        * torch.sin(4.5 * torch.pi * x2)
        * torch.exp(-((x1) ** 2 + (x2) ** 2) / (2 * sigma**2))
    )


torch.manual_seed(0)
domain_x = Square2D([(-1.0, 1.0), (-1.0, 1.0)], is_main_domain=True)
sampler = TensorizedSampler([DomainSampler(domain_x), UniformParametricSampler([])])
space1 = SpectralxSpace(
    1, "sine", 6, bounds=[(-1.0, 1.0), (-1.0, 1.0)], integrator=sampler
)
p = LinearProjector(space1, func_test)
p.solve(n_collocation=10000, verbose=True)


space2 = SpectralxSpace(
    1, "sine", 6, bounds=[(-1.0, 1.0), (-1.0, 1.0)], integrator=sampler
)
opt_1 = {
    "name": "adam",
    "optimizer_args": {"lr": 4.0e-2},
}
default_opt = OptimizerData(opt_1)
p2 = CollocationProjector(space2, func_test, optimizers=default_opt)
p2.solve(epochs=500, n_collocation=10000, verbose=True)

space3 = SpectralxSpace(
    1, "sine", 6, bounds=[(-1.0, 1.0), (-1.0, 1.0)], integrator=sampler
)
p3 = NaturalGradientProjector(space3, func_test)
p3.solve(epochs=20, n_collocation=10000, verbose=True)

plot_abstract_approx_spaces(
    (p.space, p2.space, p3.space),  # the approximation spaces
    (domain_x),  # the spatial domain
    loss=(
        None,
        p2.losses,
        p3.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,
)

plt.show()


# x, mu = sampler.sample(40000)
# x1, x2 = x.get_components()
# w_exact = func_test(x, mu)
# w1 = space1.evaluate(x, mu).w.detach().cpu()
# w2 = space2.evaluate(x, mu).w.detach().cpu()
# e1 = (space1.evaluate(x, mu).w - w_exact).detach().cpu()
# e2 = (space2.evaluate(x, mu).w - w_exact).detach().cpu()
# fig, ax = plt.subplots(2, 2, figsize=(10, 8))
# x1 = x1.detach().cpu()
# x2 = x2.detach().cpu()
## Premier scatter avec colorbar
# sc = ax[0, 0].scatter(x1, x2, c=w1, s=1, cmap="turbo")
# sc = ax[0, 1].scatter(
#    x1,
#    x2,
#    c=w1 - w_exact.detach().cpu(),
#    s=1,
#    cmap="turbo",
#    label=f"{torch.sqrt(torch.mean(e1**2)):3.2e}",
# )
# fig.colorbar(sc, ax=ax[0, 1])
# ax[0, 1].legend()
## Premier scatter avec colorbar
# sc = ax[1, 0].scatter(x1, x2, c=w2, s=1, cmap="turbo")
# fig.colorbar(sc, ax=ax[1, 0])
# sc = ax[1, 1].scatter(
#    x1,
#    x2,
#    c=w2 - w_exact.detach().cpu(),
#    s=1,
#    cmap="turbo",
#    label=f"{torch.sqrt(torch.mean(e2**2)):3.2e}",
# )
# fig.colorbar(sc, ax=ax[1, 1])
# ax[1, 1].legend()
# plt.show()