"""An example illustrating rsdf nn to learn a bean shape from a parametric curve."""

import matplotlib.pyplot as plt
import torch

from scimba_torch.geometry.parametric_hypersurface import (
    ParametricHyperSurface,
)
from scimba_torch.geometry.regularized_sdf_projectors import (
    learn_regularized_sdf,
)
from scimba_torch.integration.monte_carlo import DomainSampler, TensorizedSampler
from scimba_torch.integration.monte_carlo_parameters import UniformParametricSampler
from scimba_torch.optimizers.optimizers_data import OptimizerData
from scimba_torch.plots.plot_regularized_sdf_projector import (
    plot_regularized_sdf_projector,
)
from scimba_torch.utils import Mapping

torch.manual_seed(2)


# define the bean 2d function
def bean_2d_function(t: torch.Tensor) -> torch.Tensor:
    """The bean 2d function.

    Args:
        t: The argument.

    Returns:
        c(t).
    """
    a, b = 3, 5
    sin = torch.sin(t)
    cos = torch.cos(t)

    x = (sin**a + cos**b) * cos
    y = (sin**a + cos**b) * sin

    return torch.cat((x, y), dim=-1)


# define the bean 2d mapping, compose it with a rotation
bean_2d_mapping = Mapping(1, 2, bean_2d_function)
bean_2d_mapping = Mapping.compose(bean_2d_mapping, Mapping.rot_2d(-torch.pi / 2))
# create a parametric hyper_surface
bean_2d = ParametricHyperSurface([(0.0, 2 * torch.pi)], bean_2d_mapping)

# a bounding box for the hypersurface
# bean_2d_bb = [(-0.4, 1.2), (-1.2, 0.4)]

opt_1 = {
    "name": "adam",
    "optimizer_args": {"lr": 9.9e-3, "betas": (0.9, 0.999)},
}
opt_2 = {
    "name": "lbfgs",
    "switch_at_epoch_ratio": 0.7,
}
optimizer_arg = OptimizerData(opt_1, opt_2)

ginn = learn_regularized_sdf(
    parametric_hyper_surface=bean_2d,
    preconditioner="None",  # can be "None", "ENG" or "Anagram"
    mode="new",  # can be "new", "load" or "resume"
    load_from="bean_2d",
    save_to="bean_2d",
    optimizers=OptimizerData(opt_1, opt_2),
    epochs=2000,
    n_collocation=4000,
    n_bc_collocation=2000,
    verbose=False,
)

# plot the result
plot_regularized_sdf_projector(
    ginn,
    n_visu=512,  # number of points for the implicit plot
    draw_contours=True,
    n_drawn_contours=20,
)

plt.show()

ginn = learn_regularized_sdf(
    parametric_hyper_surface=bean_2d,
    preconditioner="ENG",
    mode="new",
    load_from="bean_2d_ENG",
    save_to="bean_2d_ENG",
    epochs=200,
    n_collocation=4000,
    n_bc_collocation=2000,
    verbose=False,
)

# plot the result
plot_regularized_sdf_projector(
    ginn,
    n_visu=512,
    draw_contours=True,
    n_drawn_contours=20,
)

plt.show()

x_sampler = DomainSampler(ginn.geometric_domain)
sampler = TensorizedSampler((x_sampler, UniformParametricSampler([])))

x, mu = sampler.sample(10)
y = ginn.evaluate(x, mu)
print("y.shape: ", y.shape)
dydx = tuple(ginn.space.grad(y, x))[0]
print("dydx.shape: ", dydx.shape)