r"""Solves the viscous Burgers advection equation in 1D using a PINN.

.. math::

    \partial_t u + \partial_x \frac {u^2}{2} - \sigma \partial_{xx} u & = f in \Omega \times (0, T) \\
    u & = g on \partial \Omega \times (0, T) \\
    u & = u_0 on \Omega \times {0}

where :math:`u: \partial \Omega \times (0, T) \to \mathbb{R}` is the unknown function, :math:`\Omega \subset \mathbb{R}` is the spatial domain and :math:`(0, T) \subset \mathbb{R}` is the time domain.
Dirichlet boundary conditions are prescribed.

The equation is solved on a segment domain; weak boundary and initial conditions are used.
The solution is approximated with a standard temporal PINN trained with ENG optimizer.
"""

import timeit

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

from scimba_jax.domains.meshless_domains.domains_1d import Segment1D
from scimba_jax.nonlinear_approximation.approximation_spaces.approximation_spaces import (
    ApproximationSpace,
)
from scimba_jax.nonlinear_approximation.integration.monte_carlo import (
    DomainSampler,
    TensorizedSampler,
)
from scimba_jax.nonlinear_approximation.integration.monte_carlo_time import (
    UniformTimeSampler,
)
from scimba_jax.nonlinear_approximation.networks.mlp import MLP
from scimba_jax.nonlinear_approximation.numerical_solvers.projectors import Projector
from scimba_jax.physical_models.temporal_pde.viscous_burgers import ViscousBurgers1D
from scimba_jax.plots.plots_nd import (
    plot_abstract_approx_space,
)

N_COLLOC = 2000
N_BC_COLLOC = 2000
N_IC_COLLOC = 2000
N_EPOCHS = 200


def f_rhs(t: jnp.ndarray, x: jnp.ndarray, sigma: float) -> jnp.ndarray:
    exp_neg_t = jnp.exp(-t)
    sin_x = jnp.sin(2 * jnp.pi * x)
    cos_x = jnp.cos(2 * jnp.pi * x)
    return (
        exp_neg_t
        * sin_x
        * (2 * jnp.pi * (cos_x * exp_neg_t + 2 * jnp.pi * sigma) - 1.0)
    )


def exact_sol(t: jnp.ndarray, x: jnp.ndarray):
    return jnp.sin(2 * jnp.pi * x) * jnp.exp(-t)


def f_init(x: jnp.ndarray):
    t = jnp.zeros_like(x)
    return exact_sol(t, x)


domain_t = (0.0, 1.0)
domain_x = [(-1.0, 1.0)]

dx = Segment1D(domain_x[0], is_main_domain=True)
sampler = TensorizedSampler(
    [
        UniformTimeSampler(domain_t),
        DomainSampler(dx),
    ],
    model_type="t_x",
    bc=True,
    ic=True,
)

# create the model
model = ViscousBurgers1D(
    main_domain=dx,
    time_domain=domain_t,
    bc="weak",
    ic="weak",
    f_rhs=f_rhs,
    f_ic_rhs=lambda *args: f_init(*args),
    sigma=1e-2 / jnp.pi,
)

key = jax.random.PRNGKey(42)
# create the pinn

print("@@@@@@@@@@@@@@@ train with ENG @@@@@@@@@@@@@@@@@@@@@@")
nn = MLP(in_size=2, out_size=1, hidden_sizes=[16, 32, 16], key=key)

space = ApproximationSpace(
    {"x": 1},
    [(nn, "scalar", None)],
    model_type="t_x",
)
pinn = Projector(model, space, sampler)

start = timeit.default_timer()
key, pinn = pinn.project(key, space, N_EPOCHS, N_COLLOC, N_BC_COLLOC, N_IC_COLLOC)
end = timeit.default_timer()

print("best loss: ", pinn.best_loss)
print("time for %d epochs: " % N_EPOCHS, end - start)

plot_abstract_approx_space(
    pinn.space,
    dx,
    time_domain=domain_t,
    solution=exact_sol,
    error=exact_sol,
    loss=pinn.losses,
    residual=pinn.model,
    titles=("Approximated solution of 1D viscous Burgers equation with source"),
)

plt.show()