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

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

Note: this is similar to projection_transport_t1x1v.py but with a velocity domain being a segment instead of a circle.
"""

import matplotlib.pyplot as plt
import torch

from scimba_torch.approximation_space.nn_space import NNxvSpace
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,
    UniformVelocitySamplerOnCuboid,
)
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_space,
    plot_abstract_approx_spaces,
)

torch.random.manual_seed(0)


def exact(t, x, v, mu):
    x_x = x.get_components()
    v_a = v.get_components()
    exp_1 = torch.exp(-((x_x - v_a * t) ** 2 * 10))
    exp_2 = torch.exp(-((v_a - 1.0) ** 2 * 0.01))
    return exp_1 * exp_2


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_t = (0.0, 0.1)
domain_x = Segment1D((-1.0, 1.0), is_main_domain=True)
domain_v = Segment1D((-2.0, 2.0))
domain_mu = []

sampler = TensorizedSampler(
    [
        DomainSampler(domain_x),
        UniformVelocitySamplerOnCuboid(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()

plot_abstract_approx_space(
    space1,
    domain_x,
    velocity_domain=domain_v,
    solution=exact_t1,
    error=exact_t1,
    derivatives=["ux", "uv"],
)

plt.show()

plot_abstract_approx_spaces(
    (
        space0,
        space1,
    ),
    domain_x,
    velocity_domains=(domain_v,),
    solution=(exact_t0, exact_t1),
    error=(exact_t0, exact_t1),
    derivatives=["ux", "uv"],
    titles=[
        "computed approximation at t=%.2e" % t_0,
        "computed approximation at t=%.2e" % t_1,
    ],
)

plt.show()