r"""Solves a 2D Poisson PDE written directly in cylindrical (polar) coordinates.

.. math::

    -\left(\frac{\partial^2 u}{\partial r^2} + \frac{1}{r} \frac{\partial u}{\partial r}
    + \frac{1}{r^2} \frac{\partial^2 u}{\partial \theta^2}\right) & = f
    \quad \text{in } \Omega = (r_{\min}, r_{\max}) \times (\theta_{\min}, \theta_{\max}) \\
    u & = g \quad \text{on } \partial \Omega

The physical domain, with coordinates :math:`(r, \theta)`, is generated from a
reference unit square through a :class:`~scimba_jax.domains.domain_mapping.DomainMapping`
-- exactly as the tokamak (Grad-Shafranov) example generates its physical
:math:`(R, Z)` coordinates from a reference unit disk. The neural network
therefore takes :math:`(r, \theta)` directly as input, and the Laplace
operator is written in its cylindrical form, with explicit :math:`1/r` and
:math:`1/r^2` factors -- exactly as the Grad-Shafranov operator carries an
explicit :math:`1/R` factor.

An exact solution :math:`u(r, \theta) = \sin(\alpha (r - r_{\min})) \cos(m \theta)`,
with :math:`\alpha = \pi / (r_{\max} - r_{\min})`, is manufactured so that the
right-hand side :math:`f` and the (non-homogeneous) Dirichlet data :math:`g`
can be computed analytically, and errors can be measured.

Boundary conditions are enforced weakly. The neural network is a simple MLP,
trained with Adam.
"""

import timeit

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

from scimba_jax.domains.domain_mapping import DomainMapping
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.networks.mlp import MLP
from scimba_jax.nonlinear_approximation.numerical_solvers.projectors import Projector
from scimba_jax.physical_models.elliptic_pde.curvilinear_poisson import (
    CylindricalPoissonDirichletND,
)
from scimba_jax.plots.plots_nd import plot_abstract_approx_space

# ── Parameters ─────────────────────────────────────────────────────────────────

R_MIN, R_MAX = 1.0, 2.0
THETA_MIN, THETA_MAX = 0.0, 2.0 * jnp.pi
M = 2
ALPHA = jnp.pi / (R_MAX - R_MIN)

N_COLLOC = 2000
N_BC_COLLOC = 2000
N_EPOCHS = 50
HIDDEN = [16] * 3

key = jax.random.PRNGKey(0)


# ── Manufactured solution and source term ──────────────────────────────────────
# u(r, theta) = sin(alpha * (r - r_min)) * cos(m * theta)
# f = -Delta u, with Delta the cylindrical Laplacian.


def exact_sol_pointwise(x: jnp.ndarray) -> jnp.ndarray:
    r, theta = x[0:1], x[1:2]
    return jnp.sin(ALPHA * (r - R_MIN)) * jnp.cos(M * theta)


def exact_sol(x: jnp.ndarray) -> jnp.ndarray:
    """Batched exact solution, for plotting."""
    r, theta = x[:, 0:1], x[:, 1:2]
    return jnp.sin(ALPHA * (r - R_MIN)) * jnp.cos(M * theta)


def f_rhs(x: jnp.ndarray) -> jnp.ndarray:
    r, theta = x[0:1], x[1:2]
    s = jnp.sin(ALPHA * (r - R_MIN))
    c = jnp.cos(ALPHA * (r - R_MIN))
    return jnp.cos(M * theta) * (ALPHA**2 * s - ALPHA * c / r + (M**2) * s / r**2)


def dirichlet_bc(x: jnp.ndarray, n: jnp.ndarray) -> jnp.ndarray:
    return exact_sol_pointwise(x)


# ── Domain: cylindrical (r, theta) coordinates generated via a mapping ────────
# A reference unit square is mapped to the physical (r, theta) rectangle,
# exactly as the tokamak example maps a reference unit disk to the physical
# (R, Z) tokamak cross-section.

mapping = DomainMapping.rectangle(
    origin=jnp.array([R_MIN, THETA_MIN]),
    x_dir=jnp.array([R_MAX - R_MIN, 0.0]),
    y_dir=jnp.array([0.0, THETA_MAX - THETA_MIN]),
)

domain = Square2D([(0.0, 1.0), (0.0, 1.0)], is_main_domain=True)
domain.set_mapping(mapping, bounds_postmap=[(R_MIN, R_MAX), (THETA_MIN, THETA_MAX)])

sampler = TensorizedSampler([DomainSampler(domain)], bc=True)

# ── Network & space ────────────────────────────────────────────────────────────

key, subkey = jax.random.split(key)
nn = MLP(in_size=2, out_size=1, hidden_sizes=HIDDEN, key=subkey)
space = ApproximationSpace({"x": 2}, [(nn, "scalar", None)], model_type="x")

model = CylindricalPoissonDirichletND(
    domain,
    f_rhs=lambda *args: f_rhs(*args),
    bc="weak",
    f_bc_rhs=lambda *args: dirichlet_bc(*args),
    model_type="x",
)

# ── Training ────────────────────────────────────────────────────────────────────

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

key, sample_dict = sampler.sample(key, N_COLLOC, N_BC_COLLOC)
print("initial loss: ", pinn.evaluate_loss(space, sample_dict))

t0 = timeit.default_timer()
key, pinn = pinn.project(key, space, N_EPOCHS, N_COLLOC, N_BC_COLLOC)
print("best loss: ", pinn.best_loss)
print("time for %d epochs: " % N_EPOCHS, timeit.default_timer() - t0)

# ── Plot ─────────────────────────────────────────────────────────────────────

plot_abstract_approx_space(
    pinn.space,
    domain,
    loss=pinn.losses,
    residual=pinn.model,
    error=exact_sol,
    draw_contours=True,
    n_drawn_contours=20,
    title="2D Poisson equation in cylindrical (r, theta) coordinates",
)
plt.show()