r"""Solves a 2D Helmholtz PDE using PINNs.

.. math::

    -\delta u + u & = f in \Omega

where :math:`x = (x_1, x_2) \in \Omega = (-1, 1) \times (-1, 1)`, :math:`f` such that :math:`u(x_1, x_2) = \sin(\pi x_1) \sin(\pi x_2)`.

The boundary conditions are either periodic or Neumann boundary conditions.

The neural network is a simple MLP (Multilayer Perceptron).

The optimization is done using Natural Gradient Descent.

The goal of this example is to show how to use feature enrichment in PINNs.
"""

import timeit

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

from scimba_jax.domains.meshless_domains.base import VolumetricDomain
from scimba_jax.domains.meshless_domains.domains_2d import Square2D
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_parameters import (
    UniformParametricSampler,
)
from scimba_jax.nonlinear_approximation.model_class.funcparam_vectorial import (
    ParamScalarFunction,
)
from scimba_jax.nonlinear_approximation.networks.mlp import MLP
from scimba_jax.nonlinear_approximation.numerical_solvers.projectors import Projector
from scimba_jax.physical_models.abstract_physical_model import (
    PHYSICAL_RESIDUALS_TYPE,
    AbstractPhysicalModel,
)
from scimba_jax.physical_models.abstract_residuals import (
    NDARRAYS_FUNC_TYPE,
    PARAM_FUNC_TYPE,
    InteriorResidual,
)
from scimba_jax.physical_models.boundary_residuals import (
    NeumannResidual,
)
from scimba_jax.plots.plots_nd import plot_abstract_approx_spaces

N_COLLOC = 2000
N_BC_COLLOC = 2000
N_EPOCHS = 150

BOUNDS = [(-1.0, 1.0), (-1.0, 1.0)]
DOM_X = Square2D(BOUNDS, is_main_domain=True)
DOM_MU = [(0.0, 1.0)]

FREQUENCY = 1
assert FREQUENCY > 0, "FREQUENCY must be positive"
assert type(FREQUENCY) in [int], "FREQUENCY must be an integer"


def f_rhs(xy: jnp.ndarray, mu: jnp.ndarray) -> jnp.ndarray:
    x, y = xy[0:1], xy[1:2]
    μ = mu[0:1]
    k = FREQUENCY
    return (
        (2 * jnp.pi**2 + 1)
        * k**2
        * jnp.cos(k * jnp.pi * x)
        * jnp.cos(k * jnp.pi * y)
        * jnp.exp(-(μ**2))
    )


def u_exact(x, y, μ):
    return (
        jnp.cos(FREQUENCY * jnp.pi * x)
        * jnp.cos(FREQUENCY * jnp.pi * y)
        * jnp.exp(-(μ**2))
    )


def f_bc_rhs_neumann(xy: jnp.ndarray, n: jnp.ndarray, mu: jnp.ndarray) -> jnp.ndarray:
    return jnp.zeros_like(xy[0:1])


def exact_sol(xy: jnp.ndarray, mu: jnp.ndarray) -> jnp.ndarray:
    x, y = xy[:, 0:1], xy[:, 1:2]
    μ = mu[:, 0:1]
    return u_exact(x, y, μ)


class HelmholtzResidual(InteriorResidual):
    def __init__(
        self,
        domain: VolumetricDomain,
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
    ):
        super().__init__(domain=domain, size=1, model_type="x_mu", f_rhs=f_rhs)

    def construct_residual(self, *vars: PARAM_FUNC_TYPE) -> PARAM_FUNC_TYPE:
        rho = vars[0]
        assert isinstance(rho, ParamScalarFunction)
        lap = rho.laplacian_x()
        return -lap + FREQUENCY**2 * rho


class HelmholtzND(AbstractPhysicalModel):
    """A n D Helmholtz equation with no parameter."""

    def __init__(
        self,
        main_domain: VolumetricDomain,
        f_rhs: NDARRAYS_FUNC_TYPE | None = None,
        bc: str = "weak",
        f_bc_rhs: NDARRAYS_FUNC_TYPE | None = None,
    ):
        super().__init__(main_domain=main_domain)

        self.physical_residuals: PHYSICAL_RESIDUALS_TYPE = {
            self.main_domain.get_label(): HelmholtzResidual(
                domain=main_domain,
                f_rhs=f_rhs,
            ),
        }

        if bc == "weak":
            for boundary in self.boundaries:
                self.physical_residuals[boundary] = NeumannResidual(
                    domain=self.boundaries[boundary],
                    model_type="x_mu",
                    f_rhs=f_bc_rhs,
                )


def create_and_train_pinn(embedding=None, **kwargs) -> Projector:
    key = jax.random.PRNGKey(0)

    sampler = TensorizedSampler(
        [DomainSampler(DOM_X), UniformParametricSampler(DOM_MU)], bc=True
    )

    key, sample_dict = sampler.sample(key, N_COLLOC)

    nn = MLP(
        in_size=2 + 1,
        out_size=1,
        hidden_sizes=[16, 16],
        key=key,
        embedding=embedding,
        **kwargs,
    )

    space = ApproximationSpace({"x": 2}, [(nn, "scalar", None)], model_type="x_mu")

    if embedding is None:
        bc = "weak"
    else:
        if "periodic" in embedding or "neumann" in embedding:
            bc = "strong"
        else:
            bc = "weak"

    model = HelmholtzND(DOM_X, f_rhs, bc, f_bc_rhs_neumann)

    weights = {"interior": [1.0], "boundary": [40.0]}
    pinn = Projector(model, space, sampler, weights=weights)

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

    print("time for %d epochs: " % N_EPOCHS, end - start)

    return pinn


pinn = create_and_train_pinn()

pinn_fourier = create_and_train_pinn(
    embedding="fourier",
    n_fourier_features=6,
    fourier_features_std=FREQUENCY,
    embedding_axes=[0, 1],
)

pinn_neumann = create_and_train_pinn(
    embedding="neumann_square",
    domain_bounds=BOUNDS,
    embedding_axes=[0, 1],
)

pinn_periodic_1 = create_and_train_pinn(
    embedding="periodic",
    periods=[2.0, 2.0],
    embedding_axes=[0, 1],
    n_periodic_features=1,
)

pinn_periodic_3 = create_and_train_pinn(
    embedding="periodic",
    periods=[2.0, 2.0],
    embedding_axes=[0, 1],
    n_periodic_features=3,
)

pinns = [pinn, pinn_fourier, pinn_neumann, pinn_periodic_1, pinn_periodic_3]
titles = [
    "No embedding",
    "Fourier embedding",
    "Neumann embedding",
    "Periodic embedding (1 feature)",
    "Periodic embedding (3 features)",
]

plot_abstract_approx_spaces(
    [p.space for p in pinns],
    DOM_X,
    DOM_MU,
    loss=[p.losses for p in pinns],
    error=exact_sol,
    draw_contours=True,
    n_drawn_contours=20,
    titles=titles,
)
plt.show()