r"""Approximation and plot of a 3D time-dependent function using a neural network.

This example uses Adam to approximate a function :math:`x, v \in \mathbb{R}^2 \times \mathbb{S}^1 \mapsto f(t, x, v) \in \mathbb{R}` at two different time instants.
"""

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

from scimba_torch.approximation_space.nn_space import NNxvSpace
from scimba_torch.domain.meshless_domain.domain_2d import Circle2D, Square2D
from scimba_torch.integration.monte_carlo import DomainSampler, TensorizedSampler
from scimba_torch.integration.monte_carlo_parameters import (
    UniformParametricSampler,
    UniformVelocitySampler,
)
from scimba_torch.neural_nets.coordinates_based_nets.mlp import GenericMLP
from scimba_torch.numerical_solvers.collocation_projector import (
    CollocationProjector,
)
from scimba_torch.plots.plots_nd import plot_abstract_approx_spaces

torch.random.manual_seed(0)


def exact(t, x, v, mu):
    x_x, x_y = x.get_components()
    v_x, v_y = v.get_components()
    return (
        torch.exp(-((x_x - v_x * t) ** 2 * 10))
        * torch.exp(-((x_y - v_y * t) ** 2 * 10))
        * torch.exp(-((v_x - 1.0) ** 2 * 0.01))
        * torch.exp(-((v_y) ** 2 * 0.01))
    )


t_0 = 0.0


def exact_t0(x, v, mu):
    return exact(t_0, x, v, mu)


t_1 = 0.5


def exact_t1(x, v, mu):
    return exact(t_1, x, v, mu)


domain_x = Square2D([(-1, 1), (-1, 1)], is_main_domain=True)
domain_v = Circle2D(torch.Tensor([0, 0]), 1.0)
domain_mu = []

sampler = TensorizedSampler(
    [
        DomainSampler(domain_x),
        UniformVelocitySampler(domain_v),
        UniformParametricSampler(domain_mu),
    ]
)

space0 = NNxvSpace(
    1, 0, GenericMLP, domain_x, domain_v, sampler, layer_sizes=[20, 20, 20]
)

proj0 = CollocationProjector(space0, exact_t0, bool_preconditioner=False)

resume_solve = False
if resume_solve or not proj0.load(__file__, "t0"):
    proj0.solve(epochs=1000, n_collocation=3000, verbose=True)
    proj0.save(__file__, "t0")

proj0.space.load_from_best_approx()

space1 = NNxvSpace(
    1, 0, GenericMLP, domain_x, domain_v, sampler, layer_sizes=[20, 20, 20]
)

proj1 = CollocationProjector(space1, exact_t1, bool_preconditioner=False)

resume_solve = False
if resume_solve or not proj1.load(__file__, "t1"):
    proj1.solve(epochs=1000, n_collocation=3000, verbose=True)
    proj1.save(__file__, "t1")

proj1.space.load_from_best_approx()

cuts = [
    ([0.0, 0.0], [-0.5, 0.5]),
    ([0.0, 0.2], [0.0, 1.0]),
]

plot_abstract_approx_spaces(
    space1,
    domain_x,
    velocity_domains=domain_v,
    solution=exact_t1,
    derivatives=["uv1"],
    cuts=cuts[0],
    velocity_values=([0.0, np.pi / 2.0],),
    velocity_strs=[r"$0\pi$", "$\\frac{\\pi}{2}$"],
)

plt.show()

plot_abstract_approx_spaces(
    (space0, space1),
    domain_x,
    velocity_domains=domain_v,
    solution=(exact_t0, exact_t1),
    derivatives=["uv1"],
    cuts=cuts[0],
    velocity_values=([np.pi / 4.0],),
    velocity_strs=["$\\frac{\\pi}{4}$"],
    titles=[
        "computed approximation at t=%.2e" % t_0,
        "computed approximation at t=%.2e" % t_1,
    ],
)

plt.show()